docs(training_guides): preemptible TrainJobs with Kueue + checkpoint resume - #260
Conversation
…resume New guide and assets for running Kubeflow Trainer v2 TrainJobs that can share a GPU cluster with online InferenceServices: a two-ClusterQueue cohort where inference owns the nominal quota and training borrows it, HuggingFace-Trainer-style checkpoint+resume so preemption costs a bounded amount of work, and the operational knobs you actually want to set when the same cluster serves both. - docs: `preemptible-trainjobs-with-kueue.mdx` walks through the cohort, the checkpoint-aware TrainingRuntime, and the inference/training coexistence checklist (peak-sized quota, one-way borrowing, reclaimWithinCohort: Any, PodsReady timeout, …). - assets: `assets/kueue/preemption/` ships ready-to-apply YAMLs for the cohort + cluster queues, namespaced LocalQueues, WorkloadPriorityClasses, the checkpoint PVC, the HF-Trainer-based TrainingRuntime, a sample low-priority TrainJob, and a sample InferenceService. - index: link the new guide alongside the existing Trainer v2 / Kueue entries. - e2e: `cases/c12_kueue_preemption.sh` end-to-end verifies the flow on HAMI vGPU — wires up the cohort, submits the TrainJob, waits for a checkpoint to land on the PVC, fires the high-priority preemptor, asserts `Preempted=True reason=InCohortReclamation` on the training Workload, lets inference finish, and asserts the resumed pod loads the checkpoint. Verified PASS in 222s on `g1-c1-x86` (Tesla P100, Kueue v0.18.0). The case skips with a clear message if the `kueue.x-k8s.io` API group is missing, so the rest of the harness keeps working on clusters without Kueue. The TrainingRuntime uses plain HuggingFace Trainer (not LlamaFactory) — LF's bootstrap hangs on toy models on Tesla P100 / HAMI in dev, while HF Trainer reproduces the same checkpoint semantics in sub-second per-step time. The recipe transfers directly to LlamaFactory, training_hub, mini_trainer or any other Trainer-based wrapper since they all expose the same `output_dir` / `save_strategy` / `resume_from_checkpoint` knobs. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (3)
WalkthroughAdds a new guide with example Kueue manifests, RWX checkpoint PVC and TrainingRuntime that supports checkpoint resume, Local/ClusterQueue and WorkloadPriority assets, and an end-to-end test script plus runner integration to validate preemption and resume behavior. ChangesKueue Preemption + Checkpoint Resume
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Deploying alauda-ai with
|
| Latest commit: |
46b613b
|
| Status: | ✅ Deploy successful! |
| Preview URL: | https://4b8bd443.alauda-ai.pages.dev |
| Branch Preview URL: | https://codex-training-guides-kueue.alauda-ai.pages.dev |
Document the namespace-level reservation pattern as a counterpart to the asymmetric inference/training cohort: each ClusterQueue carries both nominalQuota and borrowingLimit so tenants get a guaranteed floor while still bursting into shared cohort capacity when neighbours are idle. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
docs/en/training_guides/assets/kueue/preemption/training-runtime.yaml (1)
44-45: ⚡ Quick winConsider documenting or removing hardcoded imagePullSecrets.
The hardcoded
harbor-mlops-regcredsecret will cause pod failures in clusters where this secret doesn't exist. While line 59 instructs users to replace the image, it doesn't mention thatimagePullSecretsalso needs customization or removal.Consider one of:
- Add a comment on line 44 noting that users should replace or remove this based on their registry
- Comment out the
imagePullSecretsblock with an explanation- Remove it entirely if the default image is intended to be public
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/en/training_guides/assets/kueue/preemption/training-runtime.yaml` around lines 44 - 45, The hardcoded imagePullSecrets entry (imagePullSecrets: - name: harbor-mlops-regcred) can cause pod failures if that secret doesn't exist; update the training-runtime.yaml by either removing the imagePullSecrets block, commenting it out with a clear note, or replacing the hardcoded name with a placeholder and a comment instructing users to set or remove it based on their registry (reference the imagePullSecrets block and the harbor-mlops-regcred name so reviewers can locate and update that section).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/en/training_guides/preemptible-trainjobs-with-kueue.mdx`:
- Line 21: The markdown link text "Alauda Build of HAMI" currently points to the
Kueue site; update the URL target for that link in the table cell so the anchor
text "Alauda Build of HAMI" points to the correct HAMI/Alauda documentation or
repository instead of "https://kueue.sigs.k8s.io/". Locate the table row
containing "GPU device plugin" and the link with text "Alauda Build of HAMI" and
replace the href with the proper HAMI/Alauda URL (e.g., the official Alauda HAMI
docs or repo) while keeping the surrounding text and backticks
(`nvidia.com/gpualloc`, `gpucores`, `gpumem`) unchanged.
In `@e2e/run_all.sh`:
- Around line 29-31: The runner currently treats any non-zero exit as a failure;
update the test-execution logic (the loop that runs each test and checks the
variable rc) so that rc==1 is treated as a skip: increment the skip counter
(e.g., skip_count) and continue instead of counting it as a failure, while still
treating rc!=0 && rc!=1 as a genuine failure (increment fail_count). Apply this
change to the branch that runs "C12:GPU:cases/c12_kueue_preemption.sh" and the
other test-run branch handling lines 52-60 (the same rc-check/aggregation
logic).
---
Nitpick comments:
In `@docs/en/training_guides/assets/kueue/preemption/training-runtime.yaml`:
- Around line 44-45: The hardcoded imagePullSecrets entry (imagePullSecrets: -
name: harbor-mlops-regcred) can cause pod failures if that secret doesn't exist;
update the training-runtime.yaml by either removing the imagePullSecrets block,
commenting it out with a clear note, or replacing the hardcoded name with a
placeholder and a comment instructing users to set or remove it based on their
registry (reference the imagePullSecrets block and the harbor-mlops-regcred name
so reviewers can locate and update that section).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: c3eb94ca-0792-401b-b736-38c7860f8abd
📒 Files selected for processing (11)
docs/en/training_guides/assets/kueue/preemption/checkpoint-pvc.yamldocs/en/training_guides/assets/kueue/preemption/cluster-queues.yamldocs/en/training_guides/assets/kueue/preemption/inference-service.yamldocs/en/training_guides/assets/kueue/preemption/local-queues.yamldocs/en/training_guides/assets/kueue/preemption/training-runtime.yamldocs/en/training_guides/assets/kueue/preemption/trainjob-low-priority.yamldocs/en/training_guides/assets/kueue/preemption/workload-priorities.yamldocs/en/training_guides/index.mdxdocs/en/training_guides/preemptible-trainjobs-with-kueue.mdxe2e/cases/c12_kueue_preemption.she2e/run_all.sh
| | Kubeflow Trainer v2 | `trainer.kubeflow.org` API group; see [Fine-Tuning with Kubeflow Trainer v2](./fine-tune-with-trainer-v2) | | ||
| | Kueue (v0.13+ for `v1beta2` API) | See [Install Kueue](../kueue/install) | | ||
| | Shared RWX storage | The checkpoint PVC must be reachable from any node the trainer might land on after a re-admission | | ||
| | GPU device plugin | Examples use [Alauda Build of HAMI](https://kueue.sigs.k8s.io/) vGPU resources (`nvidia.com/gpualloc`, `gpucores`, `gpumem`); swap for `nvidia.com/gpu` if you use the upstream NVIDIA device plugin | |
There was a problem hiding this comment.
Fix the HAMI reference link target.
The link text says “Alauda Build of HAMI” but currently points to the Kueue site, which is misleading for readers following prerequisites.
Suggested patch
-| GPU device plugin | Examples use [Alauda Build of HAMI](https://kueue.sigs.k8s.io/) vGPU resources (`nvidia.com/gpualloc`, `gpucores`, `gpumem`); swap for `nvidia.com/gpu` if you use the upstream NVIDIA device plugin |
+| GPU device plugin | Examples use [Alauda Build of HAMI](https://github.com/Project-HAMi/HAMi) vGPU resources (`nvidia.com/gpualloc`, `gpucores`, `gpumem`); swap for `nvidia.com/gpu` if you use the upstream NVIDIA device plugin |📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| | GPU device plugin | Examples use [Alauda Build of HAMI](https://kueue.sigs.k8s.io/) vGPU resources (`nvidia.com/gpualloc`, `gpucores`, `gpumem`); swap for `nvidia.com/gpu` if you use the upstream NVIDIA device plugin | | |
| | GPU device plugin | Examples use [Alauda Build of HAMI](https://github.com/Project-HAMi/HAMi) vGPU resources (`nvidia.com/gpualloc`, `gpucores`, `gpumem`); swap for `nvidia.com/gpu` if you use the upstream NVIDIA device plugin | |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/en/training_guides/preemptible-trainjobs-with-kueue.mdx` at line 21, The
markdown link text "Alauda Build of HAMI" currently points to the Kueue site;
update the URL target for that link in the table cell so the anchor text "Alauda
Build of HAMI" points to the correct HAMI/Alauda documentation or repository
instead of "https://kueue.sigs.k8s.io/". Locate the table row containing "GPU
device plugin" and the link with text "Alauda Build of HAMI" and replace the
href with the proper HAMI/Alauda URL (e.g., the official Alauda HAMI docs or
repo) while keeping the surrounding text and backticks (`nvidia.com/gpualloc`,
`gpucores`, `gpumem`) unchanged.
Summary
e2e/cases/c12_kueue_preemption.shand wires it intorun_all.sh. The case verifies the whole flow end-to-end (cohort admit → checkpoint → preempt → re-admit → resume from checkpoint) and exits cleanly with a skip message if Kueue isn't installed.What's in the box
Doc —
docs/en/training_guides/preemptible-trainjobs-with-kueue.mdx:kubectl applyflow for the three asset bundles (cluster queues, priorities, namespace queues).save_steps,Trainer.train(resume_from_checkpoint=…),terminationGracePeriodSeconds).reclaimWithinCohort: Any,PodsReadytimeout, etc.Preempted=True reason=InCohortReclamationcondition directly off the Workload.Assets —
docs/en/training_guides/assets/kueue/preemption/:cluster-queues.yaml—ResourceFlavor+ two cohort ClusterQueues with the inference/training asymmetric borrowing config.workload-priorities.yaml—WorkloadPriorityClassfor inference (1000) and training (10).local-queues.yaml— namespaced LocalQueues, rendered withsed s/<your-namespace>/$NS/.checkpoint-pvc.yaml— RWX PVC for the trainer'soutput_dir.training-runtime.yaml— checkpoint-awareTrainingRuntimeusing HF Trainer; auto-detects the latestcheckpoint-N/and passes it to.train(resume_from_checkpoint=…).trainjob-low-priority.yaml,inference-service.yaml— sample submissions labelled for the respective queues/priorities.E2E —
e2e/cases/c12_kueue_preemption.sh:checkpoint-N/to land on the PVC (viakubectl exec ls, since HF Trainer's tqdm progress bar overwrites the "Saving model checkpoint to" log line — that nuance is captured in the script comments).Preempted=True reason=InCohortReclamation.[checkpoint] resuming from ….The TrainingRuntime uses plain HuggingFace Trainer (not LlamaFactory) — LlamaFactory's bootstrap hangs on toy models on Tesla P100 + HAMI in this dev cluster, while HF Trainer reproduces the same checkpoint semantics in sub-second per-step time on the same hardware. The recipe transfers directly to LlamaFactory / training_hub / mini_trainer since they all expose the same
output_dir/save_strategy/resume_from_checkpointknobs — the doc calls this out explicitly.Verification
yarn lint→ 0 errors, 0 warnings.bash run_all.sh C12ong1-c1-x86(NVIDIA Tesla P100, HAMI vGPU, upstream Kueue v0.18.0): PASS in 222s.Timeline from the green run:
Test plan
yarn lintbash e2e/run_all.sh C12ong1-c1-x86yarn dev🤖 Generated with Claude Code
Summary by CodeRabbit