From 5eb44a94317f963bf6defcd66987b808cb5ab604 Mon Sep 17 00:00:00 2001 From: Lucas Manning Date: Thu, 20 Jul 2023 09:59:53 -0700 Subject: [PATCH] Add config flags and sandbox chroot configuration for TPU proxying. PiperOrigin-RevId: 549662855 --- runsc/cmd/chroot.go | 55 +++++++++++++++++++++++++++++ runsc/cmd/util/BUILD | 5 ++- runsc/cmd/util/tpu.go | 78 ++++++++++++++++++++++++++++++++++++++++++ runsc/config/config.go | 3 ++ runsc/config/flags.go | 1 + 5 files changed, 141 insertions(+), 1 deletion(-) create mode 100644 runsc/cmd/util/tpu.go diff --git a/runsc/cmd/chroot.go b/runsc/cmd/chroot.go index bc0487755..8f62ef144 100644 --- a/runsc/cmd/chroot.go +++ b/runsc/cmd/chroot.go @@ -18,11 +18,14 @@ import ( "errors" "fmt" "os" + "path" "path/filepath" + "regexp" specs "github.com/opencontainers/runtime-spec/specs-go" "golang.org/x/sys/unix" "gvisor.dev/gvisor/pkg/log" + "gvisor.dev/gvisor/runsc/cmd/util" "gvisor.dev/gvisor/runsc/config" "gvisor.dev/gvisor/runsc/specutils" ) @@ -120,6 +123,9 @@ func setUpChroot(pidns bool, spec *specs.Spec, conf *config.Config) error { if err := nvproxyUpdateChroot(chroot, spec, conf); err != nil { return fmt.Errorf("error configuring chroot for Nvidia GPUs: %w", err) } + if err := tpuProxyUpdateChroot(chroot, conf); err != nil { + return fmt.Errorf("error configuring chroot for TPU devices: %w", err) + } if err := specutils.SafeMount("", chroot, "", unix.MS_REMOUNT|unix.MS_RDONLY|unix.MS_BIND, "", "/proc"); err != nil { return fmt.Errorf("error remounting chroot in read-only: %v", err) @@ -128,6 +134,55 @@ func setUpChroot(pidns bool, spec *specs.Spec, conf *config.Config) error { return pivotRoot(chroot) } +func tpuProxyUpdateChroot(chroot string, conf *config.Config) error { + if !conf.TPUProxy { + return nil + } + devices, err := util.EnumerateHostTPUDevices() + if err != nil { + return fmt.Errorf("enumerating TPU device files: %w", err) + } + for _, deviceNum := range devices { + devPath := fmt.Sprintf("/dev/accel%d", deviceNum) + if err := mountInChroot(chroot, devPath, devPath, "bind", unix.MS_BIND); err != nil { + return fmt.Errorf("error mounting %q in chroot: %v", devPath, err) + } + finfo, err := os.Stat(path.Join(chroot, devPath)) + if err != nil { + return fmt.Errorf("error statting %q: %v", devPath, err) + } + // Ensure the file mounted in was a char device file. + if finfo.Mode()&os.ModeType != os.ModeCharDevice|os.ModeDevice { + return fmt.Errorf("unexpected file type for %q, want %s, got %s", path.Join(chroot, devPath), os.ModeCharDevice|os.ModeDevice, finfo.Mode()&os.ModeType) + + } + // Multiple paths link to the /sys/devices/pci0000:00/ + // directory that contains all relevant sysfs accel device info that we need + // bind mounted into the sandbox chroot. We can construct this path by + // reading the link below, which points to + // /sys/devices/pci0000:00//accel/accel# and traversing up 2 + // directories. + sysAccelPath := fmt.Sprintf("/sys/class/accel/accel%d", deviceNum) + sysAccelLink, err := os.Readlink(sysAccelPath) + if err != nil { + return fmt.Errorf("error reading %q: %v", sysAccelPath, err) + } + // Ensure the link is in the form we expect. + sysAccelLinkMatcher := regexp.MustCompile(fmt.Sprintf(`../../devices/pci0000:00/(\d+:\d+:\d+\.\d+)/accel/accel%d`, deviceNum)) + if !sysAccelLinkMatcher.MatchString(sysAccelLink) { + return fmt.Errorf("unexpected link %q -> %q, link should have %q format", sysAccelPath, sysAccelLink, sysAccelLinkMatcher.String()) + } + sysPCIDeviceDir, err := filepath.Abs(path.Join(filepath.Dir(sysAccelPath), sysAccelLink, "../..")) + if err != nil { + return fmt.Errorf("error parsing path %q: %v", sysAccelPath, err) + } + if err := mountInChroot(chroot, sysPCIDeviceDir, sysPCIDeviceDir, "bind", unix.MS_BIND); err != nil { + return fmt.Errorf("error mounting %q in chroot: %v", sysAccelPath, err) + } + } + return nil +} + func nvproxyUpdateChroot(chroot string, spec *specs.Spec, conf *config.Config) error { if !specutils.GPUFunctionalityRequested(spec, conf) { return nil diff --git a/runsc/cmd/util/BUILD b/runsc/cmd/util/BUILD index 25eb274ff..d3dea9f8a 100644 --- a/runsc/cmd/util/BUILD +++ b/runsc/cmd/util/BUILD @@ -7,7 +7,10 @@ package( go_library( name = "util", - srcs = ["util.go"], + srcs = [ + "tpu.go", + "util.go", + ], visibility = [ "//runsc/cli:__subpackages__", "//runsc/cmd:__subpackages__", diff --git a/runsc/cmd/util/tpu.go b/runsc/cmd/util/tpu.go new file mode 100644 index 000000000..1a4c9d93a --- /dev/null +++ b/runsc/cmd/util/tpu.go @@ -0,0 +1,78 @@ +// Copyright 2023 The gVisor Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package util + +import ( + "fmt" + "os" + "path/filepath" + "regexp" + "strconv" + "strings" +) + +const googleVendorID = 0x1AE0 + +var tpuV4DeviceIDs = map[uint64]any{0x005E: nil, 0x0056: nil} + +// TODO(b/288456802): Add support for /dev/vfio controlled accelerators. +// This is required for v5+ TPUs. + +// EnumerateHostTPUDevices returns the accelerator device minor numbers of all +// TPUs on the machine. +func EnumerateHostTPUDevices() ([]uint32, error) { + paths, err := filepath.Glob("/dev/accel*") + if err != nil { + return nil, fmt.Errorf("enumerating TPU device files: %w", err) + } + + accelDeviceRegex := regexp.MustCompile(`^/dev/accel(\d+)$`) + var devMinors []uint32 + for _, path := range paths { + if ms := accelDeviceRegex.FindStringSubmatch(path); ms != nil { + index, err := strconv.ParseUint(ms[1], 10, 32) + if err != nil { + return nil, fmt.Errorf("invalid host device file %q: %w", path, err) + } + + vendor, err := readHexInt(fmt.Sprintf("/sys/class/accel/accel%d/device/vendor", index)) + if err != nil { + return nil, err + } + if vendor != googleVendorID { + continue + } + deviceID, err := readHexInt(fmt.Sprintf("/sys/class/accel/accel%d/device/device", index)) + if err != nil { + return nil, err + } + if _, ok := tpuV4DeviceIDs[deviceID]; !ok { + continue + } + + devMinors = append(devMinors, uint32(index)) + } + } + return devMinors, nil +} + +func readHexInt(path string) (uint64, error) { + data, err := os.ReadFile(path) + if err != nil { + return 0, err + } + numStr := strings.Trim(strings.TrimSpace(strings.TrimPrefix(string(data), "0x")), "\x00") + return strconv.ParseUint(numStr, 16, 64) +} diff --git a/runsc/config/config.go b/runsc/config/config.go index c50bf404f..8418c66e6 100644 --- a/runsc/config/config.go +++ b/runsc/config/config.go @@ -290,6 +290,9 @@ type Config struct { // containers or set by `docker --gpus`. NVProxyDocker bool `flag:"nvproxy-docker"` + // TPUProxy enables support for TPUs. + TPUProxy bool `flag:"tpuproxy"` + // TestOnlyAllowRunAsCurrentUserWithoutChroot should only be used in // tests. It allows runsc to start the sandbox process as the current // user, and without chrooting the sandbox process. This can be diff --git a/runsc/config/flags.go b/runsc/config/flags.go index 2ce64e031..b3d45a6ad 100644 --- a/runsc/config/flags.go +++ b/runsc/config/flags.go @@ -121,6 +121,7 @@ func RegisterFlags(flagSet *flag.FlagSet) { // Flags that control sandbox runtime behavior: accelerator related. flagSet.Bool("nvproxy", false, "EXPERIMENTAL: enable support for Nvidia GPUs") flagSet.Bool("nvproxy-docker", false, "Expose GPUs to containers based on NVIDIA_VISIBLE_DEVICES, as requested by the container or set by `docker --gpus`. Allows containers to self-serve GPU access and thus disabled by default for security. libnvidia-container must be installed on the host. No effect unless --nvproxy is enabled.") + flagSet.Bool("tpuproxy", false, "EXPERIMENTAL: enable support for TPU device passthrough.") // Test flags, not to be used outside tests, ever. flagSet.Bool("TESTONLY-unsafe-nonroot", false, "TEST ONLY; do not ever use! This skips many security measures that isolate the host from the sandbox.")