fix(snapshotter): address scheduling review feedback

Signed-off-by: Yuan Chen <yuanchen97@gmail.com>

Author: yuanchen8911

docs/user/cli-reference.md

@@ -1172,7 +1172,7 @@ spec:
       serviceAccountName: aicr
       nodeSelector:
         my-org/gpu-pool: "true"
-      tolerations:
+      tolerations:                         # [] clears the live snapshot agent's tolerate-all default
         - "gpu-type=h100:NoSchedule"
       requireGpu: true
     execution:

docs/user/validation.md

@@ -910,7 +910,7 @@ aicr validate \
   --toleration dedicated=worker-workload:NoExecute
 ```
 
-These flags affect the inner benchmark pods that run on GPU nodes (NCCL workers, Dynamo workers). When `--snapshot` is omitted, they also configure the preliminary live snapshot agent. With no toleration override, that agent tolerates all taints; an explicit `spec.validate.agent.tolerations: []` clears that default. Neither flag affects the validator orchestrator Job itself.
+These flags affect the inner benchmark pods that run on GPU nodes (NCCL workers, Dynamo workers). The example above supplies `--snapshot`, so it does not launch the live snapshot agent. When `--snapshot` is omitted, the flags also configure that preliminary agent. With no toleration override, the agent tolerates all taints; an explicit `spec.validate.agent.tolerations: []` clears that default. Neither flag affects the validator orchestrator Job itself.
 
 For `inference-perf` specifically, `--node-selector` narrows the pool of candidate GPU nodes — the validator then picks the candidate with the most free GPUs (subtracting same-ledger occupancy only — DRA allocations from DRA capacity, device-plugin requests from device-plugin capacity — and skipping DRA candidates that carry scalar `nvidia.com/gpu` workloads) and pins all Dynamo Frontend + worker pods to that node via `kubernetes.io/hostname`. The AIPerf benchmark runner pod is CPU-only, uses a tolerate-all / no-nodeSelector pod spec, and is unaffected by these flags.
 

pkg/cli/validate.go

@@ -135,7 +135,9 @@ func resolveValidateNodeSelector(cmd *cli.Command, resolved *config.ValidateReso
 // inference-perf that want to mirror the target node's taints by default
 // must distinguish "operator opted into tolerate-all" from "operator said
 // nothing". Returning nil here when neither CLI nor config set the field
-// keeps the env var unset, so the inner validator context sees nil.
+// keeps the env var unset, so the inner validator context sees nil. The live
+// snapshot path consumes that same nil as its signal to apply the agent's
+// tolerate-all default at the Job projection boundary.
 func resolveValidateTolerations(cmd *cli.Command, resolved *config.ValidateResolved) ([]corev1.Toleration, error) {
 	if cmd.IsSet("toleration") {
 		tols, err := snapshotter.ParseTolerations(cmd.StringSlice("toleration"))

pkg/snapshotter/agent.go

@@ -148,6 +148,32 @@ type AgentConfig struct {
 	Limits corev1.ResourceList
 }
 
+// buildAgentConfig projects snapshotter configuration onto the deployer's
+// Job configuration. Keep scheduling defaults at this projection boundary so
+// every snapshot-agent caller gets the same nil-versus-empty behavior.
+func buildAgentConfig(config *AgentConfig, agentOutput string) agent.Config {
+	return agent.Config{
+		Namespace:          config.Namespace,
+		ServiceAccountName: config.ServiceAccountName,
+		JobName:            config.JobName,
+		Image:              config.Image,
+		ImagePullSecrets:   config.ImagePullSecrets,
+		NodeSelector:       config.NodeSelector,
+		Tolerations:        effectiveAgentTolerations(config.Tolerations),
+		Output:             agentOutput,
+		Debug:              config.Debug,
+		Privileged:         config.Privileged,
+		RequireGPU:         config.RequireGPU,
+		RuntimeClassName:   config.RuntimeClassName,
+		MaxNodesPerEntry:   config.MaxNodesPerEntry,
+		OS:                 config.OS,
+		ClusterConfigPath:  config.ClusterConfigPath,
+		DiscoverNetwork:    config.DiscoverNetwork,
+		Requests:           config.Requests,
+		Limits:             config.Limits,
+	}
+}
+
 // deployAndWaitForResult handles the common deploy-wait-retrieve lifecycle for an agent Job.
 // It creates the deployer, deploys RBAC and the Job, streams logs, waits for completion,
 // and retrieves the snapshot data from the result ConfigMap.
@@ -170,26 +196,7 @@ func deployAndWaitForResult(ctx context.Context, clientset k8sclient.Interface,
 	// name the injected selector (TOCTOU: node may be cordoned after detection).
 	autoInjectedGPUSelector := maybeInjectGPUNodeSelector(ctx, clientset, config)
 
-	agentConfig := agent.Config{
-		Namespace:          config.Namespace,
-		ServiceAccountName: config.ServiceAccountName,
-		JobName:            config.JobName,
-		Image:              config.Image,
-		ImagePullSecrets:   config.ImagePullSecrets,
-		NodeSelector:       config.NodeSelector,
-		Tolerations:        effectiveAgentTolerations(config.Tolerations),
-		Output:             agentOutput,
-		Debug:              config.Debug,
-		Privileged:         config.Privileged,
-		RequireGPU:         config.RequireGPU,
-		RuntimeClassName:   config.RuntimeClassName,
-		MaxNodesPerEntry:   config.MaxNodesPerEntry,
-		OS:                 config.OS,
-		ClusterConfigPath:  config.ClusterConfigPath,
-		DiscoverNetwork:    config.DiscoverNetwork,
-		Requests:           config.Requests,
-		Limits:             config.Limits,
-	}
+	agentConfig := buildAgentConfig(config, agentOutput)
 
 	deployer := agent.NewDeployer(clientset, agentConfig)
 
@@ -268,6 +275,7 @@ func deployAndWaitForResult(ctx context.Context, clientset k8sclient.Interface,
 			fmt.Fprintln(logWriter(), logs)
 			fmt.Fprintln(logWriter(), "--- end logs ---")
 		}
+		isTransient := errors.IsTransient(waitErr)
 		msg := "job failed"
 		if autoInjectedGPUSelector {
 			msg = "job failed (auto-injected node selector nvidia.com/gpu.present=true — " +
@@ -276,11 +284,15 @@ func deployAndWaitForResult(ctx context.Context, clientset k8sclient.Interface,
 				"key=value:effect. To override placement, pass --node-selector " +
 				"kubernetes.io/hostname=<gpu-node>; --require-gpu selects a node " +
 				"advertising the nvidia.com/gpu resource)"
+		} else if isTransient {
+			msg = "job failed (verify target nodes are Ready and schedulable; if " +
+				"tolerations were explicitly cleared or replaced, pass a matching " +
+				"--toleration key=value:effect)"
 		}
 		// A wait that exceeded the deadline (pending pod, image pull, no schedulable
 		// node) is transient and retryable — classify it as ErrCodeTimeout rather
 		// than masking it as a deterministic ErrCodeInternal failure.
-		if errors.IsTransient(waitErr) {
+		if isTransient {
 			return nil, errors.Wrap(errors.ErrCodeTimeout, msg, waitErr)
 		}
 		return nil, errors.Wrap(errors.ErrCodeInternal, msg, waitErr)

pkg/snapshotter/agent_test.go

@@ -52,7 +52,7 @@ func TestDefaultTolerations(t *testing.T) {
 	}
 }
 
-func TestEffectiveAgentTolerations(t *testing.T) {
+func TestBuildAgentConfigTolerations(t *testing.T) {
 	explicit := []corev1.Toleration{
 		{
 			Key:      "dedicated",
@@ -84,9 +84,9 @@ func TestEffectiveAgentTolerations(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got := effectiveAgentTolerations(tt.input)
+			got := buildAgentConfig(&AgentConfig{Tolerations: tt.input}, "snapshot.yaml").Tolerations
 			if !reflect.DeepEqual(got, tt.want) {
-				t.Errorf("effectiveAgentTolerations() = %#v, want %#v", got, tt.want)
+				t.Errorf("buildAgentConfig().Tolerations = %#v, want %#v", got, tt.want)
 			}
 		})
 	}