testcluster: Remove all dependencies on cluster spec protobuffer.

Instead, cluster nodepools are reconstructed from Kubernetes client API
data and looking at their labels from that. Then these are used for
configuring pods to schedule on these nodes.

This reduces dependencies on GKE-specific implementation details and
makes benchmarks easier to run outside of GKE.

PiperOrigin-RevId: 700447594

test/kubernetes/benchmarks/abslbuild.go

@@ -53,7 +53,11 @@ func BuildABSL(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContex
 	defer cluster.DeletePersistentVolume(ctx, persistentVol)
 
 	image := imageAMD
-	if cluster.RuntimeTestNodepoolIsARM() {
+	testCPUArch, err := cluster.RuntimeTestNodepoolArchitecture(ctx)
+	if err != nil {
+		t.Fatalf("Failed to get runtime test nodepool architecture: %v", err)
+	}
+	if testCPUArch == testcluster.CPUArchitectureARM {
 		t.Skipf("Building ABSL is not supported on ARM")
 		return
 	}
@@ -91,14 +95,14 @@ func BuildABSL(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContex
 		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
-			endProfiling, err := profiling.MaybeSetup(ctx, t, cluster, benchmarkNS)
+			endProfiling, err := profiling.MaybeSetup(ctx, t, k8sCtx, cluster, benchmarkNS)
 			if err != nil {
 				t.Fatalf("Failed to setup profiling: %v", err)
 			}
 			defer endProfiling()
 
 			pod := newABSLPod(benchmarkNS, name, image, test.volume)
-			pod, err = cluster.ConfigurePodForRuntimeTestNodepool(pod)
+			pod, err = cluster.ConfigurePodForRuntimeTestNodepool(ctx, pod)
 			if err != nil {
 				t.Fatalf("Failed to set pod for test runtime: %v", err)
 			}

test/kubernetes/benchmarks/ffmpeg.go

@@ -53,9 +53,18 @@ func RunFFMPEG(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContex
 	}
 	defer cluster.DeletePersistentVolume(ctx, persistentVol)
 
-	image := imageAMD
-	if cluster.RuntimeTestNodepoolIsARM() {
+	testCPUArch, err := cluster.RuntimeTestNodepoolArchitecture(ctx)
+	if err != nil {
+		t.Fatalf("Failed to get runtime test nodepool architecture: %v", err)
+	}
+	var image string
+	switch testCPUArch {
+	case testcluster.CPUArchitectureX86:
+		image = imageAMD
+	case testcluster.CPUArchitectureARM:
 		image = imageARM
+	default:
+		t.Fatalf("Unsupported CPU architecture: %v", testCPUArch)
 	}
 	if image, err = k8sCtx.ResolveImage(ctx, image); err != nil {
 		t.Fatalf("Failed to resolve image: %v", err)
@@ -91,14 +100,14 @@ func RunFFMPEG(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContex
 		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
-			endProfiling, err := profiling.MaybeSetup(ctx, t, cluster, benchmarkNS)
+			endProfiling, err := profiling.MaybeSetup(ctx, t, k8sCtx, cluster, benchmarkNS)
 			if err != nil {
 				t.Fatalf("Failed to setup profiling: %v", err)
 			}
 			defer endProfiling()
 
 			p := newFfmpegDevPod(benchmarkNS, name, image, test.volume)
-			p, err = cluster.ConfigurePodForRuntimeTestNodepool(p)
+			p, err = cluster.ConfigurePodForRuntimeTestNodepool(ctx, p)
 			if err != nil {
 				t.Fatalf("Failed to configure pod for runtime: %v", err)
 			}

test/kubernetes/benchmarks/grpc.go

@@ -53,9 +53,18 @@ func BuildGRPC(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContex
 	}
 	defer cluster.DeletePersistentVolume(ctx, persistentVol)
 
-	image := imageAMD
-	if cluster.RuntimeTestNodepoolIsARM() {
+	testCPUArch, err := cluster.RuntimeTestNodepoolArchitecture(ctx)
+	if err != nil {
+		t.Fatalf("Failed to get runtime test nodepool architecture: %v", err)
+	}
+	var image string
+	switch testCPUArch {
+	case testcluster.CPUArchitectureX86:
+		image = imageAMD
+	case testcluster.CPUArchitectureARM:
 		image = imageARM
+	default:
+		t.Fatalf("Unsupported CPU architecture: %v", testCPUArch)
 	}
 	if image, err = k8sCtx.ResolveImage(ctx, image); err != nil {
 		t.Fatalf("Failed to resolve image: %v", err)
@@ -91,14 +100,14 @@ func BuildGRPC(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContex
 		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
-			endProfiling, err := profiling.MaybeSetup(ctx, t, cluster, benchmarkNS)
+			endProfiling, err := profiling.MaybeSetup(ctx, t, k8sCtx, cluster, benchmarkNS)
 			if err != nil {
 				t.Fatalf("Failed to setup profiling: %v", err)
 			}
 			defer endProfiling()
 
 			pod := newGRPCPod(benchmarkNS, name, image, test.volume)
-			pod, err = cluster.ConfigurePodForRuntimeTestNodepool(pod)
+			pod, err = cluster.ConfigurePodForRuntimeTestNodepool(ctx, pod)
 			if err != nil {
 				t.Fatalf("Failed to set pod for test runtime: %v", err)
 			}

test/kubernetes/benchmarks/gsutil.go

@@ -56,9 +56,18 @@ func RunGSUtil(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContex
 	}
 	defer cluster.DeletePersistentVolume(ctx, persistentVol)
 
-	image := imageAMD
-	if cluster.RuntimeTestNodepoolIsARM() {
+	testCPUArch, err := cluster.RuntimeTestNodepoolArchitecture(ctx)
+	if err != nil {
+		t.Fatalf("Failed to get runtime test nodepool architecture: %v", err)
+	}
+	var image string
+	switch testCPUArch {
+	case testcluster.CPUArchitectureX86:
+		image = imageAMD
+	case testcluster.CPUArchitectureARM:
 		image = imageARM
+	default:
+		t.Fatalf("Unsupported CPU architecture: %v", testCPUArch)
 	}
 	if image, err = k8sCtx.ResolveImage(ctx, image); err != nil {
 		t.Fatalf("Failed to resolve image: %v", err)
@@ -113,7 +122,7 @@ func RunGSUtil(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContex
 			} {
 				t.Run(slicing.name, func(t *testing.T) {
 					// Setup profiling if requested by the user.
-					endProfiling, err := profiling.MaybeSetup(ctx, t, cluster, benchmarkNS)
+					endProfiling, err := profiling.MaybeSetup(ctx, t, k8sCtx, cluster, benchmarkNS)
 					if err != nil {
 						t.Fatalf("Failed to setup profiling: %v", err)
 					}
@@ -122,7 +131,7 @@ func RunGSUtil(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContex
 					// Create a pod that performs setup, then times
 					// downloading.
 					p := newGSUtilDevPod(benchmarkNS, name, image, storage.volume, slicing.option)
-					p, err = cluster.ConfigurePodForRuntimeTestNodepool(p)
+					p, err = cluster.ConfigurePodForRuntimeTestNodepool(ctx, p)
 					if err != nil {
 						t.Fatalf("Failed to configure pod for runtime: %v", err)
 					}

test/kubernetes/benchmarks/httpbench/httpbench.go

@@ -135,7 +135,7 @@ func (h *HTTPBenchmark) runRound(ctx context.Context, t *testing.T, round Round,
 	}
 	name := fmt.Sprintf("wrk2-%dthreads-%sqps", round.NumThreads, qpsText)
 	client := h.newWrk2Client(name, ip, round)
-	client, err := h.Cluster.ConfigurePodForClientNodepool(client)
+	client, err := h.Cluster.ConfigurePodForClientNodepool(ctx, client)
 	if err != nil {
 		t.Fatalf("failed to configure wrk2 pod for client nodepool: %v", err)
 	}
@@ -243,7 +243,7 @@ func (h *HTTPBenchmark) getWgetPod(ip string) *v13.Pod {
 // waitForServer waits for an HTTP server to start responding on the given
 // IP and port.
 func (h *HTTPBenchmark) waitForServer(ctx context.Context, ip string) error {
-	wget, err := h.Cluster.ConfigurePodForClientNodepool(h.getWgetPod(ip))
+	wget, err := h.Cluster.ConfigurePodForClientNodepool(ctx, h.getWgetPod(ip))
 	if err != nil {
 		return fmt.Errorf("failed to configure wget pod for client nodepool: %v", err)
 	}

test/kubernetes/benchmarks/nginx.go

@@ -62,12 +62,20 @@ func BenchmarkNginx(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesC
 	}
 	defer benchmarkNS.Cleanup(ctx)
 
-	nginxImage := nginxImageAMD
-	if cluster.RuntimeTestNodepoolIsARM() {
-		nginxImage = nginxImageARM
-	}
-	nginxImage, err := k8sCtx.ResolveImage(ctx, nginxImage)
+	testCPUArch, err := cluster.RuntimeTestNodepoolArchitecture(ctx)
 	if err != nil {
+		t.Fatalf("Failed to get runtime test nodepool architecture: %v", err)
+	}
+	var nginxImage string
+	switch testCPUArch {
+	case testcluster.CPUArchitectureX86:
+		nginxImage = nginxImageAMD
+	case testcluster.CPUArchitectureARM:
+		nginxImage = nginxImageARM
+	default:
+		t.Fatalf("Unsupported CPU architecture: %v", testCPUArch)
+	}
+	if nginxImage, err = k8sCtx.ResolveImage(ctx, nginxImage); err != nil {
 		t.Fatalf("Failed to resolve image: %v", err)
 	}
 
@@ -114,7 +122,7 @@ func BenchmarkNginx(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesC
 		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
-			endProfiling, err := profiling.MaybeSetup(ctx, t, cluster, benchmarkNS)
+			endProfiling, err := profiling.MaybeSetup(ctx, t, k8sCtx, cluster, benchmarkNS)
 			if err != nil {
 				t.Fatalf("Failed to setup profiling: %v", err)
 			}
@@ -123,7 +131,7 @@ func BenchmarkNginx(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesC
 			name := fmt.Sprintf("nginx-%s", test.suffix)
 
 			server := newNginxServer(benchmarkNS, name, nginxImage, test.volume)
-			server, err = cluster.ConfigurePodForRuntimeTestNodepool(server)
+			server, err = cluster.ConfigurePodForRuntimeTestNodepool(ctx, server)
 			if err != nil {
 				t.Fatalf("Failed to configure pod for runtime nodepool: %v", err)
 			}

test/kubernetes/benchmarks/ollama.go

@@ -164,7 +164,7 @@ func (ops *ollamaPodServer) InstrumentedRequest(ctx context.Context, argvFn func
 			RestartPolicy: v13.RestartPolicyNever,
 		},
 	}
-	clientPod, err := ops.cluster.ConfigurePodForClientNodepool(clientPod)
+	clientPod, err := ops.cluster.ConfigurePodForClientNodepool(ctx, clientPod)
 	if err != nil {
 		return nil, fmt.Errorf("failed to configure pod: %v", err)
 	}
@@ -245,7 +245,7 @@ func BenchmarkOllama(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kubernetes
 	if err != nil {
 		t.Fatalf("Failed to resolve image: %v", err)
 	}
-	ollamaPod, err := cluster.ConfigurePodForRuntimeTestNodepool(newOllamaServerPod(benchmarkNS, serverImage))
+	ollamaPod, err := cluster.ConfigurePodForRuntimeTestNodepool(ctx, newOllamaServerPod(benchmarkNS, serverImage))
 	if err != nil {
 		t.Fatalf("Failed to configure pod for runtime nodepool: %v", err)
 	}

test/kubernetes/benchmarks/postgresql.go

@@ -56,7 +56,7 @@ func BenchmarkPostgresPGBench(ctx context.Context, t *testing.T, k8sCtx k8sctx.K
 		t.Fatalf("cannot reset namespace: %v", err)
 	}
 	defer benchmarkNS.Cleanup(ctx)
-	endProfiling, err := profiling.MaybeSetup(ctx, t, cluster, benchmarkNS)
+	endProfiling, err := profiling.MaybeSetup(ctx, t, k8sCtx, cluster, benchmarkNS)
 	if err != nil {
 		t.Fatalf("Failed to setup profiling: %v", err)
 	}
@@ -80,7 +80,7 @@ func BenchmarkPostgresPGBench(ctx context.Context, t *testing.T, k8sCtx k8sctx.K
 		server.ObjectMeta.Labels = make(map[string]string)
 	}
 	server.ObjectMeta.Labels[postgresServerLabelKey] = postgresServerLabelValue
-	server, err = cluster.ConfigurePodForRuntimeTestNodepool(server)
+	server, err = cluster.ConfigurePodForRuntimeTestNodepool(ctx, server)
 	if err != nil {
 		t.Fatalf("ConfigurePodForRuntimeTestNodepool on cluster %q: %v", cluster.GetName(), err)
 	}
@@ -127,7 +127,7 @@ func BenchmarkPostgresPGBench(ctx context.Context, t *testing.T, k8sCtx k8sctx.K
 			fmt.Sprintf("--username=%s", postgresUser),
 			fmt.Sprintf("--dbname=%s", postgresDatabase),
 		}, false /* withPort */, nil /* pvc */)
-		pgIsReady, err = cluster.ConfigurePodForClientNodepool(pgIsReady)
+		pgIsReady, err = cluster.ConfigurePodForClientNodepool(ctx, pgIsReady)
 		if err != nil {
 			return fmt.Errorf("ConfigurePodForClientNodepool on cluster %q: pod: %q: %v", cluster.GetName(), pgIsReadyName, err)
 		}
@@ -170,7 +170,7 @@ func BenchmarkPostgresPGBench(ctx context.Context, t *testing.T, k8sCtx k8sctx.K
 			fmt.Sprintf("--username=%s", postgresUser),
 			postgresDatabase,
 		}, false /* withPort */, nil /* pvc */)
-		initDB, err = cluster.ConfigurePodForClientNodepool(initDB)
+		initDB, err = cluster.ConfigurePodForClientNodepool(ctx, initDB)
 		if err != nil {
 			return fmt.Errorf("ConfigurePodForClientNodepool on cluster %q: pod: %q: %v", cluster.GetName(), initDBName, err)
 		}
@@ -206,7 +206,7 @@ func BenchmarkPostgresPGBench(ctx context.Context, t *testing.T, k8sCtx k8sctx.K
 				postgresDatabase,
 			}
 			client := newPostgresPod(benchmarkNS, "pgbench", image, clientCmd, false /* withPort */, nil /* pvc */)
-			client, err = cluster.ConfigurePodForClientNodepool(client)
+			client, err = cluster.ConfigurePodForClientNodepool(ctx, client)
 			if err != nil {
 				t.Fatalf("ConfigurePodForClientNodepool on cluster %q: pod: %q: %v", cluster.GetName(), client.GetName(), err)
 			}

test/kubernetes/benchmarks/profiling/BUILD

@@ -17,6 +17,7 @@ go_library(
         "//pkg/atomicbitops",
         "//runsc/flag",
         "//test/kubernetes",
+        "//test/kubernetes/k8sctx",
         "//test/kubernetes/testcluster",
         "//test/metricsviz",
         "@com_github_google_pprof//profile:go_default_library",

test/kubernetes/benchmarks/profiling/profiling.go

@@ -43,6 +43,7 @@ import (
 	"gvisor.dev/gvisor/pkg/atomicbitops"
 	"gvisor.dev/gvisor/runsc/flag"
 	k8s "gvisor.dev/gvisor/test/kubernetes"
+	"gvisor.dev/gvisor/test/kubernetes/k8sctx"
 	"gvisor.dev/gvisor/test/kubernetes/testcluster"
 	"gvisor.dev/gvisor/test/metricsviz"
 	appsv1 "k8s.io/api/apps/v1"
@@ -203,17 +204,29 @@ func streamDir(dirPath string) operation {
 }
 
 // startsOperations starts the given operations in a DaemonSet.
-func startOperations(ctx context.Context, c *testcluster.TestCluster, ns *testcluster.Namespace, operations []operation) (*appsv1.DaemonSet, error) {
+func startOperations(ctx context.Context, k8sCtx k8sctx.KubernetesContext, c *testcluster.TestCluster, ns *testcluster.Namespace, operations []operation) (*appsv1.DaemonSet, error) {
 	ds := profileDSTemplate(c)
 	ds.Namespace = ns.Namespace
 	ds.ObjectMeta.Namespace = ns.Namespace
 	ds.Spec.Template.Namespace = ns.Namespace
 	ds.Spec.Template.ObjectMeta.Namespace = ns.Namespace
-	c.ConfigureDaemonSetForRuntimeTestNodepool(&ds)
+	c.ConfigureDaemonSetForRuntimeTestNodepool(ctx, &ds)
 	ds.Spec.Template.Spec.RuntimeClassName = nil // Must run unsandboxed.
-	image := profileHelperImageAMD64
-	if c.RuntimeTestNodepoolIsARM() {
+	var image string
+	testCPUArch, err := c.RuntimeTestNodepoolArchitecture(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to determine test CPU architecture: %w", err)
+	}
+	switch testCPUArch {
+	case testcluster.CPUArchitectureX86:
+		image = profileHelperImageAMD64
+	case testcluster.CPUArchitectureARM:
 		image = profileHelperImageARM64
+	default:
+		return nil, fmt.Errorf("unsupported CPU architecture: %q", testCPUArch)
+	}
+	if image, err = k8sCtx.ResolveImage(ctx, image); err != nil {
+		return nil, fmt.Errorf("failed to resolve image %q: %w", image, err)
 	}
 	for i, op := range operations {
 		name := op.name
@@ -247,6 +260,7 @@ func startOperations(ctx context.Context, c *testcluster.TestCluster, ns *testcl
 // profileRun encapsulates data about a profiling run.
 // It is used after the run completes so that profiles can be retrieved.
 type profileRun struct {
+	k8sCtx                k8sctx.KubernetesContext
 	c                     *testcluster.TestCluster
 	ns                    *testcluster.Namespace
 	localProfileDir       string
@@ -256,7 +270,7 @@ type profileRun struct {
 // MaybeSetup sets up profiling if requested. It returns a cleanup function.
 // If the returned error is nil, the cleanup function is non-nil and should be
 // called regardless of whether profiling is actually enabled or not.
-func MaybeSetup(ctx context.Context, t *testing.T, c *testcluster.TestCluster, ns *testcluster.Namespace) (func(), error) {
+func MaybeSetup(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContext, c *testcluster.TestCluster, ns *testcluster.Namespace) (func(), error) {
 	profileDirName := fmt.Sprintf("%s.%s", t.Name(), time.Now().Format("20060102-150405"))
 	profileDirName = regexp.MustCompile("[^-_=.\\w]+").ReplaceAllString(profileDirName, ".")
 	hasGVisorRuntime, err := c.HasGVisorTestRuntime(ctx)
@@ -316,6 +330,7 @@ func MaybeSetup(ctx context.Context, t *testing.T, c *testcluster.TestCluster, n
 		}
 		cleanup = func() {
 			err := processProfileRun(ctx, t, &profileRun{
+				k8sCtx:                k8sCtx,
 				c:                     c,
 				ns:                    ns,
 				localProfileDir:       localProfileDir,
@@ -340,7 +355,7 @@ func MaybeSetup(ctx context.Context, t *testing.T, c *testcluster.TestCluster, n
 	if len(setupCommands) > 0 {
 		setupCtx, setupCancel := context.WithTimeout(ctx, 2*time.Minute)
 		defer setupCancel()
-		ds, err := startOperations(setupCtx, c, ns, setupCommands)
+		ds, err := startOperations(setupCtx, k8sCtx, c, ns, setupCommands)
 		if err != nil {
 			return nil, err
 		}
@@ -359,7 +374,7 @@ func processProfileRun(ctx context.Context, t *testing.T, run *profileRun) error
 	beforeSpawn := metav1.NewTime(time.Now())
 	retrievalCtx, retrievalCancel := context.WithCancel(ctx)
 	defer retrievalCancel()
-	ds, err := startOperations(retrievalCtx, run.c, run.ns, []operation{
+	ds, err := startOperations(retrievalCtx, run.k8sCtx, run.c, run.ns, []operation{
 		dirOp,
 		setFlag("profile", "false"),
 		removeFlag("profile-cpu"),

test/kubernetes/benchmarks/pytorch.go

@@ -257,7 +257,7 @@ func RunPytorch(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesConte
 // doPytorchRun runs a single PyTorch test.
 func doPytorchRun(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContext, cluster *testcluster.TestCluster, params pytorchTest) {
 	benchmarkNS := cluster.Namespace(testcluster.NamespaceBenchmark)
-	endProfiling, err := profiling.MaybeSetup(ctx, t, cluster, benchmarkNS)
+	endProfiling, err := profiling.MaybeSetup(ctx, t, k8sCtx, cluster, benchmarkNS)
 	if err != nil {
 		t.Fatalf("Failed to setup profiling: %v", err)
 	}
@@ -276,7 +276,7 @@ func doPytorchRun(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesCon
 		t.Fatalf("Failed to create pod: %v", err)
 	}
 
-	pod, err = cluster.ConfigurePodForRuntimeTestNodepool(pod)
+	pod, err = cluster.ConfigurePodForRuntimeTestNodepool(ctx, pod)
 	if err != nil {
 		t.Fatalf("Failed to configure pod for test-nodepool: %v", err)
 	}

test/kubernetes/benchmarks/redis.go

@@ -69,9 +69,18 @@ func BenchmarkRedis(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesC
 	}
 	defer cluster.DeletePersistentVolume(ctx, persistentVol)
 
-	image := redisImageAMD
-	if cluster.RuntimeTestNodepoolIsARM() {
+	testCPUArch, err := cluster.RuntimeTestNodepoolArchitecture(ctx)
+	if err != nil {
+		t.Fatalf("Failed to get runtime test nodepool architecture: %v", err)
+	}
+	var image string
+	switch testCPUArch {
+	case testcluster.CPUArchitectureX86:
+		image = redisImageAMD
+	case testcluster.CPUArchitectureARM:
 		image = redisImageARM
+	default:
+		t.Fatalf("Unsupported CPU architecture: %v", testCPUArch)
 	}
 	if image, err = k8sCtx.ResolveImage(ctx, image); err != nil {
 		t.Fatalf("Failed to resolve image: %v", err)
@@ -110,7 +119,7 @@ func BenchmarkRedis(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesC
 		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
-			endProfiling, err := profiling.MaybeSetup(ctx, t, cluster, benchmarkNS)
+			endProfiling, err := profiling.MaybeSetup(ctx, t, k8sCtx, cluster, benchmarkNS)
 			if err != nil {
 				t.Fatalf("Failed to setup profiling: %v", err)
 			}
@@ -122,7 +131,7 @@ func BenchmarkRedis(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesC
 				server.ObjectMeta.Labels = make(map[string]string)
 			}
 			server.ObjectMeta.Labels[redisServerLabelKey] = redisServerLabelValue
-			server, err = cluster.ConfigurePodForRuntimeTestNodepool(server)
+			server, err = cluster.ConfigurePodForRuntimeTestNodepool(ctx, server)
 			if err != nil {
 				t.Fatalf("ConfigurePodForRuntimeTestNodepool on cluster %q: %v", cluster.GetName(), err)
 			}
@@ -164,7 +173,7 @@ func BenchmarkRedis(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesC
 			pingCmd := []string{"redis-cli", "-h", ip, "-r", "5", "-i", "1", "ping"}
 			ensureUp := func() error {
 				pinger := newRedisPod(benchmarkNS, fmt.Sprintf("rpinger-%s", test.suffix), image, pingCmd)
-				pinger, err = cluster.ConfigurePodForClientNodepool(pinger)
+				pinger, err = cluster.ConfigurePodForClientNodepool(ctx, pinger)
 				if err != nil {
 					return fmt.Errorf("ConfigurePodForClientNodepool on cluster %q: pod: %q: %v", cluster.GetName(), pinger.GetName(), err)
 				}
@@ -226,7 +235,7 @@ func BenchmarkRedis(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesC
 								"--precision", "4", // Floating-point precision for reporting latency (in ms)
 							}
 							client := newRedisPod(benchmarkNS, "client", image, clientCmd)
-							client, err = cluster.ConfigurePodForClientNodepool(client)
+							client, err = cluster.ConfigurePodForClientNodepool(ctx, client)
 							if err != nil {
 								t.Fatalf("ConfigurePodForClientNodepool on cluster %q: pod: %q: %v", cluster.GetName(), client.GetName(), err)
 							}

test/kubernetes/benchmarks/rubydev.go

@@ -59,9 +59,18 @@ func RunRubyDev(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesConte
 	}
 	defer cluster.DeletePersistentVolume(ctx, persistentVol)
 
-	image := imageAMD
-	if cluster.RuntimeTestNodepoolIsARM() {
+	testCPUArch, err := cluster.RuntimeTestNodepoolArchitecture(ctx)
+	if err != nil {
+		t.Fatalf("Failed to get runtime test nodepool architecture: %v", err)
+	}
+	var image string
+	switch testCPUArch {
+	case testcluster.CPUArchitectureX86:
+		image = imageAMD
+	case testcluster.CPUArchitectureARM:
 		image = imageARM
+	default:
+		t.Fatalf("Unsupported CPU architecture: %v", testCPUArch)
 	}
 	if image, err = k8sCtx.ResolveImage(ctx, image); err != nil {
 		t.Fatalf("failed to resolve image: %v", err)
@@ -96,14 +105,14 @@ func RunRubyDev(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesConte
 		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
-			endProfiling, err := profiling.MaybeSetup(ctx, t, cluster, benchmarkNS)
+			endProfiling, err := profiling.MaybeSetup(ctx, t, k8sCtx, cluster, benchmarkNS)
 			if err != nil {
 				t.Fatalf("Failed to setup profiling: %v", err)
 			}
 			defer endProfiling()
 			// create a new RubyDevPod and set it to run on the runtime under test nodepool.
 			pod := newRubyDevPod(benchmarkNS, name, image, test.volume)
-			pod, err = cluster.ConfigurePodForRuntimeTestNodepool(pod)
+			pod, err = cluster.ConfigurePodForRuntimeTestNodepool(ctx, pod)
 			if err != nil {
 				t.Fatalf("failed to configure pod for test runtime node: %v", err)
 			}

test/kubernetes/benchmarks/stablediffusion.go

@@ -66,7 +66,7 @@ func (r *kubernetesPodRunner) Run(ctx context.Context, image string, argv []stri
 			RestartPolicy: v13.RestartPolicyNever,
 		},
 	}
-	stableDiffusionXLPod, err := r.cluster.ConfigurePodForRuntimeTestNodepool(stableDiffusionXLPod)
+	stableDiffusionXLPod, err := r.cluster.ConfigurePodForRuntimeTestNodepool(ctx, stableDiffusionXLPod)
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to configure pod: %v", err)
 	}

test/kubernetes/benchmarks/startup.go

@@ -48,7 +48,7 @@ func MeasureStartup(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesC
 	t.Logf("Warning: This is not a meaningful benchmark. Read the comments.")
 
 	benchmarkNS := cluster.Namespace(testcluster.NamespaceBenchmark)
-	endProfiling, err := profiling.MaybeSetup(ctx, t, cluster, benchmarkNS)
+	endProfiling, err := profiling.MaybeSetup(ctx, t, k8sCtx, cluster, benchmarkNS)
 	if err != nil {
 		t.Fatalf("Failed to setup profiling: %v", err)
 	}
@@ -63,7 +63,7 @@ func MeasureStartup(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesC
 	if err != nil {
 		t.Fatalf("Failed to resolve image: %v", err)
 	}
-	p, err := cluster.ConfigurePodForRuntimeTestNodepool(benchmarkNS.NewAlpinePod(podName, image, command))
+	p, err := cluster.ConfigurePodForRuntimeTestNodepool(ctx, benchmarkNS.NewAlpinePod(podName, image, command))
 	if err != nil {
 		t.Fatalf("failed to set pod for test nodepool: %v", err)
 	}

test/kubernetes/benchmarks/tensorflow.go

@@ -48,7 +48,7 @@ var workloads = map[string]string{
 // RunTensorflowOnCPU runs the Tensorflow example workloads on CPU.
 func RunTensorflowOnCPU(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContext, cluster *testcluster.TestCluster) {
 	benchmarkNS := cluster.Namespace(testcluster.NamespaceBenchmark)
-	endProfiling, err := profiling.MaybeSetup(ctx, t, cluster, benchmarkNS)
+	endProfiling, err := profiling.MaybeSetup(ctx, t, k8sCtx, cluster, benchmarkNS)
 	if err != nil {
 		t.Fatalf("Failed to setup profiling: %v", err)
 	}
@@ -64,9 +64,18 @@ func RunTensorflowOnCPU(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kuberne
 		t.Fatalf("Failed to initialize benchmark recorder: %v", err)
 	}
 
-	image := imageAMD
-	if cluster.RuntimeTestNodepoolIsARM() {
+	testCPUArch, err := cluster.RuntimeTestNodepoolArchitecture(ctx)
+	if err != nil {
+		t.Fatalf("Failed to get runtime test nodepool architecture: %v", err)
+	}
+	var image string
+	switch testCPUArch {
+	case testcluster.CPUArchitectureX86:
+		image = imageAMD
+	case testcluster.CPUArchitectureARM:
 		image = imageARM
+	default:
+		t.Fatalf("Unsupported CPU architecture: %v", testCPUArch)
 	}
 	if image, err = k8sCtx.ResolveImage(ctx, image); err != nil {
 		t.Fatalf("Failed to resolve image: %v", err)
@@ -85,7 +94,7 @@ func RunTensorflowOnCPU(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kuberne
 		workloadName := workloadPathToName[workloadPath]
 		t.Run(workloadName, func(t *testing.T) {
 			pod := newTensorflowOnCPUPod(benchmarkNS, name, image, workloadPath)
-			pod, err := cluster.ConfigurePodForRuntimeTestNodepool(pod)
+			pod, err := cluster.ConfigurePodForRuntimeTestNodepool(ctx, pod)
 			if err != nil {
 				t.Fatalf("Failed to set pod for test runtime: %v", err)
 			}

test/kubernetes/benchmarks/wordpress.go

@@ -60,7 +60,7 @@ var (
 // BenchmarkWordpress runs a benchmark of WordPress performance.
 func BenchmarkWordpress(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContext, cluster *testcluster.TestCluster) {
 	benchmarkNS := cluster.Namespace(testcluster.NamespaceBenchmark)
-	endProfiling, err := profiling.MaybeSetup(ctx, t, cluster, benchmarkNS)
+	endProfiling, err := profiling.MaybeSetup(ctx, t, k8sCtx, cluster, benchmarkNS)
 	if err != nil {
 		t.Fatalf("Failed to setup profiling: %v", err)
 	}
@@ -84,7 +84,7 @@ func BenchmarkWordpress(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kuberne
 		t.Fatalf("failed to resolve image: %v", err)
 	}
 	database := newMariaDBServer(benchmarkNS, databaseName, mariaDBImg, dbVolume)
-	database, err = cluster.ConfigurePodForTertiaryNodepool(database)
+	database, err = cluster.ConfigurePodForTertiaryNodepool(ctx, database)
 	if err != nil {
 		t.Fatalf("Failed to configure pod for tertiary nodepool: %v", err)
 	}
@@ -110,7 +110,7 @@ func BenchmarkWordpress(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kuberne
 		t.Fatalf("Failed to resolve image: %v", err)
 	}
 	server := newWordpressServer(benchmarkNS, name, wordpressImg, mariaDBIP)
-	server, err = cluster.ConfigurePodForRuntimeTestNodepool(server)
+	server, err = cluster.ConfigurePodForRuntimeTestNodepool(ctx, server)
 	if err != nil {
 		t.Fatalf("Failed to configure pod for runtime nodepool: %v", err)
 	}
@@ -138,7 +138,7 @@ func BenchmarkWordpress(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kuberne
 
 	// Install WordPress.
 	installWordpressPod := newWordpressInstall(benchmarkNS, "install-wordpress", wordpressIP)
-	installWordpressPod, err = cluster.ConfigurePodForClientNodepool(installWordpressPod)
+	installWordpressPod, err = cluster.ConfigurePodForClientNodepool(ctx, installWordpressPod)
 	if err != nil {
 		t.Fatalf("Failed to configure pod for client nodepool: %v", err)
 	}

test/kubernetes/testcluster/BUILD

@@ -15,6 +15,7 @@ go_library(
         "//visibility:public",
     ],
     deps = [
+        "//pkg/sync",
         "//test/kubernetes:test_range_config_go_proto",
         "@io_k8s_api//apps/v1:go_default_library",
         "@io_k8s_api//core/v1:go_default_library",

test/kubernetes/testcluster/testcluster.go

@@ -20,13 +20,13 @@ import (
 	"context"
 	"fmt"
 	"io"
-	"reflect"
 	"strconv"
 	"strings"
 	"time"
 
 	"golang.org/x/sync/errgroup"
 	cspb "google.golang.org/genproto/googleapis/container/v1"
+	"gvisor.dev/gvisor/pkg/sync"
 	testpb "gvisor.dev/gvisor/test/kubernetes/test_range_config_go_proto"
 	appsv1 "k8s.io/api/apps/v1"
 	v13 "k8s.io/api/core/v1"
@@ -40,8 +40,6 @@ import (
 const (
 	// archKey is given to nodepools to mark their architecture. Used here to mark ARM nodepools.
 	archKey = "kubernetes.io/arch"
-	// armValue marks an ARM nodepool.
-	armValue = "arm64"
 
 	// k8sApp is used as a label to distinguish between applications.
 	k8sApp = "k8s-app"
@@ -59,18 +57,21 @@ const (
 	NamespaceBenchmark = "benchmark"
 )
 
+// NodePoolType is the type of a NodePool.
+type NodePoolType string
+
 // Nodepool names.
 const (
 	// TestRuntimeNodepoolName is the value that marks a "test-runtime-nodepool", or a nodepool where
 	// w/ the runtime under test.
-	TestRuntimeNodepoolName = "test-runtime-nodepool"
+	TestRuntimeNodepoolName NodePoolType = "test-runtime-nodepool"
 	// ClientNodepoolName is the value that marks a client nodepool. Usually this is a plain GKE
 	// nodepool
-	ClientNodepoolName = "client-nodepool"
+	ClientNodepoolName NodePoolType = "client-nodepool"
 	// TertiaryNodepoolName is the value that marks the tertiary nodepool.
 	// This could either be a plain GKE nodepool or could be gVisor-enabled,
 	// as configured during test range creation.
-	TertiaryNodepoolName = "tertiary-nodepool"
+	TertiaryNodepoolName NodePoolType = "tertiary-nodepool"
 )
 
 // Nodepool keys.
@@ -83,8 +84,8 @@ const (
 	NodepoolNumAcceleratorsKey = "num-accelerators"
 	// NodepoolTPUTopologyKey is the key to mark the TPU topology used by a nodepool.
 	NodepoolTPUTopologyKey = "tpu-topology"
-	// Name of the nodepool key used in Pod.Spec.NodeSelector.
-	NodePoolSelectorKey = "cloud.google.com/gke-nodepool"
+	// NodepoolInstanceTypeKey is the key to mark the instance type used by a nodepool.
+	NodepoolInstanceTypeKey = "node.kubernetes.io/instance-type"
 	// Name of the TPU accelerator key used in Pod.Spec.NodeSelector.
 	NodepoolTPUAcceleratorSelectorKey = "cloud.google.com/gke-tpu-accelerator"
 	// Name of the TPU topology key used in Pod.Spec.NodeSelector.
@@ -111,6 +112,17 @@ const (
 	gvisorRuntimeClass = "gvisor"
 )
 
+// CPUArchitecture is the CPU architecture of a node.
+// It is stored under the archKey label in node labels.
+type CPUArchitecture string
+
+const (
+	// CPUArchitectureX86 is the x86 CPU architecture.
+	CPUArchitectureX86 = CPUArchitecture("amd64")
+	// CPUArchitectureARM is the ARM CPU architecture.
+	CPUArchitectureARM = CPUArchitecture("arm64")
+)
+
 // AcceleratorType is the gpu type to be used.
 type AcceleratorType string
 
@@ -124,40 +136,58 @@ const (
 
 // TestCluster wraps clusters with their individual ClientSets so that helper methods can be called.
 type TestCluster struct {
-	cluster *testpb.Cluster
-	client  kubernetes.Interface
+	clusterName string
+	client      kubernetes.Interface
 
 	// testNodepoolRuntimeOverride, if set, overrides the runtime used for pods
 	// running on the test nodepool. If unset, the test nodepool's default
 	// runtime is used.
 	testNodepoolRuntimeOverride RuntimeType
+
+	// nodepoolsMu controls the initialization of `nodepools`.
+	nodepoolsMu sync.Mutex
+
+	// nodepools is a map of NodePools that exist in this cluster.
+	// It is nil by default and initialized lazily.
+	nodepools map[NodePoolType]*NodePool
 }
 
-type testClusterConstructorKey int
+// NodePool is a set of nodes in a TestCluster.
+// These nodes share a set of relevant labels and are used to segment the
+// set of nodes in a Kubernetes cluster.
+// In the context of Kubernetes tests and benchmarks, these pools are used
+// to separate where workloads of each type schedule and run.
+// NodePools are expected to be uniform (i.e. same amount of resources and
+// reasonably similar hardware) so that simple pod scheduling can determine
+// where to consume resources.
+type NodePool struct {
+	// nodePoolType is the type of the nodepool.
+	// It is used to identify the nodes in the cluster, and therefore as a
+	// scheduling constraint for pods to run exclusively on these nodes.
+	nodePooltype NodePoolType
 
-const (
-	// testClusterConstructor is the key for the context value that holds the
-	// constructor function for TestCluster.
-	// Defaults to newTestCluster.
-	testClusterConstructor testClusterConstructorKey = iota
-)
+	// runtime is the container runtime to use when scheduling pods on this
+	// nodepool by default.
+	runtime RuntimeType
 
-// WithTestClusterConstructor returns a context that contains a custom
-// constructor for TestCluster.
-func WithTestClusterConstructor(ctx context.Context, constructor func(context.Context, *testpb.Cluster) (*TestCluster, error)) context.Context {
-	return context.WithValue(ctx, testClusterConstructor, constructor)
+	// cpuArchitecture is the CPU architecture of nodes in the nodepool.
+	cpuArchitecture CPUArchitecture
+
+	// acceleratorType is the accelerator type present on nodes in the nodepool.
+	// Empty string if the nodes have no accelerators.
+	acceleratorType AcceleratorType
+
+	// numAccelerators is the number of accelerators present on nodes in the
+	// nodepool.
+	numAccelerators int
+
+	// tpuTopology is the TPU topology used by the nodepool.
+	// Empty string if the nodepool has no TPU-based accelerators.
+	tpuTopology string
 }
 
-// NewTestCluster returns a new TestCluster client.
-func NewTestCluster(ctx context.Context, cluster *testpb.Cluster) (*TestCluster, error) {
-	constructor, ok := ctx.Value(testClusterConstructor).(func(context.Context, *testpb.Cluster) (*TestCluster, error))
-	if !ok || constructor == nil || reflect.ValueOf(constructor).IsNil() {
-		constructor = newTestCluster
-	}
-	return constructor(ctx, cluster)
-}
-
-func newTestCluster(_ context.Context, cluster *testpb.Cluster) (*TestCluster, error) {
+// NewTestClusterFromProto returns a new TestCluster client from a proto.
+func NewTestClusterFromProto(ctx context.Context, cluster *testpb.Cluster) (*TestCluster, error) {
 	config, err := clientcmd.BuildConfigFromFlags("" /*masterURL*/, cluster.GetCredentialFile())
 	if err != nil {
 		return nil, fmt.Errorf("BuildConfigFromFlags: %w", err)
@@ -166,37 +196,26 @@ func newTestCluster(_ context.Context, cluster *testpb.Cluster) (*TestCluster, e
 	if err != nil {
 		return nil, fmt.Errorf("kubernetes.NewForConfig: %w", err)
 	}
-	return NewTestClusterWithClient(cluster, client), nil
+	var clusterPB cspb.Cluster
+	if err := cluster.GetCluster().UnmarshalTo(&clusterPB); err != nil {
+		return nil, fmt.Errorf("cannot unmarshal cluster: %w", err)
+	}
+	clusterName := clusterPB.GetName()
+	return NewTestClusterFromClient(clusterName, client), nil
 }
 
-// NewTestClusterWithClient returns a new TestCluster client with a given client.
-func NewTestClusterWithClient(cluster *testpb.Cluster, client kubernetes.Interface) *TestCluster {
+// NewTestClusterFromClient returns a new TestCluster client with a given client.
+func NewTestClusterFromClient(clusterName string, client kubernetes.Interface) *TestCluster {
 	return &TestCluster{
-		cluster:                     cluster,
+		clusterName:                 clusterName,
 		client:                      client,
 		testNodepoolRuntimeOverride: "",
 	}
 }
 
-// Cluster returns the underlying cluster proto for tests.
-func (t *TestCluster) Cluster() *testpb.Cluster {
-	return t.cluster
-}
-
-// ContainerCluster returns the underlying container cluster proto.
-func (t *TestCluster) ContainerCluster() (*cspb.Cluster, error) {
-	var cluster cspb.Cluster
-	err := t.cluster.GetCluster().UnmarshalTo(&cluster)
-	return &cluster, err
-}
-
 // GetName returns this cluster's name.
 func (t *TestCluster) GetName() string {
-	cluster, err := t.ContainerCluster()
-	if err != nil {
-		return fmt.Sprintf("[error:%v]", err)
-	}
-	return cluster.GetName()
+	return t.clusterName
 }
 
 // GetGVisorRuntimeLabelMap returns the gVisor runtime key-value pair used
@@ -253,36 +272,95 @@ func (t *TestCluster) deleteNamespace(ctx context.Context, namespaceName string)
 	return ctx.Err()
 }
 
-// ListNodes is a helper method to list nodes in a cluster.
-func (t *TestCluster) ListNodes(ctx context.Context) (*v13.NodeList, error) {
-	return t.client.CoreV1().Nodes().List(ctx, v1.ListOptions{})
+// getNodePool returns the NodePool of the given type.
+// If nodepools have not been initialized yet, this method will initialize
+// them.
+func (t *TestCluster) getNodePool(ctx context.Context, nodepoolType NodePoolType) (*NodePool, error) {
+	t.nodepoolsMu.Lock()
+	defer t.nodepoolsMu.Unlock()
+	if t.nodepools == nil {
+		nodes, err := t.client.CoreV1().Nodes().List(ctx, v1.ListOptions{})
+		if err != nil {
+			return nil, fmt.Errorf("cannot list nodes: %w", err)
+		}
+		nodepools := make(map[NodePoolType]*NodePool, len(nodes.Items))
+		for _, node := range nodes.Items {
+			npType := NodePoolType(node.Labels[NodePoolTypeKey])
+			if npType == "" {
+				continue
+			}
+			npRuntime := RuntimeType(node.Labels[NodepoolRuntimeKey])
+			if npRuntime == "" {
+				continue
+			}
+			npArchitecture := CPUArchitecture(node.Labels[archKey])
+			if npArchitecture == "" {
+				continue
+			}
+			npAcceleratorType := AcceleratorType(node.Labels[NodepoolTPUAcceleratorSelectorKey])
+			if npAcceleratorType == "" {
+				// Attempt to derive it from instance type if possible.
+				if instanceType, hasInstanceType := node.Labels[NodepoolInstanceTypeKey]; hasInstanceType {
+					for accelType, machineType := range TPUAcceleratorMachineTypeMap {
+						if machineType == instanceType {
+							npAcceleratorType = accelType
+							break
+						}
+					}
+				}
+			}
+			npNumAccelerators := 0
+			if countStr, hasCount := node.Labels[NodepoolNumAcceleratorsKey]; hasCount {
+				if npNumAccelerators, err = strconv.Atoi(countStr); err != nil {
+					return nil, fmt.Errorf("cannot parse accelerator count (%q) value %q as an integer: %w", NodepoolNumAcceleratorsKey, countStr, err)
+				}
+			}
+			npTPUTopology := node.Labels[NodepoolTPUTopologyKey]
+			existingNodepool, ok := nodepools[npType]
+			if !ok {
+				nodepools[npType] = &NodePool{
+					nodePooltype:    npType,
+					runtime:         npRuntime,
+					cpuArchitecture: npArchitecture,
+					acceleratorType: npAcceleratorType,
+					numAccelerators: npNumAccelerators,
+					tpuTopology:     npTPUTopology,
+				}
+				continue
+			}
+			if existingNodepool.runtime != npRuntime {
+				return nil, fmt.Errorf("nodes in nodepool %q have conflicting runtimes: %v vs %v", npType, existingNodepool.runtime, npRuntime)
+			}
+			if existingNodepool.cpuArchitecture != npArchitecture {
+				return nil, fmt.Errorf("nodes in nodepool %q have conflicting architectures: %v vs %v", npType, existingNodepool.cpuArchitecture, npArchitecture)
+			}
+			if existingNodepool.acceleratorType != npAcceleratorType {
+				return nil, fmt.Errorf("nodes in nodepool %q have conflicting accelerator types: %v vs %v", npType, existingNodepool.acceleratorType, npAcceleratorType)
+			}
+			if existingNodepool.numAccelerators != npNumAccelerators {
+				return nil, fmt.Errorf("nodes in nodepool %q have conflicting accelerator counts: %v vs %v", npType, existingNodepool.numAccelerators, npNumAccelerators)
+			}
+			if existingNodepool.tpuTopology != npTPUTopology {
+				return nil, fmt.Errorf("nodes in nodepool %q have conflicting TPU topologies: %v vs %v", npType, existingNodepool.tpuTopology, npTPUTopology)
+			}
+		}
+		t.nodepools = nodepools
+	}
+	np, ok := t.nodepools[nodepoolType]
+	if !ok {
+		return nil, fmt.Errorf("cluster %q contains no %q nodepool", t.GetName(), nodepoolType)
+	}
+	return np, nil
 }
 
 // HasGVisorTestRuntime returns whether the test nodes in this cluster
 // use the gVisor runtime.
 func (t *TestCluster) HasGVisorTestRuntime(ctx context.Context) (bool, error) {
-	nodes, err := t.ListNodes(ctx)
+	testNodePool, err := t.getNodePool(ctx, TestRuntimeNodepoolName)
 	if err != nil {
-		return false, fmt.Errorf("cannot list nodes: %w", err)
+		return false, err
 	}
-	var foundRuntime RuntimeType
-	for _, n := range nodes.Items {
-		if n.Labels[NodePoolTypeKey] != TestRuntimeNodepoolName {
-			continue
-		}
-		nodeRuntime := RuntimeType(n.Labels[NodepoolRuntimeKey])
-		if nodeRuntime == "" {
-			return false, fmt.Errorf("node %q has no runtime label", n.GetName())
-		}
-		if foundRuntime == "" {
-			foundRuntime = nodeRuntime
-			continue
-		}
-		if nodeRuntime != foundRuntime {
-			return false, fmt.Errorf("found conflicting runtimes in the same cluster: %q vs %q", foundRuntime, nodeRuntime)
-		}
-	}
-	return foundRuntime == RuntimeTypeGVisor || foundRuntime == RuntimeTypeGVisorNvidia, nil
+	return testNodePool.runtime == RuntimeTypeGVisor || testNodePool.runtime == RuntimeTypeGVisorNvidia, nil
 }
 
 // CreatePod is a helper to create a pod.
@@ -399,37 +477,37 @@ func (t *TestCluster) doWaitForPod(ctx context.Context, pod *v13.Pod, phase v13.
 	}
 }
 
-// RuntimeTestNodepoolIsARM returns true if the runtime undertest nodepool is an ARM nodepool.
-func (t *TestCluster) RuntimeTestNodepoolIsARM() bool {
-	np, err := t.getNodePoolByName(TestRuntimeNodepoolName)
+// RuntimeTestNodepoolArchitecture returns the CPU architecture of the test nodepool.
+func (t *TestCluster) RuntimeTestNodepoolArchitecture(ctx context.Context) (CPUArchitecture, error) {
+	np, err := t.getNodePool(ctx, TestRuntimeNodepoolName)
 	if err != nil {
-		return false
+		return "", err
 	}
-	return strings.HasPrefix(np.GetConfig().GetMachineType(), "t2a")
+	return np.cpuArchitecture, nil
 }
 
 // configureDaemonSetForNodepool configures the DaemonSet to run on a given nodepool.
-func (t *TestCluster) configureDaemonSetForNodepool(ds *appsv1.DaemonSet, nodepoolName string) error {
-	np, err := t.getNodePoolByName(nodepoolName)
+func (t *TestCluster) configureDaemonSetForNodepool(ctx context.Context, ds *appsv1.DaemonSet, nodepoolType NodePoolType) error {
+	np, err := t.getNodePool(ctx, nodepoolType)
 	if err != nil {
 		return err
 	}
 	if ds.Labels == nil {
 		ds.Labels = make(map[string]string)
 	}
-	return t.applyCommonPodConfigurations(np, &ds.Spec.Template.Spec)
+	return t.applyCommonPodConfigurations(ctx, np, &ds.Spec.Template.Spec)
 }
 
 // configurePodForNodepool configures the pod to run on a given nodepool.
-func (t *TestCluster) configurePodForNodepool(pod *v13.Pod, nodepoolName string) (*v13.Pod, error) {
-	np, err := t.getNodePoolByName(nodepoolName)
+func (t *TestCluster) configurePodForNodepool(ctx context.Context, pod *v13.Pod, nodepoolType NodePoolType) (*v13.Pod, error) {
+	np, err := t.getNodePool(ctx, nodepoolType)
 	if err != nil {
 		return nil, err
 	}
 	if pod.Labels == nil {
 		pod.Labels = make(map[string]string)
 	}
-	if err := t.applyCommonPodConfigurations(np, &pod.Spec); err != nil {
+	if err := t.applyCommonPodConfigurations(ctx, np, &pod.Spec); err != nil {
 		return nil, err
 	}
 	return pod, nil
@@ -437,60 +515,40 @@ func (t *TestCluster) configurePodForNodepool(pod *v13.Pod, nodepoolName string)
 
 // ConfigureDaemonSetForRuntimeTestNodepool configures the DaemonSet to run
 // on the test runtime.
-func (t *TestCluster) ConfigureDaemonSetForRuntimeTestNodepool(ds *appsv1.DaemonSet) error {
-	return t.configureDaemonSetForNodepool(ds, TestRuntimeNodepoolName)
+func (t *TestCluster) ConfigureDaemonSetForRuntimeTestNodepool(ctx context.Context, ds *appsv1.DaemonSet) error {
+	return t.configureDaemonSetForNodepool(ctx, ds, TestRuntimeNodepoolName)
 }
 
 // ConfigurePodForRuntimeTestNodepool configures the pod to run on the test runtime.
-func (t *TestCluster) ConfigurePodForRuntimeTestNodepool(pod *v13.Pod) (*v13.Pod, error) {
-	return t.configurePodForNodepool(pod, TestRuntimeNodepoolName)
+func (t *TestCluster) ConfigurePodForRuntimeTestNodepool(ctx context.Context, pod *v13.Pod) (*v13.Pod, error) {
+	return t.configurePodForNodepool(ctx, pod, TestRuntimeNodepoolName)
 }
 
 // ConfigurePodForClientNodepool configures the pod to run on the client
 // nodepool.
-func (t *TestCluster) ConfigurePodForClientNodepool(pod *v13.Pod) (*v13.Pod, error) {
-	return t.configurePodForNodepool(pod, ClientNodepoolName)
+func (t *TestCluster) ConfigurePodForClientNodepool(ctx context.Context, pod *v13.Pod) (*v13.Pod, error) {
+	return t.configurePodForNodepool(ctx, pod, ClientNodepoolName)
 }
 
 // ConfigurePodForTertiaryNodepool configures the pod to run on the tertiary
 // nodepool.
-func (t *TestCluster) ConfigurePodForTertiaryNodepool(pod *v13.Pod) (*v13.Pod, error) {
-	return t.configurePodForNodepool(pod, TertiaryNodepoolName)
+func (t *TestCluster) ConfigurePodForTertiaryNodepool(ctx context.Context, pod *v13.Pod) (*v13.Pod, error) {
+	return t.configurePodForNodepool(ctx, pod, TertiaryNodepoolName)
 }
 
-func (t *TestCluster) getNodePoolByName(name string) (*cspb.NodePool, error) {
-	cluster, err := t.ContainerCluster()
-	if err != nil {
-		return nil, err
-	}
-	for _, np := range cluster.GetNodePools() {
-		if np.GetName() == name {
-			return np, nil
-		}
-	}
-	return nil, fmt.Errorf("failed to find nodepool %q: %+v", name, cluster.GetNodePools())
-}
-
-func (t *TestCluster) applyCommonPodConfigurations(np *cspb.NodePool, podSpec *v13.PodSpec) error {
-	// Apply GKE Sandbox configurations if the nodepool is a GKE Sandbox nodepool.
+func (t *TestCluster) applyCommonPodConfigurations(ctx context.Context, np *NodePool, podSpec *v13.PodSpec) error {
 	if podSpec.NodeSelector == nil {
 		podSpec.NodeSelector = make(map[string]string)
 	}
-
-	np.GetConfig().GetLabels()[NodePoolTypeKey] = np.GetName()
-
 	// Force the pod to run on this nodepool.
-	podSpec.NodeSelector[NodePoolSelectorKey] = np.GetName()
+	podSpec.NodeSelector[NodePoolTypeKey] = string(np.nodePooltype)
 
 	// Figure out which runtime to use for this pod, either by flag override or
 	// autodetection based on the nodepool configuration.
-	var applyRuntime = RuntimeTypeUnsandboxed
-	if np.GetName() == TestRuntimeNodepoolName && t.testNodepoolRuntimeOverride != "" {
+	var applyRuntime = np.runtime
+	if np.nodePooltype == TestRuntimeNodepoolName && t.testNodepoolRuntimeOverride != "" {
 		applyRuntime = t.testNodepoolRuntimeOverride
-	} else if nodePoolRuntime, ok := np.GetConfig().GetLabels()[NodepoolRuntimeKey]; ok {
-		applyRuntime = RuntimeType(nodePoolRuntime)
 	}
-
 	// Apply the runtime we've chosen, whether by override or autodetection.
 	applyRuntime.ApplyPodSpec(podSpec)
 
@@ -498,39 +556,25 @@ func (t *TestCluster) applyCommonPodConfigurations(np *cspb.NodePool, podSpec *v
 	// selector option.
 	// This doesn't really constrain the pod further, but allows
 	// this number to be carried over when setting pod resources.
-	if len(np.GetConfig().GetAccelerators()) > 0 {
-		totalAccels := 0
-		for _, accelCfg := range np.GetConfig().GetAccelerators() {
-			totalAccels += int(accelCfg.GetAcceleratorCount())
-		}
-		if accelCount, ok := np.GetConfig().GetLabels()[NodepoolNumAcceleratorsKey]; !ok || accelCount != strconv.Itoa(totalAccels) {
-			return fmt.Errorf("unexpected %s=%q label on nodepool with %d total accelerators", NodepoolNumAcceleratorsKey, accelCount, totalAccels)
-		}
-		podSpec.NodeSelector[NodepoolNumAcceleratorsKey] = strconv.Itoa(totalAccels)
-	} else {
-		for accelType, machineType := range TPUAcceleratorMachineTypeMap {
-			if machineType == np.GetConfig().GetMachineType() {
-				topology, ok := np.GetConfig().GetLabels()[NodepoolTPUTopologyKey]
-				if !ok {
-					return fmt.Errorf("unexpected %s=%q label on nodepool with no accelerators", NodepoolTPUTopologyKey, topology)
-				}
-				podSpec.NodeSelector[NodepoolTPUAcceleratorSelectorKey] = string(accelType)
-				podSpec.NodeSelector[NodepoolTPUTopologySelectorKey] = np.GetConfig().GetLabels()[NodepoolTPUTopologyKey]
-			}
-		}
+	if np.numAccelerators > 0 {
+		podSpec.NodeSelector[NodepoolNumAcceleratorsKey] = strconv.Itoa(np.numAccelerators)
+	}
+	if np.acceleratorType != "" {
+		podSpec.NodeSelector[NodepoolTPUAcceleratorSelectorKey] = string(np.acceleratorType)
+	}
+	if np.tpuTopology != "" {
+		podSpec.NodeSelector[NodepoolTPUTopologySelectorKey] = np.tpuTopology
 	}
 
 	// If the nodepool is an ARM nodepool, apply ARM tolerations.
-	for key, val := range np.GetConfig().GetLabels() {
-		if key == archKey && val == armValue {
-			podSpec.NodeSelector[archKey] = armValue
-			podSpec.Tolerations = append(podSpec.Tolerations, v13.Toleration{
-				Key:      archKey,
-				Value:    armValue,
-				Operator: v13.TolerationOpEqual,
-				Effect:   v13.TaintEffectNoSchedule,
-			})
-		}
+	if np.cpuArchitecture == CPUArchitectureARM {
+		podSpec.NodeSelector[archKey] = string(CPUArchitectureARM)
+		podSpec.Tolerations = append(podSpec.Tolerations, v13.Toleration{
+			Key:      archKey,
+			Value:    string(CPUArchitectureARM),
+			Operator: v13.TolerationOpEqual,
+			Effect:   v13.TaintEffectNoSchedule,
+		})
 	}
 	return nil
 }

test/kubernetes/tests/hello.go

@@ -36,7 +36,7 @@ func RunHello(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContext
 		t.Fatalf("Failed to resolve image: %v", err)
 	}
 	pod := ns.NewAlpinePod(fmt.Sprintf("hello-%d", time.Now().UnixNano()), image, []string{"/bin/sh", "-c", "echo hello"})
-	pod, err = cluster.ConfigurePodForRuntimeTestNodepool(pod)
+	pod, err = cluster.ConfigurePodForRuntimeTestNodepool(ctx, pod)
 	if err != nil {
 		t.Fatalf("Failed to set pod on cluster %q: %v", cluster.GetName(), err)
 	}