From d793733525a70d0a6b869793c5561cf25d093b71 Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Wed, 4 Dec 2024 18:14:28 -0800
Subject: [PATCH] Kubernetes benchmarks: Add profiling to benchmarks that were
 missing it.

(Diffbased)

PiperOrigin-RevId: 702928932
---
 test/kubernetes/benchmarks/BUILD              | 2 ++
 test/kubernetes/benchmarks/ollama.go          | 6 ++++++
 test/kubernetes/benchmarks/stablediffusion.go | 6 ++++++
 3 files changed, 14 insertions(+)

diff --git a/test/kubernetes/benchmarks/BUILD b/test/kubernetes/benchmarks/BUILD
index ffc782041..21ea89c4e 100644
--- a/test/kubernetes/benchmarks/BUILD
+++ b/test/kubernetes/benchmarks/BUILD
@@ -358,6 +358,7 @@ go_library(
     deps = [
         "//test/gpu/ollama",
         "//test/kubernetes",
+        "//test/kubernetes/benchmarks/profiling",
         "//test/kubernetes/benchmetric",
         "//test/kubernetes/k8sctx",
         "//test/kubernetes/testcluster",
@@ -392,6 +393,7 @@ go_library(
     deps = [
         "//test/gpu/stablediffusion",
         "//test/kubernetes",
+        "//test/kubernetes/benchmarks/profiling",
         "//test/kubernetes/benchmetric",
         "//test/kubernetes/k8sctx",
         "//test/kubernetes/testcluster",
diff --git a/test/kubernetes/benchmarks/ollama.go b/test/kubernetes/benchmarks/ollama.go
index 60301b9dd..be69e595a 100644
--- a/test/kubernetes/benchmarks/ollama.go
+++ b/test/kubernetes/benchmarks/ollama.go
@@ -28,6 +28,7 @@ import (
 
 	"gvisor.dev/gvisor/test/gpu/ollama"
 	k8s "gvisor.dev/gvisor/test/kubernetes"
+	"gvisor.dev/gvisor/test/kubernetes/benchmarks/profiling"
 	"gvisor.dev/gvisor/test/kubernetes/benchmetric"
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
@@ -235,6 +236,11 @@ func BenchmarkOllama(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kubernetes
 		t.Fatalf("cannot reset namespace: %v", err)
 	}
 	defer benchmarkNS.Cleanup(ctx)
+	endProfiling, err := profiling.MaybeSetup(ctx, t, k8sCtx, cluster, benchmarkNS)
+	if err != nil {
+		t.Fatalf("Failed to setup profiling: %v", err)
+	}
+	defer endProfiling()
 
 	logWithTime := func(t *testing.T, format string, values ...any) {
 		t.Logf("[%v] "+format, append([]any{time.Now().Format(time.TimeOnly)}, values...)...)
diff --git a/test/kubernetes/benchmarks/stablediffusion.go b/test/kubernetes/benchmarks/stablediffusion.go
index a87666434..5bfe239af 100644
--- a/test/kubernetes/benchmarks/stablediffusion.go
+++ b/test/kubernetes/benchmarks/stablediffusion.go
@@ -24,6 +24,7 @@ import (
 
 	"gvisor.dev/gvisor/test/gpu/stablediffusion"
 	k8s "gvisor.dev/gvisor/test/kubernetes"
+	"gvisor.dev/gvisor/test/kubernetes/benchmarks/profiling"
 	"gvisor.dev/gvisor/test/kubernetes/benchmetric"
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
@@ -112,6 +113,11 @@ func RunStableDiffusionXL(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kuber
 		t.Fatalf("cannot reset namespace: %v", err)
 	}
 	defer benchmarkNS.Cleanup(ctx)
+	endProfiling, err := profiling.MaybeSetup(ctx, t, k8sCtx, cluster, benchmarkNS)
+	if err != nil {
+		t.Fatalf("Failed to setup profiling: %v", err)
+	}
+	defer endProfiling()
 
 	imageName, err := k8sCtx.ResolveImage(ctx, stableDiffusionImage)
 	if err != nil {