Function TestRunGPU

cmd/nerdctl/container/container_run_linux_test.go:728–828

TestRunGPU tests GPU injection using the --gpus flag.

(t *testing.T)

// TestRunGPU tests GPU injection using the --gpus flag.
func TestRunGPU(t *testing.T) {
	t.Parallel()
	// Although CDI injection is supported by Docker, specifying the --cdi-spec-dirs on the command line is not.
	testutil.DockerIncompatible(t)
	const nvidiaSpec = `
cdiVersion: "0.5.0"
kind: "nvidia.com/gpu"
devices:
- name: "0"
  containerEdits:
    env:
    - NVIDIA_GPU_0=injected
- name: "1"
  containerEdits:
    env:
    - NVIDIA_GPU_1=injected
`
	const amdSpec = `
cdiVersion: "0.5.0"
kind: "amd.com/gpu"
devices:
- name: "0"
  containerEdits:
    env:
    - AMD_GPU_0=injected
- name: "1"
  containerEdits:
    env:
    - AMD_GPU_1=injected
`
	const unknownSpec = `
cdiVersion: "0.5.0"
kind: "unknown.com/gpu"
devices:
- name: "0"
  containerEdits:
    env:
    - UNKNOWN_GPU_0=injected
`

	testCases := []struct {
		name         string
		specs        map[string]string
		gpuFlags     []string
		expectedEnvs []string
		expectFail   bool
	}{
		{
			name:         "nvidia device injection",
			specs:        map[string]string{"nvidia.yaml": nvidiaSpec},
			gpuFlags:     []string{"--gpus", "2"},
			expectedEnvs: []string{"NVIDIA_GPU_0=injected", "NVIDIA_GPU_1=injected"},
		},
		{
			name:         "amd device injection",
			specs:        map[string]string{"amd.yaml": amdSpec},
			gpuFlags:     []string{"--gpus", "2"},
			expectedEnvs: []string{"AMD_GPU_0=injected", "AMD_GPU_1=injected"},