Ensure that pod_uid is set for DCGM metrics in the collector source.

Signed-off-by: Sean Holcomb <seanholcomb@gmail.com>

modules/collector-source/pkg/scrape/dcgm.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/opencost/opencost/core/pkg/clustercache"
 	"github.com/opencost/opencost/core/pkg/log"
+	"github.com/opencost/opencost/core/pkg/source"
 	"github.com/opencost/opencost/modules/collector-source/pkg/event"
 	"github.com/opencost/opencost/modules/collector-source/pkg/metric"
 	"github.com/opencost/opencost/modules/collector-source/pkg/scrape/target"
@@ -16,10 +17,10 @@ var dcgmRegex = regexp.MustCompile("(?i)(.*dcgm-exporter.*)")
 
 func newDCGMScrapper(clusterCache clustercache.ClusterCache) Scraper {
 	tp := newDCGMTargetProvider(clusterCache)
-	return newDCGMTargetScraper(tp)
+	return newDCGMTargetScraper(tp, podUIDEnricher(clusterCache))
 }
 
-func newDCGMTargetScraper(provider target.TargetProvider) *TargetScraper {
+func newDCGMTargetScraper(provider target.TargetProvider, enrich UpdateEnricher) *TargetScraper {
 	return newTargetScrapper(
 		event.DCGMScraperName,
 		provider,
@@ -27,7 +28,29 @@ func newDCGMTargetScraper(provider target.TargetProvider) *TargetScraper {
 			metric.DCGMFIPROFGRENGINEACTIVE,
 			metric.DCGMFIDEVDECUTIL,
 		},
-		true)
+		true,
+		enrich)
+}
+
+// podUIDEnricher backfills pod_uid on a DCGM update using its own namespace/pod name labels,
+// resolved against a freshly built index of the cluster's current pods. Left unset if pod_uid
+// is already present, or if namespace/pod can't be resolved to a known pod.
+func podUIDEnricher(clusterCache clustercache.ClusterCache) UpdateEnricher {
+	index := buildPodIndex(clusterCache.GetAllPods())
+	return func(update metric.Update) metric.Update {
+		if update.Labels[source.PodUIDLabel] != "" {
+			return update
+		}
+		namespace, pod := update.Labels[source.NamespaceLabel], update.Labels[source.PodLabel]
+		if namespace == "" || pod == "" {
+			return update
+		}
+
+		if uid, ok := index[podKey{namespace: namespace, name: pod}]; ok {
+			update.Labels[source.PodUIDLabel] = string(uid)
+		}
+		return update
+	}
 }
 
 type DCGMTargetProvider struct {

modules/collector-source/pkg/scrape/dcgm_test.go

@@ -2,6 +2,10 @@ package scrape
 
 import (
 	"testing"
+
+	"github.com/opencost/opencost/core/pkg/clustercache"
+	"github.com/opencost/opencost/core/pkg/source"
+	"github.com/opencost/opencost/modules/collector-source/pkg/metric"
 )
 
 func Test_isDCGM(t *testing.T) {
@@ -62,3 +66,89 @@ func Test_isDCGM(t *testing.T) {
 		})
 	}
 }
+
+func Test_podUIDEnricher(t *testing.T) {
+	cache := &clustercache.MockClusterCache{
+		Pods: []*clustercache.Pod{
+			{UID: "pod-uid-1", Name: "pod1", Namespace: "namespace1"},
+		},
+	}
+	enrich := podUIDEnricher(cache)
+
+	tests := map[string]struct {
+		update metric.Update
+		want   metric.Update
+	}{
+		"resolves pod_uid from namespace and pod": {
+			update: metric.Update{
+				Labels: map[string]string{
+					source.NamespaceLabel: "namespace1",
+					source.PodLabel:       "pod1",
+				},
+			},
+			want: metric.Update{
+				Labels: map[string]string{
+					source.NamespaceLabel: "namespace1",
+					source.PodLabel:       "pod1",
+					source.PodUIDLabel:    "pod-uid-1",
+				},
+			},
+		},
+		"unknown pod is left unset": {
+			update: metric.Update{
+				Labels: map[string]string{
+					source.NamespaceLabel: "namespace1",
+					source.PodLabel:       "unknown-pod",
+				},
+			},
+			want: metric.Update{
+				Labels: map[string]string{
+					source.NamespaceLabel: "namespace1",
+					source.PodLabel:       "unknown-pod",
+				},
+			},
+		},
+		"missing namespace or pod label is left unset": {
+			update: metric.Update{
+				Labels: map[string]string{
+					source.PodLabel: "pod1",
+				},
+			},
+			want: metric.Update{
+				Labels: map[string]string{
+					source.PodLabel: "pod1",
+				},
+			},
+		},
+		"existing non-empty pod_uid is left untouched": {
+			update: metric.Update{
+				Labels: map[string]string{
+					source.NamespaceLabel: "namespace1",
+					source.PodLabel:       "pod1",
+					source.PodUIDLabel:    "already-set",
+				},
+			},
+			want: metric.Update{
+				Labels: map[string]string{
+					source.NamespaceLabel: "namespace1",
+					source.PodLabel:       "pod1",
+					source.PodUIDLabel:    "already-set",
+				},
+			},
+		},
+	}
+
+	for name, tt := range tests {
+		t.Run(name, func(t *testing.T) {
+			got := enrich(tt.update)
+			if len(got.Labels) != len(tt.want.Labels) {
+				t.Fatalf("got labels %v, want %v", got.Labels, tt.want.Labels)
+			}
+			for k, v := range tt.want.Labels {
+				if got.Labels[k] != v {
+					t.Errorf("label %q = %q, want %q", k, got.Labels[k], v)
+				}
+			}
+		})
+	}
+}

modules/collector-source/pkg/scrape/index.go

@@ -46,3 +46,18 @@ func buildPVIndex(pvs []*clustercache.PersistentVolume) map[string]types.UID {
 	}
 	return m
 }
+
+// podKey is a composite key for a Pod (namespace + name).
+type podKey struct {
+	namespace string
+	name      string
+}
+
+// buildPodIndex returns a map from (namespace, name) to Pod UID.
+func buildPodIndex(pods []*clustercache.Pod) map[podKey]types.UID {
+	m := make(map[podKey]types.UID, len(pods))
+	for _, pod := range pods {
+		m[podKey{namespace: pod.Namespace, name: pod.Name}] = pod.UID
+	}
+	return m
+}

modules/collector-source/pkg/scrape/network.go

@@ -32,7 +32,8 @@ func newNetworkTargetScraper(provider target.TargetProvider) *TargetScraper {
 			metric.KubecostPodNetworkEgressBytesTotal,
 			metric.KubecostPodNetworkIngressBytesTotal,
 		},
-		true)
+		true,
+		nil)
 }
 
 type NetworkTargetProvider struct {

modules/collector-source/pkg/scrape/opencost.go

@@ -36,5 +36,6 @@ func newOpencostTargetScraper(provider target.TargetProvider) *TargetScraper {
 			metric.NodeGPUCount,
 			metric.KubecostNodeIsSpot,
 		},
-		true)
+		true,
+		nil)
 }

modules/collector-source/pkg/scrape/targetscraper.go

@@ -12,14 +12,19 @@ import (
 	"github.com/opencost/opencost/modules/collector-source/pkg/scrape/target"
 )
 
+// UpdateEnricher optionally transforms a single scraped update before it's returned from
+// Scrape().
+type UpdateEnricher func(update metric.Update) metric.Update
+
 type TargetScraper struct {
 	name           string // identifier for the scraper
 	targetProvider target.TargetProvider
 	metricNames    map[string]struct{} // filter for which metrics will be processed
 	includeMetrics bool                // toggle to make metrics an include or exclude list
+	enrich         UpdateEnricher      // optional per-update enrichment, nil means no-op
 }
 
-func newTargetScrapper(name string, provider target.TargetProvider, metricNames []string, includeMetrics bool) *TargetScraper {
+func newTargetScrapper(name string, provider target.TargetProvider, metricNames []string, includeMetrics bool, enrich UpdateEnricher) *TargetScraper {
 	metricSet := make(map[string]struct{})
 	for _, metricName := range metricNames {
 		metricSet[metricName] = struct{}{}
@@ -29,6 +34,7 @@ func newTargetScrapper(name string, provider target.TargetProvider, metricNames
 		targetProvider: provider,
 		metricNames:    metricSet,
 		includeMetrics: includeMetrics,
+		enrich:         enrich,
 	}
 }
 
@@ -70,11 +76,15 @@ func (s *TargetScraper) Scrape() []metric.Update {
 				if _, ok := s.metricNames[result.Name]; ok != s.includeMetrics {
 					continue
 				}
-				scrapeResults = append(scrapeResults, metric.Update{
+				update := metric.Update{
 					Name:   result.Name,
 					Labels: result.Labels,
 					Value:  result.Value,
-				})
+				}
+				if s.enrich != nil {
+					update = s.enrich(update)
+				}
+				scrapeResults = append(scrapeResults, update)
 			}
 			return scrapeResults
 		}

modules/collector-source/pkg/scrape/targetscraper_test.go

@@ -95,6 +95,13 @@ pod_pvc_allocation{namespace="namespace1",persistentvolume="pvc-1",persistentvol
 pod_pvc_allocation{namespace="namespace1",persistentvolume="pvc-2",persistentvolumeclaim="pvc2",pod="pod2"} 3.4359738368e+10
 `
 
+const enrichScrape = `
+# HELP test_metric test metric
+# TYPE test_metric gauge
+test_metric{name="a"} 1
+test_metric{name="b",extra="already-set"} 2
+`
+
 const dcgmScrape = `
 # HELP DCGM_FI_PROF_GR_ENGINE_ACTIVE Ratio of time the graphics engine is active.
 # TYPE DCGM_FI_PROF_GR_ENGINE_ACTIVE gauge
@@ -411,9 +418,11 @@ func TestTargetScraper_Scrape(t *testing.T) {
 			},
 		},
 		{
-			name:                 "GPU Metric",
-			scrapeText:           dcgmScrape,
-			targetScraperFactory: newDCGMTargetScraper,
+			name:       "GPU Metric",
+			scrapeText: dcgmScrape,
+			targetScraperFactory: func(provider target.TargetProvider) *TargetScraper {
+				return newDCGMTargetScraper(provider, nil)
+			},
 			expected: []metric.Update{
 				{
 					Name: metric.DCGMFIPROFGRENGINEACTIVE,
@@ -441,6 +450,38 @@ func TestTargetScraper_Scrape(t *testing.T) {
 				},
 			},
 		},
+		{
+			name:       "Enrichment",
+			scrapeText: enrichScrape,
+			targetScraperFactory: func(provider target.TargetProvider) *TargetScraper {
+				enrich := func(update metric.Update) metric.Update {
+					if update.Labels["extra"] != "" {
+						return update
+					}
+					update.Labels["extra"] = "enriched"
+					return update
+				}
+				return newTargetScrapper("test-enrich", provider, nil, false, enrich)
+			},
+			expected: []metric.Update{
+				{
+					Name: "test_metric",
+					Labels: map[string]string{
+						"name":  "a",
+						"extra": "enriched",
+					},
+					Value: 1,
+				},
+				{
+					Name: "test_metric",
+					Labels: map[string]string{
+						"name":  "b",
+						"extra": "already-set",
+					},
+					Value: 2,
+				},
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {