Add config flags and sandbox chroot configuration for TPU proxying.

PiperOrigin-RevId: 549662855

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/<pci_address>
+		// 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/<pci_address>/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

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__",

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)
+}

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

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.")