From 1e2be4b4f7220dbc5e6612307243bcffa98fbe2f Mon Sep 17 00:00:00 2001 From: Ayush Ranjan Date: Mon, 13 Nov 2023 11:02:07 -0800 Subject: [PATCH] Move TPU device ownership to gofer process. We still bind mount /sys/devices/pci0000:00//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 | 1 + pkg/sentry/devices/accel/device.go | 13 ++++-- runsc/boot/vfs.go | 53 +++++++++++-------------- runsc/cmd/chroot.go | 20 ++++------ runsc/cmd/gofer.go | 13 +++++- runsc/cmd/util/tpu.go | 64 ++++++++++++++---------------- runsc/container/container.go | 2 +- runsc/specutils/specutils.go | 26 +++++++++--- 8 files changed, 105 insertions(+), 87 deletions(-) diff --git a/pkg/sentry/devices/accel/BUILD b/pkg/sentry/devices/accel/BUILD index 617681d50..3e09b97b3 100644 --- a/pkg/sentry/devices/accel/BUILD +++ b/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", diff --git a/pkg/sentry/devices/accel/device.go b/pkg/sentry/devices/accel/device.go index a2d143e74..763113ae3 100644 --- a/pkg/sentry/devices/accel/device.go +++ b/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{ diff --git a/runsc/boot/vfs.go b/runsc/boot/vfs.go index 62ecb2533..4f42e1886 100644 --- a/runsc/boot/vfs.go +++ b/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//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 diff --git a/runsc/cmd/chroot.go b/runsc/cmd/chroot.go index df573d8ca..b4f2df559 100644 --- a/runsc/cmd/chroot.go +++ b/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//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/ // directory that contains all relevant sysfs accel device info that we need diff --git a/runsc/cmd/gofer.go b/runsc/cmd/gofer.go index 0968c3928..6cde79972 100644 --- a/runsc/cmd/gofer.go +++ b/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 } diff --git a/runsc/cmd/util/tpu.go b/runsc/cmd/util/tpu.go index 3f11af8a9..17d971f4f 100644 --- a/runsc/cmd/util/tpu.go +++ b/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) { diff --git a/runsc/container/container.go b/runsc/container/container.go index c799af944..b6b1b5371 100644 --- a/runsc/container/container.go +++ b/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. diff --git a/runsc/specutils/specutils.go b/runsc/specutils/specutils.go index 009735434..3609da000 100644 --- a/runsc/specutils/specutils.go +++ b/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 }