Support TPU v6e w/ gVisor.

PiperOrigin-RevId: 687514246
google · Oct 19, 2024 · 17edc72 · 17edc72
1 parent 0b59173
commit 17edc72
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/pkg/abi/tpu/tpu.go b/pkg/abi/tpu/tpu.go
@@ -44,6 +44,9 @@ const (
 
 	// TPUV5pDeviceID is the PCI device ID of TPU V5p hardware.
 	TPUV5pDeviceID = 0x0062
+
+	// TPUV6eDeviceID is the PCI device ID of TPU V6e hardware.
+	TPUV6eDeviceID = 0x006f
 )
 
 // TPUV4InterruptsMap maps BAR indices to valid register offsets.

diff --git a/runsc/cmd/chroot.go b/runsc/cmd/chroot.go
@@ -146,6 +146,7 @@ func tpuProxyUpdateChroot(hostRoot, chroot string, spec *specs.Spec, conf *confi
 		}
 		allowedDeviceIDs[tpu.TPUV5pDeviceID] = struct{}{}
 		allowedDeviceIDs[tpu.TPUV5eDeviceID] = struct{}{}
+		allowedDeviceIDs[tpu.TPUV6eDeviceID] = struct{}{}
 	}
 	if len(allowedDeviceIDs) == 0 {
 		paths, err = filepath.Glob(path.Join(hostRoot, "dev/accel*"))