chore(integration): test with nvidia sample CUDA image

johnstcn · johnstcn · commit 669f1e5e92a2 · 2025-03-06T23:16:37.000Z
diff --git a/integration/gpu_test.go b/integration/gpu_test.go
@@ -62,8 +62,10 @@ func TestDocker_Nvidia(t *testing.T) {
 		// Assert that we can run nvidia-smi in the inner container.
 		assertInnerNvidiaSMI(ctx, t, ctID)
 
-		// Make sure dnf still works.
-		out, err := execContainerCmd(ctx, t, ctID, "docker", "exec", "workspace_cvm", "dnf", "list")
+		// Make sure dnf still works. This checks for a regression due to
+		// gpuExtraRegex matching `libglib.so` in the outer container.
+		// This had a dependency on `libpcre.so.3` which would cause dnf to fail.
+		out, err := execContainerCmd(ctx, t, ctID, "docker", "exec", "workspace_cvm", "dnf")
 		if !assert.NoError(t, err, "failed to run dnf in the inner container") {
 			t.Logf("dnf output:\n%s", strings.TrimSpace(out))
 		}
@@ -113,6 +115,29 @@ func TestDocker_Nvidia(t *testing.T) {
 		assertInnerFiles(ctx, t, ctID, "/usr/lib/x86_64-linux-gnu/libnv*", ofs...)
 		assertInnerNvidiaSMI(ctx, t, ctID)
 	})
+
+	t.Run("CUDASample", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithCancel(context.Background())
+		t.Cleanup(cancel)
+
+		// Start the envbox container.
+		ctID := startEnvboxCmd(ctx, t, integrationtest.CUDASampleImage, "root",
+			"-v", "/usr/lib/x86_64-linux-gnu:/var/coder/usr/lib",
+			"--env", "CODER_ADD_GPU=true",
+			"--env", "CODER_USR_LIB_DIR=/var/coder/usr/lib",
+			"--runtime=nvidia",
+			"--gpus=all",
+		)
+
+		// Assert that we can run nvidia-smi in the inner container.
+		assertInnerNvidiaSMI(ctx, t, ctID)
+
+		// Assert that /tmp/vectorAdd runs successfully in the inner container.
+		_, err := execContainerCmd(ctx, t, ctID, "docker", "exec", "workspace_cvm", "/tmp/vectorAdd")
+		require.NoError(t, err, "failed to run /tmp/vectorAdd in the inner container")
+	})
 }
 
 // dockerRuntimes returns the list of container runtimes available on the host.
diff --git a/integration/integrationtest/docker.go b/integration/integrationtest/docker.go
@@ -42,6 +42,9 @@ const (
 	UbuntuImage = "gcr.io/coder-dev-1/sreya/ubuntu-coder"
 	// Redhat UBI9 image as of 2025-03-05
 	RedhatImage = "registry.access.redhat.com/ubi9/ubi:9.5"
+	// CUDASampleImage is a CUDA sample image from NVIDIA's container registry.
+	// It contains a binary /tmp/vectorAdd which can be run to test the CUDA setup.
+	CUDASampleImage = "nvcr.io/nvidia/k8s/cuda-sample:vectoradd-cuda10.2"
 
 	// RegistryImage is used to assert that we add certs
 	// correctly to the docker daemon when pulling an image
