feat(certification): support KAI gang scheduling - #303
Merged
Merged
Conversation
Certification had no way to select a gang-aware scheduler, so multi-node certification jobs could be partially scheduled instead of admitted as a gang. WorkloadRun already exposes spec.gangScheduler; this reuses the same GangSchedulerSpec on CertificationSpec rather than adding a second shape. When set, every category's resolved TrainingRuntime gets schedulerName on each replicatedJob pod spec and the kai.scheduler/queue label on the job template metadata, which is where the WorkloadRun path already writes them. An omitted queue defaults to default-queue. The rewrite runs after catalog and platform overrides resolve, so an override cannot restore default-scheduler. training/nemotron5-8b and nemotron5-56b hardcode schedulerName: default-scheduler and now pick up the configured scheduler instead. nvcrectl certification render applies the same rewrite in both the dry-run and offline branches, so rendered manifests match what the controller creates. An empty schedulerName is treated as "not configured", matching applyGangScheduler on the WorkloadRun path. The CRD requires a non-empty value, but nvcrectl renders straight from a file without consulting the API server, so without the guard a typo would render pod templates pinned to an empty scheduler while still carrying a queue label. Certifications without gangScheduler render byte-identically to before, verified against a binary built from main across GB200, H100 and GB300 on aws, gcp, azure and oci. Renames the local platform variable in the certification controller to detectedPlatform so it stops shadowing the pkg/platform package import, matching what workloadrun_controller.go already does. Closes #300 Signed-off-by: Ebot Ndip-Agbor <endipagbor@nvidia.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: ⛔ Files ignored due to path filters (50)
📒 Files selected for processing (12)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 WalkthroughWalkthroughCertification adds an optional Estimated code review effort: 4 (Complex) | ~45 minutes Merge Risk: ⚪ Minimal · up to The generated CRD already exposes the new gang scheduler configuration, so no actionable merge risk remains. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (4 passed)
Full details: Linked Issues checkExplanation The reviewable changes satisfy the main requirements in issue Resolution Inspect the excluded generated artifacts, especially helm/cluster-readiness-engine/crds/nvcre.nvidia.com_certifications.yaml, and confirm that the Certification CRD exposes spec.gangScheduler with the expected schedulerName and queue schema. Also verify api/v1alpha1/zz_generated.deepcopy.go if required by the repository generation process. Full details: Out of Scope Changes checkExplanation The changes remain within issue Full details: Docstring CoverageExplanation Docstring coverage is 82.35% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 17 functions across 9 files. (3 skipped: 3 unsupported.) ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
Contributor
lalitadithya
approved these changes
Sep 4, 2026
2 tasks done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
WorkloadRunhas supportedspec.gangSchedulerfor a while.Certificationhad no equivalent, so there was no way to put certification workloads on KAI Scheduler. Multi-node certification jobs could get partially scheduled instead of admitted as a gang, and the placed pods sit on GPUs waiting for peers that never arrive.This adds the same field to
CertificationSpec, reusing the existingGangSchedulerSpectype rather than inventing a second scheduler config shape:It is certification wide and applies to every entry in
spec.categories.What it does
For each category, every
replicatedJobin the resolvedTrainingRuntimegets:schedulerNameon the pod spec, atreplicatedJobs[].template.spec.template.speckai.scheduler/queuelabel atreplicatedJobs[].template.metadata.labelsThose are the same two locations
BuildTorchRuntimeandBuildMPIRuntimealready write for a WorkloadRun, so the two paths stay consistent. An omittedqueuedefaults todefault-queue. MPI categories get both thenodeand thelauncherjob.The rewrite runs after catalog and platform overrides resolve, so an override cannot put
default-schedulerback. That ordering matters:training/nemotron5-8bandtraining/nemotron5-56bboth hardcodeschedulerName: default-scheduler, and they now pick up the configured scheduler instead.nvcrectl certification renderapplies the same rewrite in both the dry-run and the offline branch, so what you inspect matches what the controller creates.One thing worth calling out
The helper treats an empty
schedulerNameas "not configured", the same wayapplyGangSchedulerdoes on the WorkloadRun path. The CRD requires a non-empty value, butnvcrectl renderreads a file and never talks to the API server, so without that guard a typo like this:renders every pod template with
schedulerName: ""plus a KAI queue label, exits 0, and says nothing. The manifest looks configured but runs on the default scheduler. Found while reviewing this branch, fixed here, and there is a test case that fails without the guard.Also renamed the local
platformvariable in the certification controller todetectedPlatform. It was shadowing thepkg/platformpackage import.workloadrun_controller.goalready calls it that.Verification
make buildandmake testare green. 25 packages, no failures.gangSchedulerrender byte identically to a binary built frommain, checked across GB200, H100 and GB300 on aws, gcp, azure and oci.make lintdoes not run locally here.bin/golangci-lintis built with go1.26 and the repo targetsgo 1.27.0. Pristinemainfails the same way, so it is not from this change, but it does mean CI is the first real lint check.No ADR. The original WorkloadRun gang scheduling shipped without one and this reuses that contract rather than deciding anything new.
Related Issue
Closes #300
Type of Change
Component(s) Affected
Testing
12 new cases, all following the
testutil.TestCaseParserand golden file convention:pkg/platform/testdata/apply-gang-scheduler-deps/(9): torch and MPI shapes, creating a missing labels map, preserving the launcher's existing labels, queue defaulting, overwriting a hardcoded scheduler, non-TrainingRuntime dependencies untouched, malformed shapes skipped, and the empty scheduler name no-op.pkg/certification/testdata/certification-render-gang-scheduler/(4): render path with an explicit queue, the nemotrondefault-schedulerreplacement, queue defaulting, and gangScheduler omitted. The offline branch ofrunCertificationRenderwas extracted intoresolveWorkflowsOfflineso the test drives the real code path instead of a copy of it.cmd/integration/testdata/validation/certification-gangscheduler-*(6): CRD admission for valid input, empty and omittedschedulerName, a queue failing the label value pattern, a queue over 63 chars, and the field absent.cmd/integration/testdata/reconcile/certification-gangscheduler-{nccl,training}(2): end to end through the controller, covering both MPI jobs and thedefault-schedulerreplacement.Manual check with
nvcrectl certification render --platform awson a cert withcommunication/nccl-all-reduceplustraining/nemotron5-8b: 3 pod specs getkai-schedulerand 3 jobs get the queue label. Droppingqueuegivesdefault-queue. DroppinggangSchedulergives byte identical output tomain.Checklist
git commit -s)make manifests generaterun (if*_types.gowas modified)