From 290789bab84aab78a4f613cd10a83bf95ebbe745 Mon Sep 17 00:00:00 2001 From: Lucas Manning Date: Fri, 13 Sep 2024 11:07:35 -0700 Subject: [PATCH] Refactor tpu chroot operations. Ubuntu TPU images do not have the vfio-dev directories that COS images do, so we need a more robust way of setting up the sandbox chroot to handle this case. This change implements a way to get devices and minor numbers into the sandbox with minimal support from the host filesystem and cleans up a few methods to reflect their current usage. Addresses #10795 PiperOrigin-RevId: 674363342 --- pkg/abi/linux/dev.go | 7 - pkg/sentry/devices/tpuproxy/BUILD | 11 +- pkg/sentry/devices/tpuproxy/accel/accel.go | 3 + pkg/sentry/devices/tpuproxy/tpuproxy.go | 136 +++++++--------- pkg/sentry/devices/tpuproxy/tpuproxy_test.go | 66 -------- pkg/sentry/devices/tpuproxy/vfio/vfio.go | 27 +++- pkg/sentry/fsimpl/sys/pci.go | 4 +- pkg/sentry/vfs/device.go | 9 ++ runsc/boot/BUILD | 1 - runsc/boot/vfs.go | 47 +++--- runsc/cmd/BUILD | 3 + runsc/cmd/chroot.go | 131 ++++++---------- runsc/cmd/chroot_test.go | 155 +++++++++++++++++++ runsc/cmd/gofer.go | 2 +- runsc/cmd/util/BUILD | 1 - runsc/cmd/util/tpu.go | 147 +++++++++++------- 16 files changed, 424 insertions(+), 326 deletions(-) delete mode 100644 pkg/sentry/devices/tpuproxy/tpuproxy_test.go create mode 100644 runsc/cmd/chroot_test.go diff --git a/pkg/abi/linux/dev.go b/pkg/abi/linux/dev.go index e03312b97..48f25e92f 100644 --- a/pkg/abi/linux/dev.go +++ b/pkg/abi/linux/dev.go @@ -70,10 +70,3 @@ const ( // ACCEL_MAJOR is the major device number for compute accelerator devices. ACCEL_MAJOR = 121 ) - -// Major device numbers for VFIO-based TPU. -const ( - // Major devices number between 243 and 254 are usually reserved for local use. - // The device number 245 is used by VFIO based TPU in GCP. - VFIO_MAJOR = 245 -) diff --git a/pkg/sentry/devices/tpuproxy/BUILD b/pkg/sentry/devices/tpuproxy/BUILD index cb64939e1..516266f3a 100644 --- a/pkg/sentry/devices/tpuproxy/BUILD +++ b/pkg/sentry/devices/tpuproxy/BUILD @@ -1,4 +1,4 @@ -load("//tools:defs.bzl", "go_library", "go_test") +load("//tools:defs.bzl", "go_library") package(default_applicable_licenses = ["//:license"]) @@ -17,16 +17,13 @@ go_library( "//pkg/abi/gasket", "//pkg/abi/linux", "//pkg/abi/tpu", + "//pkg/context", + "//pkg/fspath", "//pkg/seccomp", "//pkg/sentry/devices/tpuproxy/accel", "//pkg/sentry/devices/tpuproxy/vfio", + "//pkg/sentry/kernel/auth", "//pkg/sentry/vfs", "@org_golang_x_sys//unix:go_default_library", ], ) - -go_test( - name = "tpuproxy_test", - srcs = ["tpuproxy_test.go"], - library = ":tpuproxy", -) diff --git a/pkg/sentry/devices/tpuproxy/accel/accel.go b/pkg/sentry/devices/tpuproxy/accel/accel.go index f8a0d3e93..115739933 100644 --- a/pkg/sentry/devices/tpuproxy/accel/accel.go +++ b/pkg/sentry/devices/tpuproxy/accel/accel.go @@ -90,6 +90,9 @@ func (dev *accelDevice) Open(ctx context.Context, mnt *vfs.Mount, vfsd *vfs.Dent // RegisterTPUDevice registers all devices implemented by this package in vfsObj. func RegisterTPUDevice(vfsObj *vfs.VirtualFilesystem, minor uint32, lite bool) error { + if vfsObj.IsDeviceRegistered(vfs.CharDevice, linux.ACCEL_MAJOR, minor) { + return nil + } return vfsObj.RegisterDevice(vfs.CharDevice, linux.ACCEL_MAJOR, minor, &accelDevice{ lite: lite, minor: minor, diff --git a/pkg/sentry/devices/tpuproxy/tpuproxy.go b/pkg/sentry/devices/tpuproxy/tpuproxy.go index 7c68c16b2..0142e389a 100644 --- a/pkg/sentry/devices/tpuproxy/tpuproxy.go +++ b/pkg/sentry/devices/tpuproxy/tpuproxy.go @@ -17,6 +17,7 @@ package tpuproxy import ( + "bytes" "fmt" "os" "path" @@ -25,100 +26,71 @@ import ( "strconv" "strings" + "golang.org/x/sys/unix" "gvisor.dev/gvisor/pkg/abi/tpu" + "gvisor.dev/gvisor/pkg/context" + "gvisor.dev/gvisor/pkg/fspath" "gvisor.dev/gvisor/pkg/sentry/devices/tpuproxy/accel" "gvisor.dev/gvisor/pkg/sentry/devices/tpuproxy/vfio" + "gvisor.dev/gvisor/pkg/sentry/kernel/auth" "gvisor.dev/gvisor/pkg/sentry/vfs" ) -const ( - pciPathGlobTPUv4 = "/sys/devices/pci0000:*/**/accel/accel*" - pciPathGlobTPUv5 = "/sys/devices/pci0000:*/**/vfio-dev/vfio*" - iommuGroupPathGlob = "/sys/kernel/iommu_groups/*/devices/*" -) - var ( - // pathGlobToPathRegex is a map that points a TPU PCI path glob to its path regex. - // TPU v4 devices are accessible via /sys/devices/pci0000:00//accel/accel# on the host. - // TPU v5 devices are accessible via at /sys/devices/pci0000:00//vfio-dev/vfio# on the host. - pathGlobToPathRegex = map[string]string{ - pciPathGlobTPUv4: `^/sys/devices/pci0000:[[:xdigit:]]{2}/(0000:([[:xdigit:]]{2}|[[:xdigit:]]{4}):[[:xdigit:]]{2}\.[[:xdigit:]]{1,2}/)+accel/accel(\d+)$`, - pciPathGlobTPUv5: `^/sys/devices/pci0000:[[:xdigit:]]{2}/(0000:([[:xdigit:]]{2}|[[:xdigit:]]{4}):[[:xdigit:]]{2}\.[[:xdigit:]]{1,2}/)+vfio-dev/vfio(\d+)$`, - } + // TPUv4DeviceRegex is the regex for detecting TPUv4 device paths. + TPUv4DeviceRegex = regexp.MustCompile(`/dev/accel(\d+)`) + + // TPUv5DeviceRegex is the regex for detecting TPUv5 device paths. + TPUv5DeviceRegex = regexp.MustCompile(`/dev/vfio/(\d+)`) ) -// RegisterHostTPUDevices enumerates TPU devices on the host and registers them -// in the sandbox VFS. -func RegisterHostTPUDevices(vfsObj *vfs.VirtualFilesystem, allowedDeviceIDs map[int64]any) error { - for pciPathGlobal, pathRegex := range pathGlobToPathRegex { - pciAddrs, err := filepath.Glob(pciPathGlobal) - if err != nil { - return fmt.Errorf("enumerating PCI device files: %w", err) - } - pciPathRegex := regexp.MustCompile(pathRegex) - for _, pciPath := range pciAddrs { - ms := pciPathRegex.FindStringSubmatch(pciPath) - if ms == nil { - continue - } - minorNum, err := strconv.ParseUint(ms[len(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/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 _, ok := allowedDeviceIDs[deviceID]; !ok { - return fmt.Errorf("unsupported TPU device with ID: 0x%x", deviceID) - } - // VFIO iommu groups correspond to the device number. Use these - // paths to get the correct number for the sentry-internal TPU - // device files. - var deviceNum int - switch deviceID { - case tpu.TPUV4DeviceID, tpu.TPUV4liteDeviceID: - deviceNum = int(deviceNum) - case tpu.TPUV5eDeviceID, tpu.TPUV5pDeviceID: - groupPaths, err := filepath.Glob(iommuGroupPathGlob) - if err != nil { - return fmt.Errorf("enumerating IOMMU group files: %w", err) - } - for _, groupPath := range groupPaths { - pci := path.Base(groupPath) - if strings.Contains(pciPath, pci) { - n, err := strconv.Atoi(strings.Split(groupPath, "/")[4]) - if err != nil { - return fmt.Errorf("parsing IOMMU group minor number: %w", err) - } - deviceNum = n - break - } - } - default: - return fmt.Errorf("unsupported TPU device with ID: 0x%x", deviceID) - } - if err := registerTPUDevice(vfsObj, uint32(minorNum), uint32(deviceNum), deviceID); err != nil { - return fmt.Errorf("registering TPU driver: %w", err) - } - } +// RegisterTPUv4Device registers the TPUv4 device with the provided minor number +// where the corresponding PCI device is located at pciPath. Accel devices +// always have their device file number set to their minor number. +func RegisterTPUv4Device(ctx context.Context, creds *auth.Credentials, root vfs.VirtualDentry, vfsObj *vfs.VirtualFilesystem, devPath string, minorNum uint32) error { + // Get the PCI path from the accel device's symlink at + // /sys/class/accel/accel\d+. The link will be in the form + // "../../devices/pci0000:*/**/accel/accel\d+". + linkPath := filepath.Join("/sys/class/accel", filepath.Base(devPath)) + linkContent, err := vfsObj.ReadlinkAt(ctx, creds, &vfs.PathOperation{Root: root, Start: root, Path: fspath.Parse(linkPath)}) + if err != nil { + return fmt.Errorf("reading link %q: %w", linkPath, err) + } + // Exclude the ../../devices prefix and the accel/accel\d+ suffix. + pciPath := strings.TrimSuffix(strings.TrimPrefix(linkContent, "../../devices"), fmt.Sprintf("accel/%s", filepath.Base(devPath))) + pciDeviceIDPath := path.Join("/sys/devices", pciPath, "device") + + fd, err := unix.Openat(-1, pciDeviceIDPath, unix.O_RDONLY|unix.O_NOFOLLOW, 0) + if err != nil { + return err + } + file := os.NewFile(uintptr(fd), pciDeviceIDPath) + defer file.Close() + buf := bytes.Buffer{} + if _, err := buf.ReadFrom(file); err != nil { + return err + } + + deviceIDStr := strings.Replace(buf.String(), "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.RegisterTPUDevice(vfsObj, minorNum, deviceID == tpu.TPUV4liteDeviceID); err != nil { + return fmt.Errorf("registering TPU driver: %w", err) } return nil } -// registerTPUDevice registers a TPU device in vfsObj based on the given device ID. -func registerTPUDevice(vfsObj *vfs.VirtualFilesystem, minor, deviceNum uint32, deviceID int64) error { - switch deviceID { - case tpu.TPUV4DeviceID, tpu.TPUV4liteDeviceID: - return accel.RegisterTPUDevice(vfsObj, minor, deviceID == tpu.TPUV4liteDeviceID) - case tpu.TPUV5eDeviceID, tpu.TPUV5pDeviceID: - return vfio.RegisterTPUDevice(vfsObj, minor, deviceNum, false /* useDevGofer */) - default: - return fmt.Errorf("unsupported TPU device with ID: 0x%x", deviceID) +// RegisterTPUv5Device registers the TPUv5 device with the provided device path +// and minor number. +func RegisterTPUv5Device(vfsObj *vfs.VirtualFilesystem, devPath string, minorNum uint32) error { + deviceNum, err := strconv.ParseInt(path.Base(devPath), 10, 32) + if err != nil { + return fmt.Errorf("parsing device path number: %w", err) } + if err := vfio.RegisterTPUDevice(vfsObj, uint32(minorNum), uint32(deviceNum), true /* useDevGofer */); err != nil { + return fmt.Errorf("registering TPU driver: %w", err) + } + return nil } diff --git a/pkg/sentry/devices/tpuproxy/tpuproxy_test.go b/pkg/sentry/devices/tpuproxy/tpuproxy_test.go deleted file mode 100644 index d9ac53bff..000000000 --- a/pkg/sentry/devices/tpuproxy/tpuproxy_test.go +++ /dev/null @@ -1,66 +0,0 @@ -// Copyright 2024 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 tpuproxy - -import ( - "path/filepath" - "regexp" - "slices" - "testing" -) - -func TestTPUPath(t *testing.T) { - for _, tst := range []struct { - name string - pathGlob string - path string - submatch []string - }{ - { - name: "TPUv4PCIPathMatch", - pathGlob: pciPathGlobTPUv4, - path: "/sys/devices/pci0000:00/0000:00:01.0/accel/accel16", - submatch: []string{"/sys/devices/pci0000:00/0000:00:01.0/accel/accel16", "0000:00:01.0/", "00", "16"}, - }, - { - name: "TPUv4PCIPathNoMatch", - pathGlob: pciPathGlobTPUv4, - path: "/sys/devices/pci0000:00/0000:00:01.0/accel/123", - submatch: nil, - }, - { - name: "TPUv5PCIPathMatch", - pathGlob: pciPathGlobTPUv5, - path: "/sys/devices/pci0000:00/0000:00:05.0/vfio-dev/vfio20", - submatch: []string{"/sys/devices/pci0000:00/0000:00:05.0/vfio-dev/vfio20", "0000:00:05.0/", "00", "20"}, - }, - { - name: "TPUv5PCIPathNoMatch", - pathGlob: pciPathGlobTPUv5, - path: "/sys/devices/pci0000:00/0000:00:05.0/vfio/vfio20", - submatch: nil, - }, - } { - t.Run(tst.name, func(t *testing.T) { - if _, err := filepath.Glob(tst.pathGlob); err != nil { - t.Errorf("Malformed path glob: %v", err) - } - pathRegex := regexp.MustCompile(pathGlobToPathRegex[tst.pathGlob]) - if submatch := pathRegex.FindStringSubmatch(tst.path); !slices.Equal(submatch, tst.submatch) { - t.Errorf("Match TPU PCI path, got: %v, want: %v", submatch, tst.submatch) - } - }) - } -} diff --git a/pkg/sentry/devices/tpuproxy/vfio/vfio.go b/pkg/sentry/devices/tpuproxy/vfio/vfio.go index 49243544c..448506345 100644 --- a/pkg/sentry/devices/tpuproxy/vfio/vfio.go +++ b/pkg/sentry/devices/tpuproxy/vfio/vfio.go @@ -41,6 +41,12 @@ const ( VFIOPath = "/dev/vfio/vfio" ) +var ( + tpuDeviceMajor uint32 + tpuDeviceMajorInit sync.Once + tpuDeviceMajorInitErr error +) + // device implements TPU's vfs.Device for /dev/vfio/[0-9]+ // // +stateify savable @@ -158,7 +164,14 @@ func (dev *vfioDevice) Open(ctx context.Context, mnt *vfs.Mount, d *vfs.Dentry, // RegisterTPUDevice registers devices implemented by this package in vfsObj. func RegisterTPUDevice(vfsObj *vfs.VirtualFilesystem, minor, deviceNum uint32, useDevGofer bool) error { - return vfsObj.RegisterDevice(vfs.CharDevice, linux.VFIO_MAJOR, minor, &tpuDevice{ + major, err := GetTPUDeviceMajor(vfsObj) + if err != nil { + return err + } + if vfsObj.IsDeviceRegistered(vfs.CharDevice, major, minor) { + return nil + } + return vfsObj.RegisterDevice(vfs.CharDevice, major, minor, &tpuDevice{ minor: minor, num: deviceNum, useDevGofer: useDevGofer, @@ -171,6 +184,9 @@ func RegisterTPUDevice(vfsObj *vfs.VirtualFilesystem, minor, deviceNum uint32, u // RegisterVFIODevice registers VFIO devices that are implemented by this package in vfsObj. func RegisterVFIODevice(vfsObj *vfs.VirtualFilesystem, useDevGofer bool) error { + if vfsObj.IsDeviceRegistered(vfs.CharDevice, linux.MISC_MAJOR, VFIO_MINOR) { + return nil + } return vfsObj.RegisterDevice(vfs.CharDevice, linux.MISC_MAJOR, VFIO_MINOR, &vfioDevice{ useDevGofer: useDevGofer, }, &vfs.RegisterDeviceOptions{ @@ -179,3 +195,12 @@ func RegisterVFIODevice(vfsObj *vfs.VirtualFilesystem, useDevGofer bool) error { FilePerms: 0666, }) } + +// GetTPUDeviceMajor returns the dynamically allocated major number for the vfio +// device. +func GetTPUDeviceMajor(vfsObj *vfs.VirtualFilesystem) (uint32, error) { + tpuDeviceMajorInit.Do(func() { + tpuDeviceMajor, tpuDeviceMajorInitErr = vfsObj.GetDynamicCharDevMajor() + }) + return tpuDeviceMajor, tpuDeviceMajorInitErr +} diff --git a/pkg/sentry/fsimpl/sys/pci.go b/pkg/sentry/fsimpl/sys/pci.go index 7242b8544..c1a62a179 100644 --- a/pkg/sentry/fsimpl/sys/pci.go +++ b/pkg/sentry/fsimpl/sys/pci.go @@ -175,9 +175,11 @@ func (fs *filesystem) mirrorSysDevicesDir(ctx context.Context, creds *auth.Crede if err != nil { return nil, err } + // Remove the bus prefix. + pciPath := pciBusRegex.ReplaceAllString(pciPaths[pciDeviceName], "") // Both the device and PCI address entries are links to the original PCI // device directory that's at the same place earlier in the dir tree. - linkContent = fmt.Sprintf("../../../%s", pciPaths[pciDeviceName]) + linkContent = path.Join("../../../", pciPath) case dent == "iommu_group": pciDeviceName, err := pciDeviceName(dir) if err != nil { diff --git a/pkg/sentry/vfs/device.go b/pkg/sentry/vfs/device.go index 9ccb5796b..98f66e6f8 100644 --- a/pkg/sentry/vfs/device.go +++ b/pkg/sentry/vfs/device.go @@ -110,6 +110,15 @@ func (vfs *VirtualFilesystem) ForEachDevice(cb func(pathname string, kind Device return nil } +// IsDeviceRegistered returns true if a device that matches the +// (kind, major, minor) tuple is registered. +func (vfs *VirtualFilesystem) IsDeviceRegistered(kind DeviceKind, major, minor uint32) bool { + vfs.devicesMu.RLock() + defer vfs.devicesMu.RUnlock() + _, ok := vfs.devices[devTuple{kind, major, minor}] + return ok +} + // OpenDeviceSpecialFile returns a FileDescription representing the given // device. func (vfs *VirtualFilesystem) OpenDeviceSpecialFile(ctx context.Context, mnt *Mount, d *Dentry, kind DeviceKind, major, minor uint32, opts *OpenOptions) (*FileDescription, error) { diff --git a/runsc/boot/BUILD b/runsc/boot/BUILD index bd4a4f20e..6691cb466 100644 --- a/runsc/boot/BUILD +++ b/runsc/boot/BUILD @@ -35,7 +35,6 @@ go_library( "//pkg/abi", "//pkg/abi/linux", "//pkg/abi/nvgpu", - "//pkg/abi/tpu", "//pkg/bpf", "//pkg/cleanup", "//pkg/context", diff --git a/runsc/boot/vfs.go b/runsc/boot/vfs.go index b8ac7c313..b99ff205f 100644 --- a/runsc/boot/vfs.go +++ b/runsc/boot/vfs.go @@ -28,7 +28,6 @@ import ( specs "github.com/opencontainers/runtime-spec/specs-go" "gvisor.dev/gvisor/pkg/abi/linux" "gvisor.dev/gvisor/pkg/abi/nvgpu" - "gvisor.dev/gvisor/pkg/abi/tpu" "gvisor.dev/gvisor/pkg/cleanup" "gvisor.dev/gvisor/pkg/context" "gvisor.dev/gvisor/pkg/devutil" @@ -161,10 +160,6 @@ func registerFilesystems(k *kernel.Kernel, info *containerInfo) error { return err } - if err := tpuProxyRegisterDevices(info, vfsObj); err != nil { - return err - } - return nil } @@ -1365,7 +1360,28 @@ func createDeviceFile(ctx context.Context, creds *auth.Credentials, info *contai default: return fmt.Errorf("specified device at %q has invalid type %q", devSpec.Path, devSpec.Type) } - if devSpec.Path == "/dev/nvidia-uvm" && info.nvidiaUVMDevMajor != 0 && major != info.nvidiaUVMDevMajor { + if strings.HasPrefix(devSpec.Path, "/dev/vfio") || strings.HasPrefix(devSpec.Path, "/dev/accel") { + if devSpec.Path == "/dev/vfio/vfio" { + if err := vfio.RegisterVFIODevice(vfsObj, true /* useDevGofer */); err != nil { + return fmt.Errorf("registering vfio driver: %w", err) + } + } else if tpuproxy.TPUv4DeviceRegex.MatchString(devSpec.Path) { + + if err := tpuproxy.RegisterTPUv4Device(ctx, creds, root, vfsObj, devSpec.Path, minor); err != nil { + return fmt.Errorf("registering TPUv4 device: %w", err) + } + } else if tpuproxy.TPUv5DeviceRegex.MatchString(devSpec.Path) { + if err := tpuproxy.RegisterTPUv5Device(vfsObj, devSpec.Path, minor); err != nil { + return fmt.Errorf("registering TPUv5 device: %w", err) + } + log.Infof("Switching %v device major number from %d to %d", devSpec.Path, devSpec.Major, major) + var err error + major, err = vfio.GetTPUDeviceMajor(vfsObj) + if err != nil { + return fmt.Errorf("getting TPU device major number: %w", err) + } + } + } else if devSpec.Path == "/dev/nvidia-uvm" && info.nvidiaUVMDevMajor != 0 && major != info.nvidiaUVMDevMajor { // nvidia-uvm's major device number is dynamically assigned, so the // number that it has on the host may differ from the number that // it has in sentry VFS; switch from the former to the latter. @@ -1375,25 +1391,6 @@ func createDeviceFile(ctx context.Context, creds *auth.Credentials, info *contai return dev.CreateDeviceFile(ctx, vfsObj, creds, root, devSpec.Path, major, minor, mode, devSpec.UID, devSpec.GID) } -func tpuProxyRegisterDevices(info *containerInfo, vfsObj *vfs.VirtualFilesystem) error { - if !specutils.TPUProxyIsEnabled(info.spec, info.conf) { - return nil - } - allowedTPUDeviceIDs := map[int64]any{ - tpu.TPUV4DeviceID: nil, - tpu.TPUV4liteDeviceID: nil, - tpu.TPUV5pDeviceID: nil, - tpu.TPUV5eDeviceID: nil, - } - if err := tpuproxy.RegisterHostTPUDevices(vfsObj, allowedTPUDeviceIDs); err != nil { - return fmt.Errorf("registering host TPU devices: %w", err) - } - if err := vfio.RegisterVFIODevice(vfsObj, true /* useDevGofer */); err != nil { - return fmt.Errorf("registering vfio driver: %w", err) - } - return nil -} - func nvproxyRegisterDevices(info *containerInfo, vfsObj *vfs.VirtualFilesystem) error { if !specutils.NVProxyEnabled(info.spec, info.conf) { return nil diff --git a/runsc/cmd/BUILD b/runsc/cmd/BUILD index 1b478ac88..bf4cc4f7f 100644 --- a/runsc/cmd/BUILD +++ b/runsc/cmd/BUILD @@ -79,6 +79,7 @@ go_library( visibility = ["//runsc:__subpackages__"], deps = [ "//pkg/abi/linux", + "//pkg/abi/tpu", "//pkg/cleanup", "//pkg/coretag", "//pkg/coverage", @@ -131,6 +132,7 @@ go_test( size = "small", srcs = [ "capability_test.go", + "chroot_test.go", "delete_test.go", "exec_test.go", "gofer_test.go", @@ -158,5 +160,6 @@ go_test( "@com_github_google_subcommands//:go_default_library", "@com_github_opencontainers_runtime_spec//specs-go:go_default_library", "@com_github_syndtr_gocapability//capability:go_default_library", + "@org_golang_x_sys//unix:go_default_library", ], ) diff --git a/runsc/cmd/chroot.go b/runsc/cmd/chroot.go index a561ea6ed..5912987a7 100644 --- a/runsc/cmd/chroot.go +++ b/runsc/cmd/chroot.go @@ -19,10 +19,11 @@ import ( "os" "path" "path/filepath" - "regexp" + "strings" specs "github.com/opencontainers/runtime-spec/specs-go" "golang.org/x/sys/unix" + "gvisor.dev/gvisor/pkg/abi/tpu" "gvisor.dev/gvisor/pkg/log" "gvisor.dev/gvisor/runsc/cmd/util" "gvisor.dev/gvisor/runsc/config" @@ -113,7 +114,7 @@ func setUpChroot(spec *specs.Spec, conf *config.Config) error { return fmt.Errorf("error mounting proc in chroot: %v", err) } - if err := tpuProxyUpdateChroot(chroot, spec, conf); err != nil { + if err := tpuProxyUpdateChroot("/", chroot, spec, conf); err != nil { return fmt.Errorf("error configuring chroot for TPU devices: %w", err) } @@ -124,93 +125,59 @@ func setUpChroot(spec *specs.Spec, conf *config.Config) error { return pivotRoot(chroot) } -// Mount the path that dest points to for TPU at chroot, the mounted path is returned in absolute form. -func mountTPUSyslinkInChroot(chroot, dest, relativePath string, validator func(link string) bool) (string, error) { - src, err := os.Readlink(dest) - if err != nil { - return "", fmt.Errorf("error reading %v: %v", src, err) - } - // Ensure the link is in the form we expect. - if !validator(src) { - return "", fmt.Errorf("unexpected link %q -> %q", dest, src) - } - path, err := filepath.Abs(path.Join(filepath.Dir(dest), src, relativePath)) - if err != nil { - return "", fmt.Errorf("error parsing path %q: %v", src, err) - } - if err := mountInChroot(chroot, path, path, "bind", unix.MS_BIND|unix.MS_RDONLY); err != nil { - return "", fmt.Errorf("error mounting %q in chroot: %v", dest, err) - } - return path, nil -} - -func mountTPUDeviceInfoInChroot(chroot, devicePath, sysfsFormat, pciDeviceFormat string) error { - deviceMinor, valid, err := util.ExtractTPUDeviceMinor(devicePath) - if err != nil { - return fmt.Errorf("extracting TPU device minor: %w", err) - } - if !valid { - return nil - } - // Multiple paths link to the /sys/devices// - // directory that contains all relevant sysfs accel/vfio 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///accel/accel# - // * or /sys/devices///vfio-dev/vfio# for VFIO-based TPU - // and traversing up 2 directories. - // The sysDevicePath itself is a soft link to the device directory. - sysDevicePath := fmt.Sprintf(sysfsFormat, deviceMinor) - sysPCIDeviceDir, err := mountTPUSyslinkInChroot(chroot, sysDevicePath, "../..", func(link string) bool { - sysDeviceLinkMatcher := regexp.MustCompile(fmt.Sprintf(pciDeviceFormat, deviceMinor)) - return sysDeviceLinkMatcher.MatchString(link) - }) - if err != nil { - return err - } - - // Mount the device's IOMMU group if available. - iommuGroupPath := path.Join(sysPCIDeviceDir, "iommu_group") - if _, err := os.Stat(iommuGroupPath); err == nil { - if _, err := mountTPUSyslinkInChroot(chroot, iommuGroupPath, "", func(link string) bool { - iommuGroupPathMatcher := regexp.MustCompile(`../../../kernel/iommu_groups/\d+`) - return iommuGroupPathMatcher.MatchString(link) - }); err != nil { - return err - } - } - return nil -} - -func tpuProxyUpdateChroot(chroot string, spec *specs.Spec, conf *config.Config) error { +func tpuProxyUpdateChroot(hostRoot, chroot string, spec *specs.Spec, conf *config.Config) error { if !specutils.TPUProxyIsEnabled(spec, conf) { return nil } - // When a path glob is added to pathGlobToSysfsFormat, the corresponding pciDeviceFormat has to be added to pathGlobToPciDeviceFormat. - pathGlobToSysfsFormat := map[string]string{ - "/dev/accel*": "/sys/class/accel/accel%d", - "/dev/vfio/*": "/sys/class/vfio-dev/vfio%d"} - pathGlobToPciDeviceFormat := map[string]string{ - "/dev/accel*": `../../devices/pci0000:[[:xdigit:]]{2}/(\d+:\d+:\d+\.\d+)/accel/accel%d`, - "/dev/vfio/*": `../../devices/pci0000:[[:xdigit:]]{2}/(\d+:\d+:\d+\.\d+)/vfio-dev/vfio%d`} - // Bind mount device info directories for all TPU devices on the host. - // For v4 TPU, the directory /sys/devices///accel/accel# is mounted; - // For v5e TPU, the directory /sys/devices///vfio-dev/vfio# is mounted. - foundDevices := false - for pathGlob, sysfsFormat := range pathGlobToSysfsFormat { - paths, err := filepath.Glob(pathGlob) + allowedDeviceIDs := map[uint64]struct{}{} + paths, err := filepath.Glob(path.Join(hostRoot, "dev/vfio/*")) + if err != nil { + return fmt.Errorf("enumerating TPU device files: %w", err) + } + vfioDevicePath := path.Join(hostRoot, "dev/vfio/vfio") + for _, devPath := range paths { + if devPath == vfioDevicePath { + continue + } + devNum := path.Base(devPath) + iommuGroupPath := path.Join("/sys/kernel/iommu_groups", devNum) + if err := mountInChroot(chroot, path.Join(hostRoot, iommuGroupPath), iommuGroupPath, "bind", unix.MS_BIND|unix.MS_RDONLY); err != nil { + return fmt.Errorf("error mounting %q in chroot: %v", iommuGroupPath, err) + } + allowedDeviceIDs[tpu.TPUV5pDeviceID] = struct{}{} + allowedDeviceIDs[tpu.TPUV5eDeviceID] = struct{}{} + } + if len(allowedDeviceIDs) == 0 { + paths, err = filepath.Glob(path.Join(hostRoot, "dev/accel*")) if err != nil { return fmt.Errorf("enumerating TPU device files: %w", err) } - for _, devPath := range paths { - foundDevices = true - if err := mountTPUDeviceInfoInChroot(chroot, devPath, sysfsFormat, pathGlobToPciDeviceFormat[pathGlob]); err != nil { - return err + if len(paths) == 0 { + return fmt.Errorf("could not find any TPU devices on the host") + } + allowedDeviceIDs[tpu.TPUV4DeviceID] = struct{}{} + allowedDeviceIDs[tpu.TPUV4liteDeviceID] = struct{}{} + } + if len(allowedDeviceIDs) == 0 { + return fmt.Errorf("no TPU devices found on the host") + } + sysDevicesGlob := path.Join(hostRoot, "/sys/devices/pci*") + sysDevicesPaths, err := filepath.Glob(sysDevicesGlob) + if err != nil { + return fmt.Errorf("enumerating PCI device files: %w", err) + } + for _, sysDevicesPath := range sysDevicesPaths { + if err := filepath.WalkDir(sysDevicesPath, func(path string, d os.DirEntry, err error) error { + if d.Type().IsDir() && util.IsPCIDeviceDirTPU(path, allowedDeviceIDs) { + chrootPath := strings.Replace(path, hostRoot, "/", 1) + if err := mountInChroot(chroot, path, chrootPath, "bind", unix.MS_BIND|unix.MS_RDONLY); err != nil { + return fmt.Errorf("error mounting %q in chroot: %v", path, err) + } } + return nil + }); err != nil { + return fmt.Errorf("walking %q: %w", sysDevicesPath, err) } } - if !foundDevices { - return fmt.Errorf("could not find any TPU devices on the host") - } - return nil + return err } diff --git a/runsc/cmd/chroot_test.go b/runsc/cmd/chroot_test.go new file mode 100644 index 000000000..93f01d24f --- /dev/null +++ b/runsc/cmd/chroot_test.go @@ -0,0 +1,155 @@ +// Copyright 2024 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 cmd + +import ( + "fmt" + "os" + "path" + "testing" + + specs "github.com/opencontainers/runtime-spec/specs-go" + "golang.org/x/sys/unix" + "gvisor.dev/gvisor/runsc/config" +) + +func setup(t *testing.T) (string, string) { + t.Helper() + testDir := t.TempDir() + gvisorChroot := path.Join(testDir, "gvisor_chroot") + os.Mkdir(gvisorChroot, 0755) + + // Mounting the gvisor chroot makes the submounts easier to cleanup at the + // end of the test. + if err := unix.Mount(gvisorChroot, gvisorChroot, "", unix.MS_BIND, ""); err != nil { + t.Fatalf("failed to bind mount gvisor chroot: %v", err) + } + t.Cleanup(func() { + if err := unix.Unmount(gvisorChroot, unix.MNT_DETACH); err != nil { + t.Fatalf("failed to unmount gvisor chroot: %v", err) + } + }) + return testDir, gvisorChroot +} + +func TestTPUProxyV5(t *testing.T) { + testDir, gvisorChroot := setup(t) + + for i := 0; i < 3; i++ { + os.MkdirAll(path.Join(testDir, "sys", "kernel", "iommu_groups", fmt.Sprintf("%d", i)), 0755) + writeFile(t, path.Join(testDir, "dev", "vfio", fmt.Sprintf("%d", i)), "") + pciPath := path.Join(testDir, "sys", "devices", "pci0000:00", fmt.Sprintf("0000:00:00.%d", i)) + writeFile(t, path.Join(pciPath, "device"), "0x0062") + writeFile(t, path.Join(pciPath, "vendor"), "0x1ae0") + } + + if err := tpuProxyUpdateChroot(testDir, gvisorChroot, &specs.Spec{}, &config.Config{TPUProxy: true}); err != nil { + t.Fatalf("failed to update chroot: %v", err) + } + + for i := 0; i < 3; i++ { + if _, err := os.Stat(path.Join(gvisorChroot, "sys", "kernel", "iommu_groups", fmt.Sprintf("%d", i))); err != nil { + t.Errorf("failed to stat iommu group file: %v", err) + } + devicePath := path.Join(gvisorChroot, "sys", "devices", "pci0000:00", fmt.Sprintf("0000:00:00.%d", i), "device") + if _, err := os.ReadFile(devicePath); err != nil { + t.Errorf("failed to read device file: %v", err) + } + vendorPath := path.Join(gvisorChroot, "sys", "devices", "pci0000:00", fmt.Sprintf("0000:00:00.%d", i), "vendor") + if _, err := os.ReadFile(vendorPath); err != nil { + t.Errorf("failed to read device file: %v", err) + } + } +} + +func TestTPUProxyV5NestedPCIDevice(t *testing.T) { + testDir, gvisorChroot := setup(t) + + for i := 0; i < 3; i++ { + os.MkdirAll(path.Join(testDir, "sys", "kernel", "iommu_groups", fmt.Sprintf("%d", i)), 0755) + writeFile(t, path.Join(testDir, "dev", "vfio", fmt.Sprintf("%d", i)), "") + pciPath := path.Join(testDir, "sys", "devices", "pci0000:00", fmt.Sprintf("0000:00:00.%d", i)) + writeFile(t, path.Join(pciPath, "device"), "0x0062") + writeFile(t, path.Join(pciPath, "vendor"), "0x1ae0") + } + + nestedDeviceNum := 3 + writeFile(t, path.Join(testDir, "dev", "vfio", fmt.Sprintf("%d", nestedDeviceNum)), "") + os.MkdirAll(path.Join(testDir, "sys", "kernel", "iommu_groups", fmt.Sprintf("%d", nestedDeviceNum)), 0755) + pciPath := path.Join(testDir, "sys", "devices", "pci0000:00", "0000:00:00.2", "0000:00:00.3.0") + writeFile(t, path.Join(pciPath, "device"), "0x0062") + writeFile(t, path.Join(pciPath, "vendor"), "0x1ae0") + + if err := tpuProxyUpdateChroot(testDir, gvisorChroot, &specs.Spec{}, &config.Config{TPUProxy: true}); err != nil { + t.Fatalf("failed to update chroot: %v", err) + } + + if _, err := os.Stat(path.Join(gvisorChroot, "sys", "kernel", "iommu_groups", fmt.Sprintf("%d", nestedDeviceNum))); err != nil { + t.Errorf("failed to stat iommu group file: %v", err) + } + devicePath := path.Join(gvisorChroot, "sys", "devices", "pci0000:00", "0000:00:00.2", "0000:00:00.3.0", "device") + if _, err := os.ReadFile(devicePath); err != nil { + t.Errorf("failed to read device file: %v", err) + } +} + +func TestTPUProxyV4(t *testing.T) { + testDir, gvisorChroot := setup(t) + + for i := 0; i < 3; i++ { + os.MkdirAll(path.Join(testDir, "sys", "kernel", "iommu_groups", fmt.Sprintf("%d", i)), 0755) + writeFile(t, path.Join(testDir, "dev", fmt.Sprintf("accel%d", i)), "") + pciPath := path.Join(testDir, "sys", "devices", "pci0000:00", fmt.Sprintf("0000:00:00.%d", i)) + writeFile(t, path.Join(pciPath, "device"), "0x005e") + writeFile(t, path.Join(pciPath, "vendor"), "0x1ae0") + } + + if err := tpuProxyUpdateChroot(testDir, gvisorChroot, &specs.Spec{}, &config.Config{TPUProxy: true}); err != nil { + t.Fatalf("failed to update chroot: %v", err) + } + + for i := 0; i < 3; i++ { + devicePath := path.Join(gvisorChroot, "sys", "devices", "pci0000:00", fmt.Sprintf("0000:00:00.%d", i), "device") + if _, err := os.ReadFile(devicePath); err != nil { + t.Errorf("failed to read device file: %v", err) + } + vendorPath := path.Join(gvisorChroot, "sys", "devices", "pci0000:00", fmt.Sprintf("0000:00:00.%d", i), "vendor") + if _, err := os.ReadFile(vendorPath); err != nil { + t.Errorf("failed to read device file: %v", err) + } + } +} + +func writeFile(t *testing.T, fpath string, contents string) { + t.Helper() + dir := path.Dir(fpath) + if st, err := os.Stat(dir); err != nil { + if os.IsNotExist(err) { + if err := os.MkdirAll(dir, 0755); err != nil { + t.Fatalf("failed to create directory: %v", err) + } + } else { + t.Fatalf("failed to stat directory: %v", err) + } + } else if !st.IsDir() { + t.Fatalf("path %q is not a directory", dir) + } + f, err := os.Create(fpath) + if err != nil { + t.Fatalf("failed to create file: %v", err) + } + f.WriteString(contents) + f.Close() +} diff --git a/runsc/cmd/gofer.go b/runsc/cmd/gofer.go index 7761d5740..d92e4bdb2 100644 --- a/runsc/cmd/gofer.go +++ b/runsc/cmd/gofer.go @@ -551,7 +551,7 @@ func shouldExposeVFIODevice(path string) bool { // // Precondition: tpuproxy is enabled. func shouldExposeTpuDevice(path string) bool { - _, valid, _ := util.ExtractTPUDeviceMinor(path) + valid, _ := util.IsTPUDeviceValid(path) return valid || shouldExposeVFIODevice(path) } diff --git a/runsc/cmd/util/BUILD b/runsc/cmd/util/BUILD index fd57bd66c..2d55017dd 100644 --- a/runsc/cmd/util/BUILD +++ b/runsc/cmd/util/BUILD @@ -20,6 +20,5 @@ go_library( "//pkg/abi/tpu", "//pkg/log", "@com_github_google_subcommands//:go_default_library", - "@org_golang_x_sys//unix:go_default_library", ], ) diff --git a/runsc/cmd/util/tpu.go b/runsc/cmd/util/tpu.go index 95e170660..8c00b313e 100644 --- a/runsc/cmd/util/tpu.go +++ b/runsc/cmd/util/tpu.go @@ -17,89 +17,132 @@ package util import ( "fmt" "os" + "path" + "path/filepath" "regexp" "strconv" "strings" - "syscall" - "golang.org/x/sys/unix" "gvisor.dev/gvisor/pkg/abi/tpu" ) const ( - googleVendorID = 0x1AE0 - accelDevicePathRegex = `^/dev/accel(\d+)$` - accelSysfsFormat = "/sys/class/accel/accel%d/device/%s" - vfioDevicePathRegex = `^/dev/vfio/(\d+)$` - vfioSysfsFormat = "/sys/class/vfio-dev/vfio%d/device/%s" - vendorFile = "vendor" - deviceFile = "device" + googleVendorID = 0x1AE0 + accelDevicePathRegex = `^/dev/accel(\d+)$` + accelSysfsFormat = "/sys/class/accel/accel%d/device/%s" + vfioDevicePathRegex = `^/dev/vfio/(\d+)$` + iommuGroupSysfsGlobFormat = "/sys/kernel/iommu_groups/%s/devices/*" + vendorFile = "vendor" + deviceFile = "device" + pciAddressMaxLength = 13 ) -var tpuV4DeviceIDs = map[uint64]any{tpu.TPUV4DeviceID: nil, tpu.TPUV4liteDeviceID: nil} -var tpuV5DeviceIDs = map[uint64]any{tpu.TPUV5eDeviceID: nil, tpu.TPUV5pDeviceID: nil} +var ( + tpuV4DeviceIDs = map[uint64]struct{}{tpu.TPUV4DeviceID: struct{}{}, tpu.TPUV4liteDeviceID: struct{}{}} + tpuV5DeviceIDs = map[uint64]struct{}{tpu.TPUV5eDeviceID: struct{}{}, tpu.TPUV5pDeviceID: struct{}{}} + pciDeviceRegex = regexp.MustCompile(`0000:([[:xdigit:]]{2}|[[:xdigit:]]{4}):[[:xdigit:]]{2}\.[[:xdigit:]]{1,2}`) +) -// 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. -func ExtractTPUDeviceMinor(path string) (uint32, bool, error) { - devNum, valid, err := tpuV4DeviceMinor(path) +// IsPCIDeviceDirTPU returns if the given PCI device sysfs path is a TPU device +// with one of the allowed device IDs. +func IsPCIDeviceDirTPU(sysfsPath string, allowedDeviceIDs map[uint64]struct{}) bool { + dir := path.Base(sysfsPath) + if !pciDeviceRegex.MatchString(dir) || len(dir) > pciAddressMaxLength { + return false + } + vendor, err := readHexInt(path.Join(sysfsPath, vendorFile)) if err != nil { - return 0, false, err - } - if valid { - return devNum, valid, err - } - return tpuV5DeviceMinor(path) -} - -// tpuDeviceMinor 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. -func tpuDeviceMinor(devicePath, devicePathRegex, sysfsFormat string, allowedDeviceIDs map[uint64]any) (uint32, bool, error) { - deviceRegex := regexp.MustCompile(devicePathRegex) - matches := deviceRegex.FindStringSubmatch(devicePath) - if matches == nil { - return 0, false, nil - } - var st syscall.Stat_t - if err := syscall.Stat(devicePath, &st); err != nil { - return 0, false, err - } - minor := unix.Minor(st.Rdev) - vendor, err := readHexInt(fmt.Sprintf(sysfsFormat, minor, vendorFile)) - if err != nil { - return 0, false, err + return false } if vendor != googleVendorID { - return 0, false, nil + return false } - deviceID, err := readHexInt(fmt.Sprintf(sysfsFormat, minor, deviceFile)) + deviceID, err := readHexInt(path.Join(sysfsPath, deviceFile)) if err != nil { - return 0, false, err + return false } if _, ok := allowedDeviceIDs[deviceID]; !ok { - return 0, false, nil + return false } - return minor, true, nil + return true } -// tpuv4DeviceMinor returns v4 and v4lite TPU device minor number for the given path. +// IsTPUDeviceValid returns if the accelerator device is valid. +func IsTPUDeviceValid(path string) (bool, error) { + valid, err := tpuV4DeviceValid(path) + if err != nil { + return false, err + } + if valid { + return valid, err + } + return tpuV5DeviceValid(path) +} + +// tpuV4DeviceValid returns v4 and v4lite TPU device minor number for the given path. // A valid v4 TPU device is defined as: // * Path is /dev/accel#. // * Vendor is googleVendorID. // * Device ID is one of tpuV4DeviceIDs. -func tpuV4DeviceMinor(path string) (uint32, bool, error) { - return tpuDeviceMinor(path, accelDevicePathRegex, accelSysfsFormat, tpuV4DeviceIDs) +func tpuV4DeviceValid(devPath string) (bool, error) { + deviceRegex := regexp.MustCompile(accelDevicePathRegex) + matches := deviceRegex.FindStringSubmatch(devPath) + if matches == nil { + return false, nil + } + if len(matches) < 1 { + return false, fmt.Errorf("found %d matches for %s", len(matches), devPath) + } + devNum, err := strconv.ParseUint(matches[1], 10, 32) + if err != nil { + return false, err + } + vendor, err := readHexInt(fmt.Sprintf(accelSysfsFormat, devNum, vendorFile)) + if err != nil { + return false, err + } + if vendor != googleVendorID { + return false, nil + } + deviceID, err := readHexInt(fmt.Sprintf(accelSysfsFormat, devNum, deviceFile)) + if err != nil { + return false, err + } + if _, ok := tpuV4DeviceIDs[deviceID]; !ok { + return false, nil + } + return true, nil } -// tpuV5DeviceMinor returns the v5e TPU device minor number for te given path. +// tpuV5DeviceValid returns the v5e TPU device minor number for te given path. // A valid v5 TPU device is defined as: // * Path is /dev/vfio/#. // * Vendor is googleVendorID. // * Device ID is one of tpuV5DeviceIDs. -func tpuV5DeviceMinor(path string) (uint32, bool, error) { - return tpuDeviceMinor(path, vfioDevicePathRegex, vfioSysfsFormat, tpuV5DeviceIDs) +func tpuV5DeviceValid(devPath string) (bool, error) { + paths, err := filepath.Glob(fmt.Sprintf(iommuGroupSysfsGlobFormat, path.Base(devPath))) + if err != nil { + return false, err + } + if len(paths) != 1 { + return false, fmt.Errorf("found %d paths for %s", len(paths), devPath) + } + sysfsPath := paths[0] + vendor, err := readHexInt(path.Join(sysfsPath, vendorFile)) + if err != nil { + return false, err + } + if vendor != googleVendorID { + return false, nil + } + deviceID, err := readHexInt(path.Join(sysfsPath, deviceFile)) + if err != nil { + return false, err + } + if _, ok := tpuV5DeviceIDs[deviceID]; !ok { + return false, nil + } + return true, nil } func readHexInt(path string) (uint64, error) {