Probe and wait for GPUs to be available before Kubernetes GPU benchmarks.

At least on GKE, GPU resources may not be immediately available after the
cluster is created; it can take a few seconds before the nodes advertise
their availability of GPU resources and the control plane to notice this.

This change adds a probe loop at the beginning of GPU tests to wait for a
GPU-requesting pod to function before proceeding with further benchmarking.

This also renames `MaybeSetContainerResources` to `SetContainerResources`.
`MaybeSetContainerResources` used to rely on flags to determine whether to
change the container's resource requests. This is no longer the case, so
there is no more "maybe" about the function's behavior.

PiperOrigin-RevId: 707307396

test/kubernetes/benchmarks/abslbuild.go

@@ -108,7 +108,7 @@ func BuildABSL(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContex
 				t.Fatalf("Failed to set pod for test runtime: %v", err)
 			}
 
-			pod, err = testcluster.MaybeSetContainerResources(pod, name, testcluster.ContainerResourcesRequest{})
+			pod, err = testcluster.SetContainerResources(pod, "", testcluster.ContainerResourcesRequest{})
 			if err != nil {
 				t.Fatalf("Failed to set container resources: %v", err)
 			}

test/kubernetes/benchmarks/ffmpeg.go

@@ -111,7 +111,7 @@ func RunFFMPEG(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContex
 			if err != nil {
 				t.Fatalf("Failed to configure pod for runtime: %v", err)
 			}
-			p, err = testcluster.MaybeSetContainerResources(p, ffmpegContainerName, testcluster.ContainerResourcesRequest{})
+			p, err = testcluster.SetContainerResources(p, ffmpegContainerName, testcluster.ContainerResourcesRequest{})
 			if err != nil {
 				t.Fatalf("Failed to set container resources: %v", err)
 			}

test/kubernetes/benchmarks/grpc.go

@@ -113,7 +113,7 @@ func BuildGRPC(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContex
 				t.Fatalf("Failed to set pod for test runtime: %v", err)
 			}
 
-			pod, err = testcluster.MaybeSetContainerResources(pod, name, testcluster.ContainerResourcesRequest{})
+			pod, err = testcluster.SetContainerResources(pod, "", testcluster.ContainerResourcesRequest{})
 			if err != nil {
 				t.Fatalf("Failed to set container resources: %v", err)
 			}

test/kubernetes/benchmarks/gsutil.go

@@ -133,7 +133,7 @@ func RunGSUtil(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesContex
 					if err != nil {
 						t.Fatalf("Failed to configure pod for runtime: %v", err)
 					}
-					p, err = testcluster.MaybeSetContainerResources(p, name, testcluster.ContainerResourcesRequest{})
+					p, err = testcluster.SetContainerResources(p, "", testcluster.ContainerResourcesRequest{})
 					if err != nil {
 						t.Fatalf("Failed to set container resources: %v", err)
 					}

test/kubernetes/benchmarks/nginx.go

@@ -135,7 +135,7 @@ func BenchmarkNginx(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesC
 			if err != nil {
 				t.Fatalf("Failed to configure pod for runtime nodepool: %v", err)
 			}
-			server, err = testcluster.MaybeSetContainerResources(server, name, testcluster.ContainerResourcesRequest{})
+			server, err = testcluster.SetContainerResources(server, "", testcluster.ContainerResourcesRequest{})
 			if err != nil {
 				t.Fatalf("Failed to set container resources: %v", err)
 			}

test/kubernetes/benchmarks/ollama.go

@@ -235,6 +235,11 @@ func BenchmarkOllama(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kubernetes
 		t.Fatalf("cannot reset namespace: %v", err)
 	}
 	defer benchmarkNS.Cleanup(ctx)
+	reqWaitCtx, reqWaitCancel := context.WithTimeout(ctx, 5*time.Minute)
+	defer reqWaitCancel()
+	if err := benchmarkNS.WaitForResources(reqWaitCtx, testcluster.ContainerResourcesRequest{GPU: true}); err != nil {
+		t.Fatalf("failed to wait for resources: %v", err)
+	}
 	endProfiling, err := profiling.MaybeSetup(ctx, t, k8sCtx, cluster, benchmarkNS)
 	if err != nil {
 		t.Fatalf("Failed to setup profiling: %v", err)
@@ -254,7 +259,7 @@ func BenchmarkOllama(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kubernetes
 	if err != nil {
 		t.Fatalf("Failed to configure pod for runtime nodepool: %v", err)
 	}
-	ollamaPod, err = testcluster.MaybeSetContainerResources(ollamaPod, ollamaPod.ObjectMeta.Name, testcluster.ContainerResourcesRequest{GPU: true})
+	ollamaPod, err = testcluster.SetContainerResources(ollamaPod, "", testcluster.ContainerResourcesRequest{GPU: true})
 	if err != nil {
 		t.Fatalf("Failed to set container resources: %v", err)
 	}

test/kubernetes/benchmarks/postgresql.go

@@ -85,9 +85,9 @@ func BenchmarkPostgresPGBench(ctx context.Context, t *testing.T, k8sCtx k8sctx.K
 		t.Fatalf("ConfigurePodForRuntimeTestNodepool on cluster %q: %v", cluster.GetName(), err)
 	}
 
-	server, err = testcluster.MaybeSetContainerResources(server, server.Spec.Containers[0].Name, testcluster.ContainerResourcesRequest{})
+	server, err = testcluster.SetContainerResources(server, "", testcluster.ContainerResourcesRequest{})
 	if err != nil {
-		t.Fatalf("MaybeSetContainerResources on cluster %q: %v", cluster.GetName(), err)
+		t.Fatalf("SetContainerResources on cluster %q: %v", cluster.GetName(), err)
 	}
 
 	server, err = cluster.CreatePod(ctx, server)

test/kubernetes/benchmarks/pytorch.go

@@ -215,6 +215,12 @@ func doPytorchRun(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesCon
 		t.Fatalf("Failed to reset namespace: %v", err)
 	}
 	defer benchmarkNS.Cleanup(ctx)
+	reqWaitCtx, reqWaitCancel := context.WithTimeout(ctx, 5*time.Minute)
+	defer reqWaitCancel()
+	if err := benchmarkNS.WaitForResources(reqWaitCtx, testcluster.ContainerResourcesRequest{GPU: true}); err != nil {
+		t.Fatalf("failed to wait for resources: %v", err)
+	}
+
 	endProfiling, err := profiling.MaybeSetup(ctx, t, k8sCtx, cluster, benchmarkNS)
 	if err != nil {
 		t.Fatalf("Failed to setup profiling: %v", err)
@@ -235,7 +241,7 @@ func doPytorchRun(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesCon
 		t.Fatalf("Failed to configure pod for test-nodepool: %v", err)
 	}
 
-	pod, err = testcluster.MaybeSetContainerResources(pod, pod.Name, testcluster.ContainerResourcesRequest{GPU: true})
+	pod, err = testcluster.SetContainerResources(pod, "", testcluster.ContainerResourcesRequest{GPU: true})
 	if err != nil {
 		t.Fatalf("Failed to set container resources: %v", err)
 	}

test/kubernetes/benchmarks/redis.go

@@ -136,9 +136,9 @@ func BenchmarkRedis(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesC
 				t.Fatalf("ConfigurePodForRuntimeTestNodepool on cluster %q: %v", cluster.GetName(), err)
 			}
 
-			server, err = testcluster.MaybeSetContainerResources(server, server.Spec.Containers[0].Name, testcluster.ContainerResourcesRequest{})
+			server, err = testcluster.SetContainerResources(server, "", testcluster.ContainerResourcesRequest{})
 			if err != nil {
-				t.Fatalf("MaybeSetContainerResources on cluster %q: %v", cluster.GetName(), err)
+				t.Fatalf("SetContainerResources on cluster %q: %v", cluster.GetName(), err)
 			}
 
 			server, err = cluster.CreatePod(ctx, server)

test/kubernetes/benchmarks/rubydev.go

@@ -117,7 +117,7 @@ func RunRubyDev(ctx context.Context, t *testing.T, k8sCtx k8sctx.KubernetesConte
 				t.Fatalf("failed to configure pod for test runtime node: %v", err)
 			}
 
-			pod, err = testcluster.MaybeSetContainerResources(pod, builderContainerName, testcluster.ContainerResourcesRequest{})
+			pod, err = testcluster.SetContainerResources(pod, "", testcluster.ContainerResourcesRequest{})
 			if err != nil {
 				t.Fatalf("failed to set container resources: %v", err)
 			}

test/kubernetes/benchmarks/stablediffusion.go

@@ -71,7 +71,7 @@ func (r *kubernetesPodRunner) Run(ctx context.Context, image string, argv []stri
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to configure pod: %v", err)
 	}
-	stableDiffusionXLPod, err = testcluster.MaybeSetContainerResources(stableDiffusionXLPod, stableDiffusionXLPod.ObjectMeta.Name, testcluster.ContainerResourcesRequest{GPU: true})
+	stableDiffusionXLPod, err = testcluster.SetContainerResources(stableDiffusionXLPod, "", testcluster.ContainerResourcesRequest{GPU: true})
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to set container resources: %v", err)
 	}
@@ -113,6 +113,11 @@ func RunStableDiffusionXL(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kuber
 		t.Fatalf("cannot reset namespace: %v", err)
 	}
 	defer benchmarkNS.Cleanup(ctx)
+	reqWaitCtx, reqWaitCancel := context.WithTimeout(ctx, 5*time.Minute)
+	defer reqWaitCancel()
+	if err := benchmarkNS.WaitForResources(reqWaitCtx, testcluster.ContainerResourcesRequest{GPU: true}); err != nil {
+		t.Fatalf("failed to wait for resources: %v", err)
+	}
 	endProfiling, err := profiling.MaybeSetup(ctx, t, k8sCtx, cluster, benchmarkNS)
 	if err != nil {
 		t.Fatalf("Failed to setup profiling: %v", err)

test/kubernetes/benchmarks/tensorflow.go

@@ -99,7 +99,7 @@ func RunTensorflowOnCPU(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kuberne
 				t.Fatalf("Failed to set pod for test runtime: %v", err)
 			}
 
-			pod, err = testcluster.MaybeSetContainerResources(pod, name, testcluster.ContainerResourcesRequest{})
+			pod, err = testcluster.SetContainerResources(pod, "", testcluster.ContainerResourcesRequest{})
 			if err != nil {
 				t.Fatalf("Failed to set container resources: %v", err)
 			}

test/kubernetes/benchmarks/wordpress.go

@@ -114,7 +114,7 @@ func BenchmarkWordpress(ctx context.Context, t *testing.T, k8sCtx k8sctx.Kuberne
 	if err != nil {
 		t.Fatalf("Failed to configure pod for runtime nodepool: %v", err)
 	}
-	server, err = testcluster.MaybeSetContainerResources(server, name, testcluster.ContainerResourcesRequest{})
+	server, err = testcluster.SetContainerResources(server, "", testcluster.ContainerResourcesRequest{})
 	if err != nil {
 		t.Fatalf("Failed to set container resources: %v", err)
 	}

test/kubernetes/testcluster/objects.go

@@ -19,6 +19,7 @@ import (
 	"errors"
 	"fmt"
 	"reflect"
+	"time"
 
 	cspb "google.golang.org/genproto/googleapis/container/v1"
 	"google.golang.org/protobuf/proto"
@@ -142,9 +143,16 @@ type ContainerResourcesRequest struct {
 	GPU             bool
 }
 
-// MaybeSetContainerResources sets container resources if flags are given. Sets both the resource
-// limits and requests as container runtimes honor them differently.
-func MaybeSetContainerResources(pod *v13.Pod, containerName string, requests ContainerResourcesRequest) (*v13.Pod, error) {
+// String returns a string representation of the `ContainerResourcesRequest`.
+func (crr ContainerResourcesRequest) String() string {
+	return fmt.Sprintf("cpu=%q memory=%q gpu=%v", crr.CPUResources, crr.MemoryResources, crr.GPU)
+}
+
+// SetContainerResources sets container resources.
+// Sets both the resource limits and requests as container runtimes honor
+// them differently.
+// `containerName` is optional if the pod has exactly one container.
+func SetContainerResources(pod *v13.Pod, containerName string, requests ContainerResourcesRequest) (*v13.Pod, error) {
 	resourceList := v13.ResourceList{}
 	if requests.CPUResources != "" {
 		resourceList[v13.ResourceCPU] = resource.MustParse(requests.CPUResources)
@@ -166,13 +174,102 @@ func MaybeSetContainerResources(pod *v13.Pod, containerName string, requests Con
 		Requests: resourceList,
 	}
 
-	for i := range pod.Spec.Containers {
-		if pod.Spec.Containers[i].Name == containerName {
-			pod.Spec.Containers[i].Resources = requirements
-			return pod, nil
+	var containerToChange *v13.Container
+	if containerName == "" {
+		switch len(pod.Spec.Containers) {
+		case 0:
+			return nil, fmt.Errorf("no containers found in pod")
+		case 1:
+			containerToChange = &pod.Spec.Containers[0]
+		default:
+			return nil, fmt.Errorf("multiple containers found in pod %v, please specify a container name", pod)
+		}
+	} else {
+		for i := range pod.Spec.Containers {
+			if pod.Spec.Containers[i].Name == containerName {
+				containerToChange = &pod.Spec.Containers[i]
+			}
+		}
+	}
+	if containerToChange == nil {
+		return nil, fmt.Errorf("container %q not found", containerName)
+	}
+	containerToChange.Resources = requirements
+	return pod, nil
+}
+
+// ProbeResources verifies that a pod requesting the given resources can be
+// scheduled.
+func (n *Namespace) ProbeResources(ctx context.Context, requests ContainerResourcesRequest) error {
+	// Configure the pod.
+	probe := n.NewAlpinePod("resource-probe", "alpine", []string{"/bin/true"})
+	probe, err := n.testCluster.ConfigurePodForRuntimeTestNodepool(ctx, probe)
+	if err != nil {
+		return fmt.Errorf("failed to configure pod for runtime test nodepool: %w", err)
+	}
+	if probe, err = SetContainerResources(probe, "", requests); err != nil {
+		return fmt.Errorf("failed to set container resources: %w", err)
+	}
+	resources := probe.Spec.Containers[0].Resources
+
+	// If a pod already exists (e.g. from a past probe), best-effort attempt to delete it:
+	deleteCtx, deleteCancel := context.WithTimeout(ctx, 20*time.Second)
+	_ = n.testCluster.DeletePod(deleteCtx, probe)
+	deleteCancel()
+
+	// Create the pod.
+	if _, err := n.testCluster.CreatePod(ctx, probe); err != nil {
+		return fmt.Errorf("failed to create probe pod with resources %v: %w", resources, err)
+	}
+
+	// Wait up to (ctx - 30 seconds) for the pod to run.
+	waitCtx := ctx
+	if ctxDeadline, hasDeadline := waitCtx.Deadline(); hasDeadline {
+		var waitCancel context.CancelFunc
+		waitCtx, waitCancel = context.WithDeadline(ctx, ctxDeadline.Add(-30*time.Second))
+		defer waitCancel()
+	}
+	if err := n.testCluster.WaitForPodRunning(waitCtx, probe); err != nil {
+		// Best-effort time-limited deletion.
+		deleteCtx, deleteCancel := context.WithTimeout(ctx, 30*time.Second)
+		_ = n.testCluster.DeletePod(deleteCtx, probe)
+		deleteCancel()
+
+		return fmt.Errorf("probe pod with resources %v did not run: %w", resources, err)
+	}
+
+	// Delete the pod.
+	if err := n.testCluster.DeletePod(ctx, probe); err != nil {
+		return fmt.Errorf("failed to delete probe pod: %w", err)
+	}
+	return nil
+}
+
+// WaitForResources checks that a pod requesting the given resources can be
+// scheduled. If they cannot, it will loop until the given context expire or
+// the resources become available.
+func (n *Namespace) WaitForResources(ctx context.Context, requests ContainerResourcesRequest) error {
+	var lastErr error
+	for {
+		probeCtx, probeCancel := context.WithTimeout(ctx, 90*time.Second)
+		err := n.ProbeResources(probeCtx, requests)
+		probeCtxErr := probeCtx.Err()
+		probeCancel()
+		if err == nil {
+			return nil
+		}
+		select {
+		case <-ctx.Done():
+			if lastErr != nil {
+				return fmt.Errorf("context expired while waiting for resources %v to be available; last error: %w", requests, lastErr)
+			}
+			return fmt.Errorf("context expired before ever checking that resources %v are available: %w", requests, err)
+		case <-time.After(pollInterval):
+			if lastErr == nil || !errors.Is(probeCtxErr, context.DeadlineExceeded) {
+				lastErr = err
+			}
 		}
 	}
-	return nil, fmt.Errorf("container %q not found", containerName)
 }
 
 // RuntimeType is a supported runtime for the test nodepool.