From 641a1a56b8c59bc4c1507d98adb6e93ea79e4955 Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Thu, 21 Mar 2024 14:01:29 -0700
Subject: [PATCH] `dockerutil.GPURunOpts`: Expose all GPUs to containers, not
 just the first.

PiperOrigin-RevId: 617956471
---
 pkg/test/dockerutil/gpu.go | 28 ++++++++++++++++++++++++----
 1 file changed, 24 insertions(+), 4 deletions(-)

diff --git a/pkg/test/dockerutil/gpu.go b/pkg/test/dockerutil/gpu.go
index 76c8de32f..f7370a0cf 100644
--- a/pkg/test/dockerutil/gpu.go
+++ b/pkg/test/dockerutil/gpu.go
@@ -17,6 +17,7 @@ package dockerutil
 
 import (
 	"flag"
+	"fmt"
 	"os"
 
 	"github.com/docker/docker/api/types/container"
@@ -45,10 +46,16 @@ func GPURunOpts() RunOpts {
 	// COS has specific settings since it has a custom installer for GPU drivers.
 	// See: https://cloud.google.com/container-optimized-os/docs/how-to/run-gpus#install-driver
 	devices := []container.DeviceMapping{}
-	nvidia0Device := "/dev/nvidia0"
-	nvidiaUvmDevice := "/dev/nvidia-uvm"
-	nvidiactlDevice := "/dev/nvidiactl"
-	for _, device := range []string{nvidia0Device, nvidiaUvmDevice, nvidiactlDevice} {
+	var nvidiaDevices []string
+	for i := 0; true; i++ {
+		devicePath := fmt.Sprintf("/dev/nvidia%d", i)
+		if _, err := os.Stat(devicePath); err != nil {
+			break
+		}
+		nvidiaDevices = append(nvidiaDevices, devicePath)
+	}
+	nvidiaDevices = append(nvidiaDevices, "/dev/nvidia-uvm", "/dev/nvidiactl")
+	for _, device := range nvidiaDevices {
 		devices = append(devices, container.DeviceMapping{
 			PathOnHost:        device,
 			PathInContainer:   device,
@@ -89,3 +96,16 @@ func GPURunOpts() RunOpts {
 		Devices: devices,
 	}
 }
+
+// NumGPU crudely estimates the number of NVIDIA GPUs on the host.
+func NumGPU() int {
+	numGPU := 0
+	for {
+		_, err := os.Stat(fmt.Sprintf("/dev/nvidia%d", numGPU))
+		if err != nil {
+			break
+		}
+		numGPU++
+	}
+	return numGPU
+}