From e1cb1a70b453029ec642fc46ff1d101eb21588fc Mon Sep 17 00:00:00 2001 From: Etienne Perot Date: Tue, 26 Nov 2024 13:42:40 -0800 Subject: [PATCH] `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 | 10 +- test/kubernetes/benchmarks/ffmpeg.go | 17 +- test/kubernetes/benchmarks/grpc.go | 17 +- test/kubernetes/benchmarks/gsutil.go | 17 +- .../benchmarks/httpbench/httpbench.go | 4 +- test/kubernetes/benchmarks/nginx.go | 22 +- test/kubernetes/benchmarks/ollama.go | 4 +- test/kubernetes/benchmarks/postgresql.go | 10 +- test/kubernetes/benchmarks/profiling/BUILD | 1 + .../benchmarks/profiling/profiling.go | 29 +- test/kubernetes/benchmarks/pytorch.go | 4 +- test/kubernetes/benchmarks/redis.go | 21 +- test/kubernetes/benchmarks/rubydev.go | 17 +- test/kubernetes/benchmarks/stablediffusion.go | 2 +- test/kubernetes/benchmarks/startup.go | 4 +- test/kubernetes/benchmarks/tensorflow.go | 17 +- test/kubernetes/benchmarks/wordpress.go | 8 +- test/kubernetes/testcluster/BUILD | 1 + test/kubernetes/testcluster/testcluster.go | 340 ++++++++++-------- test/kubernetes/tests/hello.go | 2 +- 20 files changed, 337 insertions(+), 210 deletions(-) diff --git a/test/kubernetes/benchmarks/abslbuild.go b/test/kubernetes/benchmarks/abslbuild.go index f82d939cf..ac3cbeb3f 100644 --- a/test/kubernetes/benchmarks/abslbuild.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/ffmpeg.go b/test/kubernetes/benchmarks/ffmpeg.go index 6980cc7b1..0f79a9360 100644 --- a/test/kubernetes/benchmarks/ffmpeg.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/grpc.go b/test/kubernetes/benchmarks/grpc.go index a1814e565..c68772022 100644 --- a/test/kubernetes/benchmarks/grpc.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/gsutil.go b/test/kubernetes/benchmarks/gsutil.go index 6a0425cfd..72b9c44f5 100644 --- a/test/kubernetes/benchmarks/gsutil.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/httpbench/httpbench.go b/test/kubernetes/benchmarks/httpbench/httpbench.go index 72a500d0f..757e4ef46 100644 --- a/test/kubernetes/benchmarks/httpbench/httpbench.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/nginx.go b/test/kubernetes/benchmarks/nginx.go index e3f10532b..31bc3af7c 100644 --- a/test/kubernetes/benchmarks/nginx.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/ollama.go b/test/kubernetes/benchmarks/ollama.go index 89c511683..60301b9dd 100644 --- a/test/kubernetes/benchmarks/ollama.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/postgresql.go b/test/kubernetes/benchmarks/postgresql.go index d24759268..7b3d0ec44 100644 --- a/test/kubernetes/benchmarks/postgresql.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/profiling/BUILD b/test/kubernetes/benchmarks/profiling/BUILD index c769860bc..048f436e2 100644 --- a/test/kubernetes/benchmarks/profiling/BUILD +++ b/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", diff --git a/test/kubernetes/benchmarks/profiling/profiling.go b/test/kubernetes/benchmarks/profiling/profiling.go index 3feb736a3..0a33e2b9e 100644 --- a/test/kubernetes/benchmarks/profiling/profiling.go +++ b/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"), diff --git a/test/kubernetes/benchmarks/pytorch.go b/test/kubernetes/benchmarks/pytorch.go index 696374537..c53f49e4b 100644 --- a/test/kubernetes/benchmarks/pytorch.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/redis.go b/test/kubernetes/benchmarks/redis.go index 03977b8a2..cfe0a8302 100644 --- a/test/kubernetes/benchmarks/redis.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/rubydev.go b/test/kubernetes/benchmarks/rubydev.go index 46261bdb6..c0de48536 100644 --- a/test/kubernetes/benchmarks/rubydev.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/stablediffusion.go b/test/kubernetes/benchmarks/stablediffusion.go index b377ab982..a87666434 100644 --- a/test/kubernetes/benchmarks/stablediffusion.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/startup.go b/test/kubernetes/benchmarks/startup.go index e0ef3bb06..a57ca9126 100644 --- a/test/kubernetes/benchmarks/startup.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/tensorflow.go b/test/kubernetes/benchmarks/tensorflow.go index 6fa08aa47..da8e4f869 100644 --- a/test/kubernetes/benchmarks/tensorflow.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/wordpress.go b/test/kubernetes/benchmarks/wordpress.go index daa5fdf8a..b515ffbce 100644 --- a/test/kubernetes/benchmarks/wordpress.go +++ b/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) } diff --git a/test/kubernetes/testcluster/BUILD b/test/kubernetes/testcluster/BUILD index b71824554..d1357804d 100644 --- a/test/kubernetes/testcluster/BUILD +++ b/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", diff --git a/test/kubernetes/testcluster/testcluster.go b/test/kubernetes/testcluster/testcluster.go index fde5f3f71..b8e0220ae 100644 --- a/test/kubernetes/testcluster/testcluster.go +++ b/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 } diff --git a/test/kubernetes/tests/hello.go b/test/kubernetes/tests/hello.go index c059bdfbd..324ba0e14 100644 --- a/test/kubernetes/tests/hello.go +++ b/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) }