Move TPU device ownership to gofer process.

We still bind mount /sys/devices/pci0000:00/<pci_address>/accel/accel# files
into the sandbox process in read-only mode for all TPU devices on host.

PiperOrigin-RevId: 582021030

pkg/sentry/devices/accel/BUILD

@@ -22,6 +22,7 @@ go_library(
         "//pkg/abi/tpu",
         "//pkg/cleanup",
         "//pkg/context",
+        "//pkg/devutil",
         "//pkg/errors/linuxerr",
         "//pkg/fdnotifier",
         "//pkg/hostarch",

pkg/sentry/devices/accel/device.go

@@ -20,8 +20,10 @@ import (
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/abi/linux"
 	"gvisor.dev/gvisor/pkg/context"
+	"gvisor.dev/gvisor/pkg/devutil"
 	"gvisor.dev/gvisor/pkg/errors/linuxerr"
 	"gvisor.dev/gvisor/pkg/fdnotifier"
+	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/sentry/kernel"
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
 	"gvisor.dev/gvisor/pkg/sync"
@@ -44,12 +46,17 @@ type tpuV4Device struct {
 }
 
 func (dev *tpuV4Device) Open(ctx context.Context, mnt *vfs.Mount, vfsd *vfs.Dentry, opts vfs.OpenOptions) (*vfs.FileDescription, error) {
+	devClient := devutil.GoferClientFromContext(ctx)
+	if devClient == nil {
+		log.Warningf("devutil.CtxDevGoferClient is not set")
+		return nil, linuxerr.ENOENT
+	}
 	dev.mu.Lock()
 	defer dev.mu.Unlock()
-	hostPath := fmt.Sprintf("/dev/accel%d", dev.minor)
-	hostFD, err := unix.Openat(-1, hostPath, int((opts.Flags&unix.O_ACCMODE)|unix.O_NOFOLLOW), 0)
+	name := fmt.Sprintf("accel%d", dev.minor)
+	hostFD, err := devClient.OpenAt(ctx, name, opts.Flags)
 	if err != nil {
-		ctx.Warningf("accelDevice: failed to open host %s: %v", hostPath, err)
+		ctx.Warningf("accelDevice: failed to open device %s: %v", name, err)
 		return nil, err
 	}
 	fd := &tpuV4FD{

runsc/boot/vfs.go

@@ -1186,40 +1186,33 @@ func tpuProxyRegisterDevices(info *containerInfo, vfsObj *vfs.VirtualFilesystem)
 	if !specutils.TPUProxyIsEnabled(info.spec, info.conf) {
 		return nil
 	}
-	// At this point /dev/accel just contains the TPU devices have been mounted
-	// into the sandbox chroot. Enumerate all of them and create sentry devices.
-	paths, err := filepath.Glob("/dev/accel*")
-	if err != nil {
-		return fmt.Errorf("enumerating accel device files: %w", err)
-	}
+	// At this point /sys/devices/pci0000:00/<pci_address>/accel/accel# contains
+	// all the TPU devices on the host. Enumerate them and register TPU devices.
 	pciAddrs, err := filepath.Glob("/sys/devices/pci0000:00/*")
 	if err != nil {
 		return fmt.Errorf("enumerating PCI device files: %w", err)
 	}
-	for _, accelPath := range paths {
-		accelDeviceRegex := regexp.MustCompile(`^/dev/accel(\d+)$`)
-		if ms := accelDeviceRegex.FindStringSubmatch(accelPath); ms != nil {
-			deviceNum, _ := strconv.ParseUint(ms[1], 10, 32)
-
-			var pciDevicePath string
-			for _, pciPath := range pciAddrs {
-				if _, err := os.Stat(path.Join(pciPath, fmt.Sprintf("accel/accel%d", deviceNum))); err == nil {
-					pciDevicePath = pciPath
-				}
-			}
-			var deviceIDBytes []byte
-			if deviceIDBytes, err = os.ReadFile(path.Join(pciDevicePath, "device")); err != nil {
-				return fmt.Errorf("reading PCI device ID: %w", err)
-			}
-			deviceIDStr := strings.Replace(string(deviceIDBytes), "0x", "", -1)
-			deviceID, err := strconv.ParseInt(strings.TrimSpace(deviceIDStr), 16, 64)
-			if err != nil {
-				return fmt.Errorf("parsing PCI device ID: %w", err)
-			}
-
-			if err := accel.RegisterTPUV4Device(vfsObj, uint32(deviceNum), deviceID == tpu.TPUV4liteDeviceID); err != nil {
-				return fmt.Errorf("registering accel driver: %w", err)
-			}
+	pciPathRegex := regexp.MustCompile(`^/sys/devices/pci0000:00/\d+:\d+:\d+\.\d+/accel/accel(\d+)$`)
+	for _, pciPath := range pciAddrs {
+		ms := pciPathRegex.FindStringSubmatch(pciPath)
+		if ms == nil {
+			continue
+		}
+		deviceNum, err := strconv.ParseUint(ms[1], 10, 32)
+		if err != nil {
+			return fmt.Errorf("parsing PCI device number: %w", err)
+		}
+		var deviceIDBytes []byte
+		if deviceIDBytes, err = os.ReadFile(path.Join(pciPath, "device")); err != nil {
+			return fmt.Errorf("reading PCI device ID: %w", err)
+		}
+		deviceIDStr := strings.Replace(string(deviceIDBytes), "0x", "", -1)
+		deviceID, err := strconv.ParseInt(strings.TrimSpace(deviceIDStr), 16, 64)
+		if err != nil {
+			return fmt.Errorf("parsing PCI device ID: %w", err)
+		}
+		if err := accel.RegisterTPUV4Device(vfsObj, uint32(deviceNum), deviceID == tpu.TPUV4liteDeviceID); err != nil {
+			return fmt.Errorf("registering accel driver: %w", err)
 		}
 	}
 	return nil

runsc/cmd/chroot.go

@@ -134,23 +134,19 @@ func tpuProxyUpdateChroot(chroot string, spec *specs.Spec, conf *config.Config)
 	if !specutils.TPUProxyIsEnabled(spec, conf) {
 		return nil
 	}
-	devices, err := util.EnumerateHostTPUDevices()
+	// Bind mount /sys/devices/pci0000:00/<pci_address>/accel/accel# for all
+	// TPU devices on the host.
+	paths, err := filepath.Glob("/dev/accel*")
 	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))
+	for _, devPath := range paths {
+		deviceNum, valid, err := util.ExtractTpuDeviceMinor(devPath)
 		if err != nil {
-			return fmt.Errorf("error statting %q: %v", devPath, err)
+			return fmt.Errorf("extracting TPU device minor: %w", 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)
-
+		if !valid {
+			continue
 		}
 		// Multiple paths link to the /sys/devices/pci0000:00/<pci_address>
 		// directory that contains all relevant sysfs accel device info that we need

runsc/cmd/gofer.go

@@ -528,6 +528,15 @@ func shouldExposeNvidiaDevice(path string) bool {
 	return nvidiaDevPathReg.MatchString(path)
 }
 
+// shouldExposeTpuDevice returns true if path refers to a TPU device which
+// should be exposed to the container.
+//
+// Precondition: tpuproxy is enabled.
+func shouldExposeTpuDevice(path string) bool {
+	_, valid, _ := util.ExtractTpuDeviceMinor(path)
+	return valid
+}
+
 func (g *Gofer) setupDev(spec *specs.Spec, conf *config.Config, root, procPath string) error {
 	if err := os.MkdirAll(filepath.Join(root, "dev"), 0777); err != nil {
 		return fmt.Errorf("creating dev directory: %v", err)
@@ -537,8 +546,10 @@ func (g *Gofer) setupDev(spec *specs.Spec, conf *config.Config, root, procPath s
 		return nil
 	}
 	nvproxyEnabled := specutils.NVProxyEnabled(spec, conf)
+	tpuproxyEnabled := specutils.TPUProxyIsEnabled(spec, conf)
 	for _, dev := range spec.Linux.Devices {
-		shouldMount := nvproxyEnabled && shouldExposeNvidiaDevice(dev.Path)
+		shouldMount := (nvproxyEnabled && shouldExposeNvidiaDevice(dev.Path)) ||
+			(tpuproxyEnabled && shouldExposeTpuDevice(dev.Path))
 		if !shouldMount {
 			continue
 		}

runsc/cmd/util/tpu.go

@@ -17,7 +17,6 @@ package util
 import (
 	"fmt"
 	"os"
-	"path/filepath"
 	"regexp"
 	"strconv"
 	"strings"
@@ -32,42 +31,37 @@ var tpuV4DeviceIDs = map[uint64]any{tpu.TPUV4DeviceID: nil, tpu.TPUV4liteDeviceI
 // 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)
-	}
-
+// ExtractTpuDeviceMinor returns the accelerator device minor number for that
+// the passed device path. If the passed device is not a valid TPU device, then
+// it returns false. TPU device is defined as:
+// * Path is /dev/accel#.
+// * Vendor is googleVendorID.
+// * Device ID is one of tpuV4DeviceIDs.
+func ExtractTpuDeviceMinor(path string) (uint32, bool, error) {
 	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))
-		}
+	ms := accelDeviceRegex.FindStringSubmatch(path)
+	if ms == nil {
+		return 0, false, nil
 	}
-	return devMinors, nil
+	index, err := strconv.ParseUint(ms[1], 10, 32)
+	if err != nil {
+		return 0, false, 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 0, false, err
+	}
+	if vendor != googleVendorID {
+		return 0, false, nil
+	}
+	deviceID, err := readHexInt(fmt.Sprintf("/sys/class/accel/accel%d/device/device", index))
+	if err != nil {
+		return 0, false, err
+	}
+	if _, ok := tpuV4DeviceIDs[deviceID]; !ok {
+		return 0, false, nil
+	}
+	return uint32(index), true, nil
 }
 
 func readHexInt(path string) (uint64, error) {

runsc/container/container.go

@@ -1162,7 +1162,7 @@ func (c *Container) waitForStopped() error {
 // shouldCreateDeviceGofer indicates whether a device gofer connection should
 // be created.
 func shouldCreateDeviceGofer(spec *specs.Spec, conf *config.Config) bool {
-	return specutils.GPUFunctionalityRequested(spec, conf)
+	return specutils.GPUFunctionalityRequested(spec, conf) || specutils.TPUFunctionalityRequested(spec, conf)
 }
 
 // shouldSpawnGofer indicates whether the gofer process should be spawned.

runsc/specutils/specutils.go

@@ -575,12 +575,28 @@ func TPUProxyIsEnabled(spec *specs.Spec, conf *config.Config) bool {
 		return true
 	}
 	val, ok := spec.Annotations[annotationTPU]
-	if ok {
-		ret, err := strconv.ParseBool(val)
-		if val != "" && err != nil {
-			log.Warningf("tpuproxy annotation set to invalid value %q. Skipping.", val)
+	if !ok {
+		return false
+	}
+	ret, err := strconv.ParseBool(val)
+	if err != nil {
+		log.Warningf("tpuproxy annotation set to invalid value %q: %w. Skipping.", val, err)
+	}
+	return ret
+}
+
+// TPUFunctionalityRequested returns true if the container should have access
+// to TPU functionality.
+func TPUFunctionalityRequested(spec *specs.Spec, conf *config.Config) bool {
+	if !TPUProxyIsEnabled(spec, conf) {
+		return false
+	}
+	if spec.Linux != nil {
+		for _, dev := range spec.Linux.Devices {
+			if strings.HasPrefix(dev.Path, "/dev/accel") {
+				return true
+			}
 		}
-		return ret
 	}
 	return false
 }