Batch retry schedulability probes per scheduling round

[dejanzele]

When the retry engine is on, a mutation retry probes schedulability through SubmitChecker.Check once per job, inside the scheduling cycle. Under a mass failure, hundreds of failed jobs produce hundreds of probe calls in one cycle, and cycle time is the scheduler's main performance limit. Each consumer in the cycle also rebuilds the queue-to-policy map.

Earlier work:

• Wire the retry policy engine into the scheduler behind a feature flag #5001 wired the retry engine into the scheduler behind a feature flag.
• Add avoidSameNode and memory bump mutations to retry policy rules #5012 added the memory bump and avoidSameNode mutations that need the probe.
• The review of Wire the retry policy engine into the scheduler behind a feature flag #5001 flagged the per-job probe cost, and this change is the agreed follow-up.

The solution moves the retry decision into a planning pass. planRetryDecisions resolves the decision for every failed run in the round: one engine evaluation per job, then at most two batched Check calls for the whole round, one for memory-grown candidates and one for node-anti-affinity candidates. The two calls stay separate because each gate has its own terminal reason. generateUpdateMessagesFromJob consumes the plans. A job without a plan falls back to the per-job path and logs a warning. A future drift between the planning guard and the failure branch thus degrades to the slow path, and the behavior stays correct. The checker returns partial results when it reaches its time limits. The limits apply per call, so the probe asks again with a fresh budget, up to three calls, for the representatives an earlier call did not reach. A job that stays unprobed after that keeps its granted retry and its mutation: an absent result means "not probed", not "unschedulable". The probe also checks one representative per queue and scheduling key and applies the verdict to the whole class, because the key covers the placement requirements and the checker applies a per-queue resource limit. A mass failure affects many jobs of few distinct shapes, so the probe cost scales with the number of shapes, not with the number of jobs, and the time limits stay out of reach. The cycle builds the queue-to-policy map once and passes it to the update messages and the expiry sweep.

Tests:

• A round-batch test asserts that three failed jobs produce one probe call for a memory bump rule, and two calls when the rule also sets avoidSameNode.
• A full-cycle test drives a categorized failure through Scheduler.cycle with the flag on: the engine grants the retry, a retryable failure event goes out, and the job requeues.
• The retry suite now runs through the batch entry point, so the existing assertions on terminal reasons, compounding bumps, and budgets pin the refactor.

Validation on a live stack (compose plus kind): I built the scheduler from this branch and ran the four retry testsuite cases from #5007 against it. All four pass in 1m34s, including the OOM memory bump to success and the retry-limit exhaustion walk.

[datadog-armadaproject]

⚠️ Warnings

🚦 1 Pipeline job failed

CI | lint / Lint Go     

See error

Lint errors due to unused variable 'txn' in multiple lines: retry_policy_test.go:306, 329, 348, 372, 392, 461, 644, 787, 864, 907.

Useful? React with 👍 / 👎

_{This comment will be updated automatically if new data arrives.
🔗 Commit SHA: b6347c7 | Docs | Datadog PR Page | Give us feedback!}

[greptile-apps]

Greptile Summary

The PR moves retry schedulability decisions into a per-round planning pass, reducing repeated checker work while preserving retry mutations and terminal reasons.

• Evaluates retry policy decisions once per failed job.
• Batches memory-bump and node-anti-affinity probes separately.
• Deduplicates equivalent candidates by queue and scheduling key.
• Retries partial probes with fresh time budgets and lets unresolved retries proceed.
• Reuses the queue-to-policy mapping across update generation and lease expiry.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; the prior deadline-omission and cross-queue verdict-leakage issues are addressed in the current code.

Important Files Changed

Filename	Overview
internal/scheduler/scheduler.go	Introduces round-level retry planning, batched representative probes, partial-result handling, and shared queue-policy lookup without leaving the previously reported failures reachable.
internal/scheduler/retry_policy_test.go	Routes retry tests through the batch entry point and adds coverage for batching, queue-separated classes, and repeated partial probes.
internal/scheduler/scheduler_test.go	Adds full-cycle retry-policy coverage and expands the submit-checker fixture to model call counts and omitted partial results.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Failed jobs in scheduling round] --> B[Evaluate retry policy per job]
  B --> C[Collect memory-bump candidates]
  C --> D[Probe representatives by queue and scheduling key]
  D --> E[Collect node-affinity candidates]
  E --> F[Probe representatives by queue and scheduling key]
  F --> G[Generate retry or terminal events from plans]
  D -->|No result after retries| H[Keep granted retry and mutation]
  H --> E

Loading

_{Reviews (8): Last reviewed commit: "Abort test transactions with t.Cleanup i..." | Re-trigger Greptile}