From 45b346e8fd424c9219fd67a0c992d08ddfb56263 Mon Sep 17 00:00:00 2001
From: Jing Chen <chjing@google.com>
Date: Tue, 12 Nov 2024 10:42:27 -0800
Subject: [PATCH] Use AR images on the domain pkg.dev.

Use pkg.dev consistently since GCR is deprecated. We don't have to
manage mirrored image repos for gcr.io and pkg.dev.

PiperOrigin-RevId: 695791369
---
 test/kubernetes/BUILD                         | 10 +++++++++-
 test/kubernetes/benchmarks/BUILD              | 10 ++++++++++
 test/kubernetes/benchmarks/abslbuild_test.go  |  3 ++-
 test/kubernetes/benchmarks/ffmpeg_test.go     |  5 +++--
 test/kubernetes/benchmarks/grpc_test.go       |  5 +++--
 test/kubernetes/benchmarks/httpbench/BUILD    |  1 +
 .../benchmarks/httpbench/httpbench.go         |  3 ++-
 test/kubernetes/benchmarks/nginx_test.go      |  5 +++--
 test/kubernetes/benchmarks/ollama_test.go     |  5 +++--
 test/kubernetes/benchmarks/profiling/BUILD    |  1 +
 .../benchmarks/profiling/profiling.go         |  5 +++--
 test/kubernetes/benchmarks/pytorch_test.go    |  7 +++----
 test/kubernetes/benchmarks/redis_test.go      |  5 +++--
 test/kubernetes/benchmarks/ruby_dev_test.go   |  5 +++--
 .../benchmarks/stablediffusion_test.go        |  3 ++-
 test/kubernetes/benchmarks/tensorflow_test.go |  5 +++--
 test/kubernetes/helpers.go                    | 20 +++++++++++++++++++
 17 files changed, 74 insertions(+), 24 deletions(-)
 create mode 100644 test/kubernetes/helpers.go

diff --git a/test/kubernetes/BUILD b/test/kubernetes/BUILD
index 373e8a41bf..fe0b35606a 100644
--- a/test/kubernetes/BUILD
+++ b/test/kubernetes/BUILD
@@ -1,4 +1,4 @@
-load("//tools:defs.bzl", "proto_library")
+load("//tools:defs.bzl", "go_library", "proto_library")
 
 package(
     default_applicable_licenses = ["//:license"],
@@ -16,3 +16,11 @@ proto_library(
         "@com_google_protobuf//:any_proto",
     ],
 )
+
+go_library(
+    name = "kubernetes",
+    srcs = ["helpers.go"],
+    visibility = [
+        "//visibility:public",
+    ],
+)
diff --git a/test/kubernetes/benchmarks/BUILD b/test/kubernetes/benchmarks/BUILD
index 8593cdabc2..6ac7d55436 100644
--- a/test/kubernetes/benchmarks/BUILD
+++ b/test/kubernetes/benchmarks/BUILD
@@ -14,6 +14,7 @@ go_test(
         "notap",
     ],
     deps = [
+        "//test/kubernetes",
         "//test/kubernetes/benchmarks/profiling",
         "//test/kubernetes/benchmetric",
         "//test/kubernetes/k8sctx",
@@ -49,6 +50,7 @@ go_test(
         "notap",
     ],
     deps = [
+        "//test/kubernetes",
         "//test/kubernetes/benchmarks/profiling",
         "//test/kubernetes/benchmetric",
         "//test/kubernetes/k8sctx",
@@ -68,6 +70,7 @@ go_test(
     ],
     deps = [
         "//test/benchmarks/tools",
+        "//test/kubernetes",
         "//test/kubernetes/benchmarks/profiling",
         "//test/kubernetes/benchmetric",
         "//test/kubernetes/k8sctx",
@@ -86,6 +89,7 @@ go_test(
         "notap",
     ],
     deps = [
+        "//test/kubernetes",
         "//test/kubernetes/benchmarks/profiling",
         "//test/kubernetes/benchmetric",
         "//test/kubernetes/k8sctx",
@@ -104,6 +108,7 @@ go_test(
         "notap",
     ],
     deps = [
+        "//test/kubernetes",
         "//test/kubernetes/benchmarks/profiling",
         "//test/kubernetes/benchmetric",
         "//test/kubernetes/k8sctx",
@@ -122,6 +127,7 @@ go_test(
         "notap",
     ],
     deps = [
+        "//test/kubernetes",
         "//test/kubernetes/benchmarks/httpbench",
         "//test/kubernetes/benchmarks/profiling",
         "//test/kubernetes/k8sctx",
@@ -159,6 +165,7 @@ go_test(
         "notap",
     ],
     deps = [
+        "//test/kubernetes",
         "//test/kubernetes/benchmarks/profiling",
         "//test/kubernetes/benchmetric",
         "//test/kubernetes/k8sctx",
@@ -196,6 +203,7 @@ go_test(
         "notap",
     ],
     deps = [
+        "//test/kubernetes",
         "//test/kubernetes/benchmarks/profiling",
         "//test/kubernetes/benchmetric",
         "//test/kubernetes/k8sctx",
@@ -217,6 +225,7 @@ go_test(
     ],
     deps = [
         "//test/gpu/ollama",
+        "//test/kubernetes",
         "//test/kubernetes/benchmetric",
         "//test/kubernetes/k8sctx",
         "//test/kubernetes/testcluster",
@@ -237,6 +246,7 @@ go_test(
     ],
     deps = [
         "//test/gpu/stablediffusion",
+        "//test/kubernetes",
         "//test/kubernetes/benchmetric",
         "//test/kubernetes/k8sctx",
         "//test/kubernetes/testcluster",
diff --git a/test/kubernetes/benchmarks/abslbuild_test.go b/test/kubernetes/benchmarks/abslbuild_test.go
index 204f6be6a6..d3e9d0bca8 100644
--- a/test/kubernetes/benchmarks/abslbuild_test.go
+++ b/test/kubernetes/benchmarks/abslbuild_test.go
@@ -21,6 +21,7 @@ import (
 	"strings"
 	"testing"
 
+	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"
@@ -30,7 +31,7 @@ import (
 )
 
 const (
-	imageAMD = "gcr.io/gvisor-presubmit/benchmarks/absl_x86_64:latest"
+	imageAMD = k8s.ImageRepoPrefix + "benchmarks/absl_x86_64:latest"
 )
 
 // TestABSLBuild benchmarks building various Abseil C++ targets.
diff --git a/test/kubernetes/benchmarks/ffmpeg_test.go b/test/kubernetes/benchmarks/ffmpeg_test.go
index b2c18e5365..e102483ef9 100644
--- a/test/kubernetes/benchmarks/ffmpeg_test.go
+++ b/test/kubernetes/benchmarks/ffmpeg_test.go
@@ -20,6 +20,7 @@ import (
 	"strings"
 	"testing"
 
+	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"
@@ -29,9 +30,9 @@ import (
 )
 
 const (
-	imageAMD            = "gcr.io/gvisor-presubmit/benchmarks/ffmpeg_x86_64:latest"
-	imageARM            = "gcr.io/gvisor-presubmit/benchmarks/ffmpeg_aarch64:latest"
 	ffmpegContainerName = "ffmpeg"
+	imageAMD            = k8s.ImageRepoPrefix + "benchmarks/ffmpeg_x86_64:latest"
+	imageARM            = k8s.ImageRepoPrefix + "benchmarks/ffmpeg_aarch64:latest"
 )
 
 func TestFfmpeg(t *testing.T) {
diff --git a/test/kubernetes/benchmarks/grpc_test.go b/test/kubernetes/benchmarks/grpc_test.go
index 7440bd4daf..55ec2a8574 100644
--- a/test/kubernetes/benchmarks/grpc_test.go
+++ b/test/kubernetes/benchmarks/grpc_test.go
@@ -21,6 +21,7 @@ import (
 	"strings"
 	"testing"
 
+	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"
@@ -30,8 +31,8 @@ import (
 )
 
 const (
-	imageAMD = "gcr.io/gvisor-presubmit/benchmarks/build-grpc_x86_64:latest"
-	imageARM = "gcr.io/gvisor-presubmit/benchmarks/build-grpc_aarch64:latest"
+	imageAMD = k8s.ImageRepoPrefix + "benchmarks/build-grpc_x86_64:latest"
+	imageARM = k8s.ImageRepoPrefix + "benchmarks/build-grpc_aarch64:latest"
 )
 
 func TestGRPCBuild(t *testing.T) {
diff --git a/test/kubernetes/benchmarks/httpbench/BUILD b/test/kubernetes/benchmarks/httpbench/BUILD
index af0802c79b..6b2dca74e5 100644
--- a/test/kubernetes/benchmarks/httpbench/BUILD
+++ b/test/kubernetes/benchmarks/httpbench/BUILD
@@ -14,6 +14,7 @@ go_library(
         "//test/kubernetes/benchmarks:__subpackages__",
     ],
     deps = [
+        "//test/kubernetes",
         "//test/kubernetes/benchmetric",
         "//test/kubernetes/testcluster",
         "@io_k8s_api//core/v1:go_default_library",
diff --git a/test/kubernetes/benchmarks/httpbench/httpbench.go b/test/kubernetes/benchmarks/httpbench/httpbench.go
index 2a5d013e13..72a500d0f5 100644
--- a/test/kubernetes/benchmarks/httpbench/httpbench.go
+++ b/test/kubernetes/benchmarks/httpbench/httpbench.go
@@ -25,6 +25,7 @@ import (
 	"testing"
 	"time"
 
+	k8s "gvisor.dev/gvisor/test/kubernetes"
 	"gvisor.dev/gvisor/test/kubernetes/benchmetric"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
 	v13 "k8s.io/api/core/v1"
@@ -37,7 +38,7 @@ import (
 // the HTTP server with load on a single connection.
 const InfiniteQPS = 1 << 16
 
-const wrk2ImageAMD = "gcr.io/gvisor-presubmit/benchmarks/wrk2_x86_64:latest"
+const wrk2ImageAMD = k8s.ImageRepoPrefix + "benchmarks/wrk2_x86_64:latest"
 
 // MetricType is a type of metric to report.
 type MetricType int
diff --git a/test/kubernetes/benchmarks/nginx_test.go b/test/kubernetes/benchmarks/nginx_test.go
index aa9f146978..2220010b7b 100644
--- a/test/kubernetes/benchmarks/nginx_test.go
+++ b/test/kubernetes/benchmarks/nginx_test.go
@@ -21,6 +21,7 @@ import (
 	"testing"
 	"time"
 
+	k8s "gvisor.dev/gvisor/test/kubernetes"
 	"gvisor.dev/gvisor/test/kubernetes/benchmarks/httpbench"
 	"gvisor.dev/gvisor/test/kubernetes/benchmarks/profiling"
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
@@ -31,8 +32,6 @@ import (
 )
 
 const (
-	nginxImageAMD          = "gcr.io/gvisor-presubmit/benchmarks/nginx_x86_64:latest"
-	nginxImageARM          = "gcr.io/gvisor-presubmit/benchmarks/nginx_aarch64:latest"
 	nginxPort              = 80
 	nginxBenchmarkDuration = 70 * time.Second
 	nginxRequestTimeout    = 3 * time.Second
@@ -40,6 +39,8 @@ const (
 
 	nginxServerLabelKey   = "app.kubernetes.io/name"
 	nginxServerLabelValue = "nginx-server"
+	nginxImageAMD         = k8s.ImageRepoPrefix + "benchmarks/nginx_x86_64:latest"
+	nginxImageARM         = k8s.ImageRepoPrefix + "benchmarks/nginx_aarch64:latest"
 )
 
 var (
diff --git a/test/kubernetes/benchmarks/ollama_test.go b/test/kubernetes/benchmarks/ollama_test.go
index b496391d6f..88b833c6b1 100644
--- a/test/kubernetes/benchmarks/ollama_test.go
+++ b/test/kubernetes/benchmarks/ollama_test.go
@@ -28,6 +28,7 @@ import (
 	"unicode"
 
 	"gvisor.dev/gvisor/test/gpu/ollama"
+	k8s "gvisor.dev/gvisor/test/kubernetes"
 	"gvisor.dev/gvisor/test/kubernetes/benchmetric"
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
@@ -816,11 +817,11 @@ func doOllamaTest(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesCon
 const (
 	ollamaServerLabelKey   = "app.kubernetes.io/name"
 	ollamaServerLabelValue = "ollama-server"
-	ollamaBenchImage       = "gcr.io/gvisor-presubmit/benchmarks/gpu/ollama/bench:latest"
-	ollamaBenchClientImage = "gcr.io/gvisor-presubmit/benchmarks/gpu/ollama/client:latest"
 	ollamaPort             = 11434
 	ollamaPodName          = "ollama-server"
 	ollamaServiceName      = "ollama-service"
+	ollamaBenchImage       = k8s.ImageRepoPrefix + "benchmarks/gpu/ollama/bench:latest"
+	ollamaBenchClientImage = k8s.ImageRepoPrefix + "benchmarks/gpu/ollama/client:latest"
 )
 
 // newOllamaServerPod returns the pod spec for an ollama server.
diff --git a/test/kubernetes/benchmarks/profiling/BUILD b/test/kubernetes/benchmarks/profiling/BUILD
index 0adc683fae..c769860bc7 100644
--- a/test/kubernetes/benchmarks/profiling/BUILD
+++ b/test/kubernetes/benchmarks/profiling/BUILD
@@ -16,6 +16,7 @@ go_library(
     deps = [
         "//pkg/atomicbitops",
         "//runsc/flag",
+        "//test/kubernetes",
         "//test/kubernetes/testcluster",
         "//test/metricsviz",
         "@com_github_google_pprof//profile:go_default_library",
diff --git a/test/kubernetes/benchmarks/profiling/profiling.go b/test/kubernetes/benchmarks/profiling/profiling.go
index db29d27893..3feb736a3d 100644
--- a/test/kubernetes/benchmarks/profiling/profiling.go
+++ b/test/kubernetes/benchmarks/profiling/profiling.go
@@ -42,6 +42,7 @@ import (
 	"google.golang.org/protobuf/proto"
 	"gvisor.dev/gvisor/pkg/atomicbitops"
 	"gvisor.dev/gvisor/runsc/flag"
+	k8s "gvisor.dev/gvisor/test/kubernetes"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
 	"gvisor.dev/gvisor/test/metricsviz"
 	appsv1 "k8s.io/api/apps/v1"
@@ -61,10 +62,10 @@ var (
 
 const (
 	setupPodName            = "runsc-profiling-setup"
-	profileHelperImageAMD64 = "gcr.io/gvisor-presubmit/benchmarks/profile-helper_x86_64:latest"
-	profileHelperImageARM64 = "gcr.io/gvisor-presubmit/benchmarks/profile-helper_aarch64:latest"
 	k8sApp                  = "k8s-app"
 	hostMountDir            = "/host"
+	profileHelperImageAMD64 = k8s.ImageRepoPrefix + "benchmarks/profile-helper_x86_64:latest"
+	profileHelperImageARM64 = k8s.ImageRepoPrefix + "benchmarks/profile-helper_aarch64:latest"
 )
 
 // postProcessor is a function that is called after a profiling run completes.
diff --git a/test/kubernetes/benchmarks/pytorch_test.go b/test/kubernetes/benchmarks/pytorch_test.go
index 4cb1ab5d29..b904da867a 100644
--- a/test/kubernetes/benchmarks/pytorch_test.go
+++ b/test/kubernetes/benchmarks/pytorch_test.go
@@ -34,6 +34,7 @@ import (
 	"testing"
 	"time"
 
+	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"
@@ -42,16 +43,14 @@ import (
 	v13 "k8s.io/api/core/v1"
 )
 
-const (
-	pytorchImage = "gcr.io/gvisor-presubmit/benchmarks/pytorch_x86_64:f6f280aeb1b07989"
-)
-
 // pytorchTestType is the method used, either training or evaluation, for the model.
 type pytorchTestType string
 
 const (
 	train = pytorchTestType("train")
 	eval  = pytorchTestType("eval")
+
+	pytorchImage = k8s.ImageRepoPrefix + "benchmarks/pytorch_x86_64:f6f280aeb1b07989"
 )
 
 type pytorchMode string
diff --git a/test/kubernetes/benchmarks/redis_test.go b/test/kubernetes/benchmarks/redis_test.go
index 17b682cc75..5bf4363867 100644
--- a/test/kubernetes/benchmarks/redis_test.go
+++ b/test/kubernetes/benchmarks/redis_test.go
@@ -27,6 +27,7 @@ import (
 	"testing"
 	"time"
 
+	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"
@@ -39,12 +40,12 @@ const (
 	redisPort                    = 6379
 	defaultRequestsPerConnection = 50000
 
-	redisImageAMD         = "gcr.io/gvisor-presubmit/benchmarks/redis_x86_64:latest"
-	redisImageARM         = "gcr.io/gvisor-presubmit/benchmarks/redis_aarch64:latest"
 	redisServerLabelKey   = "app.kubernetes.io/name"
 	redisServerLabelValue = "redis-server"
 	redisVolumeName       = "redis-data"
 	redisDataDirectory    = "/redis-data"
+	redisImageAMD         = k8s.ImageRepoPrefix + "benchmarks/redis_x86_64:latest"
+	redisImageARM         = k8s.ImageRepoPrefix + "benchmarks/redis_aarch64:latest"
 )
 
 var (
diff --git a/test/kubernetes/benchmarks/ruby_dev_test.go b/test/kubernetes/benchmarks/ruby_dev_test.go
index e261241b84..17ca95d209 100644
--- a/test/kubernetes/benchmarks/ruby_dev_test.go
+++ b/test/kubernetes/benchmarks/ruby_dev_test.go
@@ -25,6 +25,7 @@ import (
 	"testing"
 
 	"gvisor.dev/gvisor/test/benchmarks/tools"
+	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"
@@ -34,9 +35,9 @@ import (
 )
 
 const (
-	imageAMD             = "gcr.io/gvisor-presubmit/benchmarks/rubydev_x86_64:latest"
-	imageARM             = "gcr.io/gvisor-presubmit/benchmarks/rubydev_aarch64:latest"
 	builderContainerName = "builder"
+	imageAMD             = k8s.ImageRepoPrefix + "benchmarks/rubydev_x86_64:latest"
+	imageARM             = k8s.ImageRepoPrefix + "benchmarks/rubydev_aarch64:latest"
 )
 
 // TestRubyDev benchmarks a build job on k8s clusters.
diff --git a/test/kubernetes/benchmarks/stablediffusion_test.go b/test/kubernetes/benchmarks/stablediffusion_test.go
index ce03612e9e..a61eb70579 100644
--- a/test/kubernetes/benchmarks/stablediffusion_test.go
+++ b/test/kubernetes/benchmarks/stablediffusion_test.go
@@ -23,6 +23,7 @@ import (
 	"time"
 
 	"gvisor.dev/gvisor/test/gpu/stablediffusion"
+	k8s "gvisor.dev/gvisor/test/kubernetes"
 	"gvisor.dev/gvisor/test/kubernetes/benchmetric"
 	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
@@ -32,7 +33,7 @@ import (
 
 const (
 	// Container image for Stable Diffusion XL.
-	stableDiffusionImage = "gcr.io/gvisor-presubmit/gpu/stable-diffusion-xl"
+	stableDiffusionImage = k8s.ImageRepoPrefix + "gpu/stable-diffusion-xl"
 )
 
 func TestStableDiffusionXL(t *testing.T) {
diff --git a/test/kubernetes/benchmarks/tensorflow_test.go b/test/kubernetes/benchmarks/tensorflow_test.go
index 03dd6f1754..39707d405e 100644
--- a/test/kubernetes/benchmarks/tensorflow_test.go
+++ b/test/kubernetes/benchmarks/tensorflow_test.go
@@ -21,6 +21,7 @@ import (
 	"testing"
 	"time"
 
+	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"
@@ -30,8 +31,8 @@ import (
 )
 
 const (
-	imageAMD = "gcr.io/gvisor-presubmit/benchmarks/tensorflow_x86_64:latest"
-	imageARM = "gcr.io/gvisor-presubmit/benchmarks/tensorflow_aarch64:latest"
+	imageAMD = k8s.ImageRepoPrefix + "benchmarks/tensorflow_x86_64:latest"
+	imageARM = k8s.ImageRepoPrefix + "benchmarks/tensorflow_aarch64:latest"
 )
 
 var workloads = map[string]string{
diff --git a/test/kubernetes/helpers.go b/test/kubernetes/helpers.go
new file mode 100644
index 0000000000..08336ea3be
--- /dev/null
+++ b/test/kubernetes/helpers.go
@@ -0,0 +1,20 @@
+// Copyright 2024 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package kubernetes contains utilities for gVisor tests.
+package kubernetes
+
+// ImageRepoPrefix is the image repository prefix of images that are
+// used by gVisor tests.
+const ImageRepoPrefix = "us-central1-docker.pkg.dev/gvisor-presubmit/"