From 7321e426cef0c7ded0263394a588fd5ec8a2424e Mon Sep 17 00:00:00 2001 From: Etienne Perot Date: Tue, 17 Dec 2024 17:08:07 -0800 Subject: [PATCH] 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 | 2 +- test/kubernetes/benchmarks/ffmpeg.go | 2 +- test/kubernetes/benchmarks/grpc.go | 2 +- test/kubernetes/benchmarks/gsutil.go | 2 +- test/kubernetes/benchmarks/nginx.go | 2 +- test/kubernetes/benchmarks/ollama.go | 7 +- test/kubernetes/benchmarks/postgresql.go | 4 +- test/kubernetes/benchmarks/pytorch.go | 8 +- test/kubernetes/benchmarks/redis.go | 4 +- test/kubernetes/benchmarks/rubydev.go | 2 +- test/kubernetes/benchmarks/stablediffusion.go | 7 +- test/kubernetes/benchmarks/tensorflow.go | 2 +- test/kubernetes/benchmarks/wordpress.go | 2 +- test/kubernetes/testcluster/objects.go | 113 ++++++++++++++++-- 14 files changed, 136 insertions(+), 23 deletions(-) diff --git a/test/kubernetes/benchmarks/abslbuild.go b/test/kubernetes/benchmarks/abslbuild.go index ee5efabb6..aeaf9a35b 100644 --- a/test/kubernetes/benchmarks/abslbuild.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/ffmpeg.go b/test/kubernetes/benchmarks/ffmpeg.go index 0f79a9360..2aa6f2a0d 100644 --- a/test/kubernetes/benchmarks/ffmpeg.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/grpc.go b/test/kubernetes/benchmarks/grpc.go index 816645791..165ed167d 100644 --- a/test/kubernetes/benchmarks/grpc.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/gsutil.go b/test/kubernetes/benchmarks/gsutil.go index 78af16460..26fdfa738 100644 --- a/test/kubernetes/benchmarks/gsutil.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/nginx.go b/test/kubernetes/benchmarks/nginx.go index 9bf12a22f..3772edde7 100644 --- a/test/kubernetes/benchmarks/nginx.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/ollama.go b/test/kubernetes/benchmarks/ollama.go index 39398b1e7..5f6e447c8 100644 --- a/test/kubernetes/benchmarks/ollama.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/postgresql.go b/test/kubernetes/benchmarks/postgresql.go index d6b127f2e..d04e35352 100644 --- a/test/kubernetes/benchmarks/postgresql.go +++ b/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) diff --git a/test/kubernetes/benchmarks/pytorch.go b/test/kubernetes/benchmarks/pytorch.go index 2ddc55709..4d382fcac 100644 --- a/test/kubernetes/benchmarks/pytorch.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/redis.go b/test/kubernetes/benchmarks/redis.go index 73ec1e937..efd06a113 100644 --- a/test/kubernetes/benchmarks/redis.go +++ b/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) diff --git a/test/kubernetes/benchmarks/rubydev.go b/test/kubernetes/benchmarks/rubydev.go index c0de48536..2261aff88 100644 --- a/test/kubernetes/benchmarks/rubydev.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/stablediffusion.go b/test/kubernetes/benchmarks/stablediffusion.go index a6fdc6ab5..2e7d6c69e 100644 --- a/test/kubernetes/benchmarks/stablediffusion.go +++ b/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) diff --git a/test/kubernetes/benchmarks/tensorflow.go b/test/kubernetes/benchmarks/tensorflow.go index 57b232863..83570a675 100644 --- a/test/kubernetes/benchmarks/tensorflow.go +++ b/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) } diff --git a/test/kubernetes/benchmarks/wordpress.go b/test/kubernetes/benchmarks/wordpress.go index 6642e258a..b3eb2400c 100644 --- a/test/kubernetes/benchmarks/wordpress.go +++ b/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) } diff --git a/test/kubernetes/testcluster/objects.go b/test/kubernetes/testcluster/objects.go index 07a4e15f4..82956817b 100644 --- a/test/kubernetes/testcluster/objects.go +++ b/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.