fix(validator): wait for KAI deployments ready in gang-scheduling check (#1514)

Signed-off-by: Mark Chmarny <mark@chmarny.com>

Author: mchmarny

validators/conformance/gang_scheduling_check.go

@@ -109,13 +109,20 @@ func CheckGangScheduling(ctx *validators.Context) error {
 		return validators.Skip("KAI scheduler not found — cluster may use a different scheduler")
 	}
 
-	// 1. All KAI scheduler deployments available.
+	// 1. All KAI scheduler deployments available. Wait (bounded) for each to
+	// become Available rather than sampling instantaneously: the single-replica
+	// admission webhook can be transiently unavailable (rollout, restart, node
+	// pressure), and a point-in-time 0/1 read would flake the whole conformance
+	// phase. A genuinely-down deployment still fails after the bound.
 	var deploymentsSummary strings.Builder
 	for _, name := range kaiSchedulerDeployments {
-		deploy, err := getDeploymentIfAvailable(ctx, "kai-scheduler", name)
+		deploy, err := waitForDeploymentAvailable(ctx, "kai-scheduler", name, defaults.K8sPodReadyTimeout)
 		if err != nil {
-			return errors.Wrap(errors.ErrCodeNotFound,
-				fmt.Sprintf("KAI scheduler component %s check failed", name), err)
+			// Preserve the helper's code (NotFound for missing, Internal for
+			// not-available/API failure, Timeout for cancellation) instead of
+			// flattening every failure to NotFound.
+			return errors.PropagateOrWrap(err, errors.ErrCodeNotFound,
+				fmt.Sprintf("KAI scheduler component %s check failed", name))
 		}
 		expected := int32(1)
 		if deploy.Spec.Replicas != nil {

validators/conformance/helpers.go

@@ -21,6 +21,7 @@ import (
 	"log/slog"
 	"net/http"
 	"strings"
+	"time"
 
 	"github.com/NVIDIA/aicr/pkg/defaults"
 	"github.com/NVIDIA/aicr/pkg/errors"
@@ -155,6 +156,74 @@ func getDeploymentIfAvailable(ctx *validators.Context, namespace, name string) (
 	return deploy, nil
 }
 
+// waitForDeploymentAvailable polls until the named Deployment reports at least
+// one available replica, or the timeout elapses. Returns the last-observed
+// Deployment so callers can capture diagnostics.
+//
+// Use this instead of getDeploymentIfAvailable for components that can be
+// transiently unavailable — single-replica webhooks, rolling updates,
+// restarts. An instantaneous check fails closed on a momentary 0/N blip; the
+// bounded wait tolerates the blip but still fails (after the bound) when a
+// deployment is genuinely down. A not-yet-created deployment is treated as
+// "keep waiting" within the bound rather than an immediate NotFound.
+//
+//nolint:unparam // namespace is part of the general helper API (mirrors getDeploymentIfAvailable); the sole caller today happens to use one namespace.
+func waitForDeploymentAvailable(ctx *validators.Context, namespace, name string, timeout time.Duration) (*appsv1.Deployment, error) {
+	pollCtx, cancel := context.WithTimeout(ctx.Ctx, timeout)
+	defer cancel()
+
+	var last *appsv1.Deployment
+	err := wait.PollUntilContextCancel(pollCtx, defaults.PodPollInterval, true,
+		func(c context.Context) (bool, error) {
+			deploy, getErr := ctx.Clientset.AppsV1().Deployments(namespace).Get(c, name, metav1.GetOptions{})
+			if getErr != nil {
+				if k8serrors.IsNotFound(getErr) {
+					return false, nil // not created yet — keep waiting within the bound
+				}
+				return false, errors.Wrap(errors.ErrCodeInternal,
+					fmt.Sprintf("failed to get deployment %s/%s", namespace, name), getErr)
+			}
+			last = deploy
+			return deploy.Status.AvailableReplicas >= 1, nil
+		},
+	)
+	if err == nil {
+		return last, nil
+	}
+
+	// Caller cancellation is an external abort, not a readiness failure.
+	// pollCtx derives from ctx.Ctx, so pollCtx.Err() is also set when the parent
+	// is canceled — check the parent explicitly first and propagate it as a
+	// transient timeout rather than reporting the deployment as not-available.
+	if ctx.Ctx.Err() != nil {
+		return last, errors.Wrap(errors.ErrCodeTimeout,
+			fmt.Sprintf("waiting for deployment %s/%s canceled", namespace, name), ctx.Ctx.Err())
+	}
+
+	expected := int32(1)
+	var avail int32
+	if last != nil {
+		avail = last.Status.AvailableReplicas
+		if last.Spec.Replicas != nil {
+			expected = *last.Spec.Replicas
+		}
+	}
+	// Our own bound elapsed (parent still live): the deployment never became
+	// available in time. Surface the NotFound-shaped not-available message the
+	// caller wraps. A non-deadline error is a genuine API failure — propagate it.
+	if pollCtx.Err() != nil {
+		if last == nil {
+			return nil, errors.New(errors.ErrCodeNotFound,
+				fmt.Sprintf("deployment %s/%s not found after %s", namespace, name, timeout))
+		}
+		return last, errors.New(errors.ErrCodeInternal,
+			fmt.Sprintf("deployment %s/%s not available after %s: %d/%d replicas",
+				namespace, name, timeout, avail, expected))
+	}
+	return last, errors.PropagateOrWrap(err, errors.ErrCodeInternal,
+		fmt.Sprintf("failed waiting for deployment %s/%s", namespace, name))
+}
+
 // verifyDaemonSetReady checks that a DaemonSet has at least one ready pod.
 func verifyDaemonSetReady(ctx *validators.Context, namespace, name string) error {
 	_, err := getDaemonSetIfReady(ctx, namespace, name)

validators/conformance/helpers_readiness_test.go

@@ -0,0 +1,140 @@
+// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"context"
+	stderrors "errors"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/NVIDIA/aicr/pkg/errors"
+	"github.com/NVIDIA/aicr/validators"
+	appsv1 "k8s.io/api/apps/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	k8sfake "k8s.io/client-go/kubernetes/fake"
+	k8stesting "k8s.io/client-go/testing"
+)
+
+//nolint:unparam // name is a meaningful fixture input even though current cases all use "admission".
+func deployWithAvailable(name string, avail int32) *appsv1.Deployment {
+	r := int32(1)
+	return &appsv1.Deployment{
+		ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: "kai-scheduler"},
+		Spec:       appsv1.DeploymentSpec{Replicas: &r},
+		Status:     appsv1.DeploymentStatus{AvailableReplicas: avail},
+	}
+}
+
+// TestWaitForDeploymentAvailable_ImmediatelyReady returns without waiting when
+// the deployment is already Available.
+func TestWaitForDeploymentAvailable_ImmediatelyReady(t *testing.T) {
+	t.Parallel()
+	client := k8sfake.NewClientset(deployWithAvailable("admission", 1))
+	vctx := &validators.Context{Ctx: context.Background(), Clientset: client}
+
+	got, err := waitForDeploymentAvailable(vctx, "kai-scheduler", "admission", 5*time.Second)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if got == nil || got.Status.AvailableReplicas != 1 {
+		t.Fatalf("expected available deployment, got %+v", got)
+	}
+}
+
+// TestWaitForDeploymentAvailable_TransientBlip is the regression test for the
+// gang-scheduling flake: a deployment that reports 0/1 on the first sample and
+// becomes Available shortly after must pass, not fail closed.
+func TestWaitForDeploymentAvailable_TransientBlip(t *testing.T) {
+	t.Parallel()
+	client := k8sfake.NewClientset(deployWithAvailable("admission", 0))
+
+	var calls int32
+	client.PrependReactor("get", "deployments", func(k8stesting.Action) (bool, runtime.Object, error) {
+		n := atomic.AddInt32(&calls, 1)
+		avail := int32(0)
+		if n >= 2 { // available from the second poll onward
+			avail = 1
+		}
+		return true, deployWithAvailable("admission", avail), nil
+	})
+
+	vctx := &validators.Context{Ctx: context.Background(), Clientset: client}
+	got, err := waitForDeploymentAvailable(vctx, "kai-scheduler", "admission", 5*time.Second)
+	if err != nil {
+		t.Fatalf("expected success after transient 0/1, got error: %v", err)
+	}
+	if got == nil || got.Status.AvailableReplicas != 1 {
+		t.Fatalf("expected available deployment, got %+v", got)
+	}
+	if atomic.LoadInt32(&calls) < 2 {
+		t.Errorf("expected the helper to re-poll past the initial 0/1 (calls=%d)", calls)
+	}
+}
+
+// TestWaitForDeploymentAvailable_NeverReady fails closed (after the bound) when
+// the deployment stays unavailable.
+func TestWaitForDeploymentAvailable_NeverReady(t *testing.T) {
+	t.Parallel()
+	client := k8sfake.NewClientset(deployWithAvailable("admission", 0))
+	vctx := &validators.Context{Ctx: context.Background(), Clientset: client}
+
+	_, err := waitForDeploymentAvailable(vctx, "kai-scheduler", "admission", 50*time.Millisecond)
+	if err == nil {
+		t.Fatal("expected error for never-ready deployment")
+	}
+	if !stderrors.Is(err, errors.New(errors.ErrCodeInternal, "")) {
+		t.Errorf("expected ErrCodeInternal, got %v", err)
+	}
+}
+
+// TestWaitForDeploymentAvailable_ParentCanceled ensures an external abort
+// (parent context canceled) is propagated as a transient timeout, not
+// misclassified as a deployment readiness failure.
+func TestWaitForDeploymentAvailable_ParentCanceled(t *testing.T) {
+	t.Parallel()
+	client := k8sfake.NewClientset(deployWithAvailable("admission", 0))
+	parent, cancel := context.WithCancel(context.Background())
+	cancel() // external abort before the wait runs
+	vctx := &validators.Context{Ctx: parent, Clientset: client}
+
+	// Generous bound so any failure must come from the parent cancellation,
+	// not our own deadline.
+	_, err := waitForDeploymentAvailable(vctx, "kai-scheduler", "admission", 5*time.Second)
+	if err == nil {
+		t.Fatal("expected error on canceled parent context")
+	}
+	if !stderrors.Is(err, errors.New(errors.ErrCodeTimeout, "")) {
+		t.Errorf("expected ErrCodeTimeout for cancellation, got %v", err)
+	}
+}
+
+// TestWaitForDeploymentAvailable_NotFound surfaces a missing deployment as
+// NotFound after the bound (not a generic internal error).
+func TestWaitForDeploymentAvailable_NotFound(t *testing.T) {
+	t.Parallel()
+	client := k8sfake.NewClientset() // no deployments
+	vctx := &validators.Context{Ctx: context.Background(), Clientset: client}
+
+	_, err := waitForDeploymentAvailable(vctx, "kai-scheduler", "admission", 50*time.Millisecond)
+	if err == nil {
+		t.Fatal("expected error for missing deployment")
+	}
+	if !stderrors.Is(err, errors.New(errors.ErrCodeNotFound, "")) {
+		t.Errorf("expected ErrCodeNotFound, got %v", err)
+	}
+}