From 75c9597d8ec85b3e8c13989bf5b3f59d221be6ce Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Tue, 18 Jun 2024 15:31:13 -0700
Subject: [PATCH] Add `nvproxy` and `tpuproxy` annotation metrics for GKE
 Sandbox pods.

PiperOrigin-RevId: 644530875
---
 runsc/metricserver/containermetrics/BUILD               | 1 +
 runsc/metricserver/containermetrics/containermetrics.go | 7 +++++++
 2 files changed, 8 insertions(+)

diff --git a/runsc/metricserver/containermetrics/BUILD b/runsc/metricserver/containermetrics/BUILD
index 8f4b681d6..7fdd2fdb0 100644
--- a/runsc/metricserver/containermetrics/BUILD
+++ b/runsc/metricserver/containermetrics/BUILD
@@ -14,5 +14,6 @@ go_library(
     deps = [
         "//pkg/prometheus",
         "//runsc/container",
+        "//runsc/specutils",
     ],
 )
diff --git a/runsc/metricserver/containermetrics/containermetrics.go b/runsc/metricserver/containermetrics/containermetrics.go
index 9d9ce543c..4961ed9d9 100644
--- a/runsc/metricserver/containermetrics/containermetrics.go
+++ b/runsc/metricserver/containermetrics/containermetrics.go
@@ -24,6 +24,7 @@ import (
 
 	"gvisor.dev/gvisor/pkg/prometheus"
 	"gvisor.dev/gvisor/runsc/container"
+	"gvisor.dev/gvisor/runsc/specutils"
 )
 
 // SandboxPrometheusLabels returns a set of Prometheus labels that identifies the sandbox running
@@ -70,6 +71,8 @@ func ComputeSpecMetadata(allContainers []*container.Container) map[string]string
 
 	hasUID0Container := false
 	ociVersion := unknownOCIVersion
+	hasNVProxy := false
+	hasTPUProxy := false
 	for _, cont := range allContainers {
 		if cont.RunsAsUID0() {
 			hasUID0Container = true
@@ -79,9 +82,13 @@ func ComputeSpecMetadata(allContainers []*container.Container) map[string]string
 		} else if ociVersion != cont.Spec.Version {
 			ociVersion = inconsistentOCIVersion
 		}
+		hasNVProxy = hasNVProxy || cont.Spec.Annotations[specutils.AnnotationNVProxy] == "true"
+		hasTPUProxy = hasTPUProxy || cont.Spec.Annotations[specutils.AnnotationTPU] == "true"
 	}
 	return map[string]string{
 		"hasuid0":    strconv.FormatBool(hasUID0Container),
 		"ociversion": ociVersion,
+		"nvproxy":    strconv.FormatBool(hasNVProxy),
+		"tpuproxy":   strconv.FormatBool(hasTPUProxy),
 	}
 }