From 61cf8cd54df6e0d45190c90e150fd8c9a280f530 Mon Sep 17 00:00:00 2001
From: Gabriel Dos Santos <gabe.dossantos@datadoghq.com>
Date: Wed, 16 Oct 2024 19:56:17 +0000
Subject: [PATCH] Logs and tests for GPU pod tag

---
 .../collectors/workloadmeta_test.go           | 53 +++++++++++++++++++
 comp/core/workloadmeta/def/types.go           |  1 +
 2 files changed, 54 insertions(+)

diff --git a/comp/core/tagger/taggerimpl/collectors/workloadmeta_test.go b/comp/core/tagger/taggerimpl/collectors/workloadmeta_test.go
index 49078f94ac2c72..3064c9e787f25a 100644
--- a/comp/core/tagger/taggerimpl/collectors/workloadmeta_test.go
+++ b/comp/core/tagger/taggerimpl/collectors/workloadmeta_test.go
@@ -844,6 +844,59 @@ func TestHandleKubePod(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "pod with containers requesting gpu resources",
+			pod: workloadmeta.KubernetesPod{
+				EntityID: podEntityID,
+				EntityMeta: workloadmeta.EntityMeta{
+					Name:      podName,
+					Namespace: podNamespace,
+				},
+				GPUType: "nvidia",
+				Containers: []workloadmeta.OrchestratorContainer{
+					{
+						ID:    fullyFleshedContainerID,
+						Name:  containerName,
+						Image: image,
+					},
+				},
+			},
+			expected: []*types.TagInfo{
+				{
+					Source:       podSource,
+					EntityID:     podTaggerEntityID,
+					HighCardTags: []string{},
+					OrchestratorCardTags: []string{
+						fmt.Sprintf("pod_name:%s", podName),
+					},
+					LowCardTags: []string{
+						fmt.Sprintf("kube_namespace:%s", podNamespace),
+						"kube_gpu_type:nvidia",
+					},
+					StandardTags: []string{},
+				},
+				{
+					Source:   podSource,
+					EntityID: fullyFleshedContainerTaggerEntityID,
+					HighCardTags: []string{
+						fmt.Sprintf("container_id:%s", fullyFleshedContainerID),
+						fmt.Sprintf("display_container_name:%s_%s", runtimeContainerName, podName),
+					},
+					OrchestratorCardTags: []string{
+						fmt.Sprintf("pod_name:%s", podName),
+					},
+					LowCardTags: append([]string{
+						fmt.Sprintf("kube_namespace:%s", podNamespace),
+						fmt.Sprintf("kube_container_name:%s", containerName),
+						"image_id:datadog/agent@sha256:a63d3f66fb2f69d955d4f2ca0b229385537a77872ffc04290acae65aed5317d2",
+						"image_name:datadog/agent",
+						"image_tag:latest",
+						"short_image:agent",
+					}, standardTags...),
+					StandardTags: standardTags,
+				},
+			},
+		},
 	}
 
 	for _, tt := range tests {
diff --git a/comp/core/workloadmeta/def/types.go b/comp/core/workloadmeta/def/types.go
index 51213b463fa375..fa5199bfb301be 100644
--- a/comp/core/workloadmeta/def/types.go
+++ b/comp/core/workloadmeta/def/types.go
@@ -746,6 +746,7 @@ func (p KubernetesPod) String(verbose bool) string {
 	if verbose {
 		_, _ = fmt.Fprintln(&sb, "Priority Class:", p.PriorityClass)
 		_, _ = fmt.Fprintln(&sb, "QOS Class:", p.QOSClass)
+		_, _ = fmt.Fprintln(&sb, "GPU Type:", p.GPUType)
 		_, _ = fmt.Fprintln(&sb, "Runtime Class:", p.RuntimeClass)
 		_, _ = fmt.Fprintln(&sb, "PVCs:", sliceToString(p.PersistentVolumeClaimNames))
 		_, _ = fmt.Fprintln(&sb, "Kube Services:", sliceToString(p.KubeServices))