Extract GPU device id from container runtime

Author: zhuminyi

comp/core/workloadmeta/collectors/internal/containerd/container_builder.go

@@ -154,6 +154,7 @@ func buildWorkloadMetaContainer(namespace string, container containerd.Container
 
 		workloadContainer.EnvVars = envs
 		workloadContainer.Hostname = spec.Hostname
+		workloadContainer.GPUDeviceIDs = util.ExtractGPUDeviceIDsFromEnvMap(envs)
 		if spec.Linux != nil {
 			workloadContainer.CgroupPath = extractCgroupPath(spec.Linux.CgroupsPath)
 		}

comp/core/workloadmeta/collectors/internal/docker/docker.go

@@ -45,13 +45,6 @@ import (
 const (
 	collectorID   = "docker"
 	componentName = "workloadmeta-docker"
-
-	// nvidiaVisibleDevicesEnvVar is the environment variable set by NVIDIA container runtime
-	// to specify which GPUs are visible to the container. Values can be:
-	// - GPU UUIDs: "GPU-uuid" or "GPU-uuid1,GPU-uuid2" (ECS, some K8s setups)
-	// - Device indices: "0", "1", "0,1" (local Docker)
-	// - Special values: "all", "none", "void"
-	nvidiaVisibleDevicesEnvVar = "NVIDIA_VISIBLE_DEVICES"
 )
 
 // imageEventActionSbom is an event that we set to create a fake docker event.
@@ -336,7 +329,7 @@ func (c *collector) buildCollectorEvent(ctx context.Context, ev *docker.Containe
 			PID:          container.State.Pid,
 			RestartCount: container.RestartCount,
 			Resources:    extractResources(container),
-			GPUDeviceIDs: extractGPUDeviceIDsForECS(container.Config.Env),
+			GPUDeviceIDs: util.ExtractGPUDeviceIDsFromEnvVars(container.Config.Env),
 		}
 
 	case events.ActionDie, docker.ActionDied:
@@ -747,39 +740,6 @@ func layersFromDockerHistoryAndInspect(history []image.HistoryResponseItem, insp
 	return layers
 }
 
-// extractGPUDeviceIDsForECS extracts GPU device identifiers from NVIDIA_VISIBLE_DEVICES environment variable,
-// but ONLY when running in ECS. For regular Docker containers, the NVIDIA container toolkit adds
-// NVIDIA_VISIBLE_DEVICES in a way that's not visible in container.Config.Env (it's added by the
-// runtime, not the container config), so we must rely on reading from procfs at metric collection time.
-// In ECS, the env var IS visible in container.Config.Env because ECS sets it directly.
-// ECS typically sets GPU UUIDs (e.g., "GPU-uuid1,GPU-uuid2"), but users can also set "all" for GPU sharing.
-func extractGPUDeviceIDsForECS(envVars []string) []string {
-	// Only extract from container config in ECS.
-	// For regular Docker, NVIDIA_VISIBLE_DEVICES is added by the container runtime
-	// and won't be visible here - the GPU probe will read it from procfs instead.
-	if !env.IsECS() {
-		return nil
-	}
-	return extractGPUDeviceIDs(envVars)
-}
-
-// extractGPUDeviceIDs parses GPU device identifiers from NVIDIA_VISIBLE_DEVICES environment variable.
-// ECS typically sets GPU UUIDs (e.g., "GPU-uuid1,GPU-uuid2"), but users can also set "all" for GPU sharing.
-// Special values "all", "none", "void" are preserved and handled in matchByGPUDeviceIDs().
-// Empty value returns nil (env var set but empty).
-func extractGPUDeviceIDs(envVars []string) []string {
-	prefix := nvidiaVisibleDevicesEnvVar + "="
-	for _, e := range envVars {
-		if value, found := strings.CutPrefix(e, prefix); found {
-			if value == "" {
-				return nil
-			}
-			return strings.Split(value, ",")
-		}
-	}
-	return nil
-}
-
 func extractResources(container container.InspectResponse) workloadmeta.ContainerResources {
 	var resources workloadmeta.ContainerResources
 

comp/core/workloadmeta/collectors/internal/docker/docker_test.go

@@ -18,62 +18,6 @@ import (
 	workloadmeta "github.com/DataDog/datadog-agent/comp/core/workloadmeta/def"
 )
 
-func Test_extractGPUDeviceIDs(t *testing.T) {
-	tests := []struct {
-		name     string
-		envVars  []string
-		expected []string
-	}{
-		{
-			name:     "single GPU UUID",
-			envVars:  []string{"PATH=/usr/bin", "NVIDIA_VISIBLE_DEVICES=GPU-aec058b1-c18e-236e-c14d-49d2990fda0f"},
-			expected: []string{"GPU-aec058b1-c18e-236e-c14d-49d2990fda0f"},
-		},
-		{
-			name:     "multiple GPU UUIDs",
-			envVars:  []string{"NVIDIA_VISIBLE_DEVICES=GPU-aec058b1-c18e-236e-c14d-49d2990fda0f,GPU-bec058b1-d18e-336e-d14d-59d2990fda1f"},
-			expected: []string{"GPU-aec058b1-c18e-236e-c14d-49d2990fda0f", "GPU-bec058b1-d18e-336e-d14d-59d2990fda1f"},
-		},
-		{
-			name:     "all GPUs",
-			envVars:  []string{"NVIDIA_VISIBLE_DEVICES=all"},
-			expected: []string{"all"},
-		},
-		{
-			name:     "none",
-			envVars:  []string{"NVIDIA_VISIBLE_DEVICES=none"},
-			expected: []string{"none"},
-		},
-		{
-			name:     "void",
-			envVars:  []string{"NVIDIA_VISIBLE_DEVICES=void"},
-			expected: []string{"void"},
-		},
-		{
-			name:     "empty value",
-			envVars:  []string{"NVIDIA_VISIBLE_DEVICES="},
-			expected: nil,
-		},
-		{
-			name:     "no NVIDIA_VISIBLE_DEVICES",
-			envVars:  []string{"PATH=/usr/bin", "HOME=/root"},
-			expected: nil,
-		},
-		{
-			name:     "empty env vars",
-			envVars:  []string{},
-			expected: nil,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			result := extractGPUDeviceIDs(tt.envVars)
-			assert.Equal(t, tt.expected, result)
-		})
-	}
-}
-
 func Test_LayersFromDockerHistoryAndInspect(t *testing.T) {
 	var emptySize int64
 	var noDiffCmd = "ENV var=dummy"

comp/core/workloadmeta/collectors/util/gpu_util.go

@@ -0,0 +1,82 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016-present Datadog, Inc.
+
+package util
+
+import (
+	"strings"
+
+	"github.com/DataDog/datadog-agent/pkg/config/env"
+)
+
+// NVIDIAVisibleDevicesEnvVar is the environment variable set by NVIDIA container runtime
+// or Kubernetes device plugin to specify which GPUs are visible to the container.
+// Values can be GPU UUIDs (e.g., "GPU-uuid1,GPU-uuid2"), special values like "all", "none", "void",
+// or MIG instance identifiers.
+const NVIDIAVisibleDevicesEnvVar = "NVIDIA_VISIBLE_DEVICES"
+
+// ExtractGPUDeviceIDs parses GPU device identifiers from NVIDIA_VISIBLE_DEVICES environment variable.
+// This is the core parsing function that extracts GPU IDs from a list of environment variables.
+//
+// Returns:
+//   - GPU UUIDs as a slice (e.g., ["GPU-uuid1", "GPU-uuid2"])
+//   - Special values like ["all"], ["none"], ["void"] are preserved
+//   - nil if the env var is not set or has an empty value
+func ExtractGPUDeviceIDs(envVars []string) []string {
+	prefix := NVIDIAVisibleDevicesEnvVar + "="
+	for _, e := range envVars {
+		if value, found := strings.CutPrefix(e, prefix); found {
+			if value == "" {
+				return nil
+			}
+			return strings.Split(value, ",")
+		}
+	}
+	return nil
+}
+
+// ShouldExtractGPUDeviceIDsFromConfig returns true if GPU device IDs should be extracted
+// from container config/spec based on the current environment.
+//
+// Supported environments:
+//   - ECS: env var set by ECS agent in task definition
+//   - Kubernetes (non-GKE): env var set by NVIDIA device plugin via Allocate() API
+//
+// Not supported (returns false):
+//   - Docker standalone: env var injected at runtime by nvidia-container-toolkit, not in config
+//   - Standalone containerd: env var may be injected at runtime, not in spec
+//   - GKE: uses custom device plugin + gVisor runtime that ignores NVIDIA_VISIBLE_DEVICES
+//
+// Note: GKE is not explicitly detected here. On GKE, the env var is either not set in a useful way
+// or ignored by the runtime, so extraction returns nil and falls back to PodResources API.
+func ShouldExtractGPUDeviceIDsFromConfig() bool {
+	return env.IsECS() || env.IsKubernetes()
+}
+
+// ExtractGPUDeviceIDsFromEnvVars extracts GPU device IDs from environment variables
+// if the current environment supports it. This combines the environment check
+// with the extraction logic for convenience.
+//
+// Use this function when you have a list of environment variable strings (e.g., from container config).
+func ExtractGPUDeviceIDsFromEnvVars(envVars []string) []string {
+	if !ShouldExtractGPUDeviceIDsFromConfig() {
+		return nil
+	}
+	return ExtractGPUDeviceIDs(envVars)
+}
+
+// ExtractGPUDeviceIDsFromEnvMap extracts GPU device IDs from an environment variable map
+// if the current environment supports it.
+//
+// Use this function when you have a map of environment variables (e.g., from containerd spec parsing).
+func ExtractGPUDeviceIDsFromEnvMap(envs map[string]string) []string {
+	if !ShouldExtractGPUDeviceIDsFromConfig() {
+		return nil
+	}
+	if val, ok := envs[NVIDIAVisibleDevicesEnvVar]; ok && val != "" {
+		return strings.Split(val, ",")
+	}
+	return nil
+}

comp/core/workloadmeta/collectors/util/gpu_util_test.go

@@ -0,0 +1,130 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016-present Datadog, Inc.
+
+package util
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+// envMapToSlice converts a map of env vars to slice format (KEY=VALUE)
+func envMapToSlice(envs map[string]string) []string {
+	if envs == nil {
+		return nil
+	}
+	result := make([]string, 0, len(envs))
+	for k, v := range envs {
+		result = append(result, k+"="+v)
+	}
+	return result
+}
+
+// TestExtractGPUDeviceIDs tests ExtractGPUDeviceIDsFromEnvVars and ExtractGPUDeviceIDsFromEnvMap
+// with the same test cases, ensuring consistent behavior across different environments.
+func TestExtractGPUDeviceIDs(t *testing.T) {
+	tests := []struct {
+		name     string
+		envMap   map[string]string
+		isECS    bool
+		isK8s    bool
+		expected []string
+	}{
+		// Environment detection tests
+		{
+			name:     "ECS extracts GPU",
+			envMap:   map[string]string{"NVIDIA_VISIBLE_DEVICES": "GPU-xxx"},
+			isECS:    true,
+			expected: []string{"GPU-xxx"},
+		},
+		{
+			name:     "Kubernetes extracts GPU",
+			envMap:   map[string]string{"NVIDIA_VISIBLE_DEVICES": "GPU-aaa"},
+			isK8s:    true,
+			expected: []string{"GPU-aaa"},
+		},
+		{
+			name:     "Standalone returns nil",
+			envMap:   map[string]string{"NVIDIA_VISIBLE_DEVICES": "GPU-xxx"},
+			expected: nil,
+		},
+		// Parsing tests (in K8s environment)
+		{
+			name:     "single GPU UUID",
+			envMap:   map[string]string{"PATH": "/usr/bin", "NVIDIA_VISIBLE_DEVICES": "GPU-aec058b1-c18e-236e-c14d-49d2990fda0f"},
+			isK8s:    true,
+			expected: []string{"GPU-aec058b1-c18e-236e-c14d-49d2990fda0f"},
+		},
+		{
+			name:     "multiple GPU UUIDs",
+			envMap:   map[string]string{"NVIDIA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb,GPU-ccc"},
+			isK8s:    true,
+			expected: []string{"GPU-aaa", "GPU-bbb", "GPU-ccc"},
+		},
+		{
+			name:     "all GPUs",
+			envMap:   map[string]string{"NVIDIA_VISIBLE_DEVICES": "all"},
+			isK8s:    true,
+			expected: []string{"all"},
+		},
+		{
+			name:     "none",
+			envMap:   map[string]string{"NVIDIA_VISIBLE_DEVICES": "none"},
+			isK8s:    true,
+			expected: []string{"none"},
+		},
+		{
+			name:     "void",
+			envMap:   map[string]string{"NVIDIA_VISIBLE_DEVICES": "void"},
+			isK8s:    true,
+			expected: []string{"void"},
+		},
+		{
+			name:     "MIG instance",
+			envMap:   map[string]string{"NVIDIA_VISIBLE_DEVICES": "MIG-abc123-def456"},
+			isK8s:    true,
+			expected: []string{"MIG-abc123-def456"},
+		},
+		{
+			name:     "empty value",
+			envMap:   map[string]string{"NVIDIA_VISIBLE_DEVICES": ""},
+			isK8s:    true,
+			expected: nil,
+		},
+		{
+			name:     "no NVIDIA_VISIBLE_DEVICES",
+			envMap:   map[string]string{"PATH": "/usr/bin"},
+			isK8s:    true,
+			expected: nil,
+		},
+		{
+			name:     "nil map",
+			envMap:   nil,
+			isK8s:    true,
+			expected: nil,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.isECS {
+				t.Setenv("ECS_CONTAINER_METADATA_URI_V4", "http://169.254.170.2/v4")
+			}
+			if tt.isK8s {
+				t.Setenv("KUBERNETES_SERVICE_PORT", "443")
+			}
+
+			// Test map-based function (used by containerd)
+			resultMap := ExtractGPUDeviceIDsFromEnvMap(tt.envMap)
+			assert.Equal(t, tt.expected, resultMap, "ExtractGPUDeviceIDsFromEnvMap")
+
+			// Test slice-based function (used by docker)
+			envSlice := envMapToSlice(tt.envMap)
+			resultSlice := ExtractGPUDeviceIDsFromEnvVars(envSlice)
+			assert.Equal(t, tt.expected, resultSlice, "ExtractGPUDeviceIDsFromEnvVars")
+		})
+	}
+}