Skip to content

docs(training_guides): preemptible TrainJobs with Kueue + checkpoint resume - #260

Merged
typhoonzero merged 3 commits into
masterfrom
codex/training-guides-kueue-preemption
Jun 10, 2026
Merged

typhoonzero merged 3 commits into
masterfrom
codex/training-guides-kueue-preemption

Conversation

@typhoonzero

@typhoonzero typhoonzero commented Jun 9, 2026

Copy link
Copy Markdown
Contributor

Summary

  • New training_guides doc + assets for running Kubeflow Trainer v2 TrainJobs that share a GPU cluster with online InferenceServices, using a two-ClusterQueue cohort (inference owns nominal quota; training borrows) and HuggingFace-Trainer-style checkpoint+resume.
  • Adds e2e/cases/c12_kueue_preemption.sh and wires it into run_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

Docdocs/en/training_guides/preemptible-trainjobs-with-kueue.mdx:

  • ASCII diagram of the cohort: how inference reserves quota and training borrows it.
  • The exact kubectl apply flow for the three asset bundles (cluster queues, priorities, namespace queues).
  • The four properties a preemption-safe TrainJob needs (RWX checkpoint PVC, frequent save_steps, Trainer.train(resume_from_checkpoint=…), terminationGracePeriodSeconds).
  • A coexistence-with-inference checklist — peak-sized nominal quota, one-way borrowing, reclaimWithinCohort: Any, PodsReady timeout, etc.
  • A verification recipe that reads the Preempted=True reason=InCohortReclamation condition directly off the Workload.

Assetsdocs/en/training_guides/assets/kueue/preemption/:

  • cluster-queues.yamlResourceFlavor + two cohort ClusterQueues with the inference/training asymmetric borrowing config.
  • workload-priorities.yamlWorkloadPriorityClass for inference (1000) and training (10).
  • local-queues.yaml — namespaced LocalQueues, rendered with sed s/<your-namespace>/$NS/.
  • checkpoint-pvc.yaml — RWX PVC for the trainer's output_dir.
  • training-runtime.yaml — checkpoint-aware TrainingRuntime using HF Trainer; auto-detects the latest checkpoint-N/ and passes it to .train(resume_from_checkpoint=…).
  • trainjob-low-priority.yaml, inference-service.yaml — sample submissions labelled for the respective queues/priorities.

E2Ee2e/cases/c12_kueue_preemption.sh:

  • Applies the cohort + queues + priorities + PVC + TrainingRuntime, all named with a per-run ID so concurrent runs don't clash.
  • Waits for the first checkpoint-N/ to land on the PVC (via kubectl 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).
  • Submits the high-priority preemptor; asserts the training Workload reaches Preempted=True reason=InCohortReclamation.
  • Waits for inference to finish; asserts the resumed pod logs [checkpoint] resuming from ….
  • Cleans up all cluster-scoped Kueue objects on EXIT.

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_checkpoint knobs — the doc calls this out explicitly.

Verification

  • yarn lint → 0 errors, 0 warnings.
  • bash run_all.sh C12 on g1-c1-x86 (NVIDIA Tesla P100, HAMI vGPU, upstream Kueue v0.18.0): PASS in 222s.

Timeline from the green run:

C12: checkpoint visible on PVC                          (training step 4)
C12: submitting high-priority preemptor Job
C12: preempted workload=trainjob-… reason=InCohortReclamation   (1s after submit)
C12: inference job complete                             (113s after submit)
C12: resumed trainer pod=…
C12: end-to-end preemption + resume verified

Test plan

  • yarn lint
  • bash e2e/run_all.sh C12 on g1-c1-x86
  • Spot-check the rendered MDX in yarn dev

🤖 Generated with Claude Code

Summary by CodeRabbit

  • Documentation
    • New guide: preemptible training with Kueue, checkpoint resume, and inference coexistence; includes step‑by‑step instructions and example manifests (cohorted queues, local queues, resource flavors, workload priority classes, inference service, training runtime, checkpoint PVC, and sample TrainJob).
  • Tests
    • Added end‑to‑end test validating preemption and checkpoint resume, plus runner support to mark and report skipped cases when Kueue is unavailable.

…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>
@coderabbitai

coderabbitai Bot commented Jun 9, 2026

Copy link
Copy Markdown

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: d9c452ab-31d5-4a21-a4ba-def9e287a134

📥 Commits

Reviewing files that changed from the base of the PR and between 268460a and 46b613b.

📒 Files selected for processing (4)
  • docs/en/training_guides/assets/kueue/preemption/cluster-queues.yaml
  • docs/en/training_guides/preemptible-trainjobs-with-kueue.mdx
  • e2e/cases/c12_kueue_preemption.sh
  • e2e/run_all.sh
✅ Files skipped from review due to trivial changes (1)
  • docs/en/training_guides/preemptible-trainjobs-with-kueue.mdx
🚧 Files skipped from review as they are similar to previous changes (3)
  • e2e/run_all.sh
  • docs/en/training_guides/assets/kueue/preemption/cluster-queues.yaml
  • e2e/cases/c12_kueue_preemption.sh

Walkthrough

Adds 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.

Changes

Kueue Preemption + Checkpoint Resume

Layer / File(s) Summary
Cohort Configuration: Flavors, Queues, and Priorities
docs/en/training_guides/assets/kueue/preemption/cluster-queues.yaml, docs/en/training_guides/assets/kueue/preemption/workload-priorities.yaml, docs/en/training_guides/assets/kueue/preemption/local-queues.yaml, docs/en/training_guides/assets/kueue/preemption/checkpoint-pvc.yaml
Core Kueue setup with c12-default ResourceFlavor, two ClusterQueues sharing cohort c12-shared (inference nominal quota with zero borrowing; training borrows and is reclaimable), LocalQueues, workload priorities, and RWX checkpoint PVC.
Example Workloads: Training Runtime and Jobs
docs/en/training_guides/assets/kueue/preemption/training-runtime.yaml, docs/en/training_guides/assets/kueue/preemption/trainjob-low-priority.yaml, docs/en/training_guides/assets/kueue/preemption/inference-service.yaml
TrainingRuntime with HuggingFace checkpoint persistence and resume detection, low-priority TrainJob, and high-priority InferenceService labeled for Kueue admission and GPU resource requests.
User Guide and Operational Documentation
docs/en/training_guides/preemptible-trainjobs-with-kueue.mdx, docs/en/training_guides/index.mdx
User-facing guide covering prerequisites, cohort design, checkpoint/resume mechanics and operational knobs, plus an index entry linking to the guide.
End-to-End Test: Preemption and Resume Validation
e2e/cases/c12_kueue_preemption.sh, e2e/run_all.sh
Automated e2e test that deploys the cohort, creates PVC and runtime, submits training then inference, verifies Preempted=True with InCohortReclamation, waits for re-admission, and confirms checkpoint resume via pod logs; runner updated to add the case and handle skips.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

  • alauda/aml-docs#259: Runner integration for e2e cases; related to adding the C12 preemption test and run_all.sh changes.
  • alauda/aml-docs#256: Changes to training guides index; related to adding the new guide entry in index.mdx.

Poem

🐰 I stitched a PVC with care,
I queued the jobs with nimble flair,
When inference knocks the training down,
Checkpoints whisper, "resume your crown",
Hopping logs confirm the sweet repair.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title 'docs(training_guides): preemptible TrainJobs with Kueue + checkpoint resume' clearly and directly summarizes the main change: adding documentation for preemptible TrainJobs with Kueue and checkpoint resume functionality.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch codex/training-guides-kueue-preemption

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@cloudflare-workers-and-pages

cloudflare-workers-and-pages Bot commented Jun 9, 2026

Copy link
Copy Markdown

Deploying alauda-ai with  Cloudflare Pages  Cloudflare Pages

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

View logs

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>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

🧹 Nitpick comments (1)
docs/en/training_guides/assets/kueue/preemption/training-runtime.yaml (1)

44-45: ⚡ Quick win

Consider documenting or removing hardcoded imagePullSecrets.

The hardcoded harbor-mlops-regcred secret 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 that imagePullSecrets also 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 imagePullSecrets block 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

📥 Commits

Reviewing files that changed from the base of the PR and between 07baf8e and 268460a.

📒 Files selected for processing (11)
  • docs/en/training_guides/assets/kueue/preemption/checkpoint-pvc.yaml
  • docs/en/training_guides/assets/kueue/preemption/cluster-queues.yaml
  • docs/en/training_guides/assets/kueue/preemption/inference-service.yaml
  • docs/en/training_guides/assets/kueue/preemption/local-queues.yaml
  • docs/en/training_guides/assets/kueue/preemption/training-runtime.yaml
  • docs/en/training_guides/assets/kueue/preemption/trainjob-low-priority.yaml
  • docs/en/training_guides/assets/kueue/preemption/workload-priorities.yaml
  • docs/en/training_guides/index.mdx
  • docs/en/training_guides/preemptible-trainjobs-with-kueue.mdx
  • e2e/cases/c12_kueue_preemption.sh
  • e2e/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 |

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

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.

Suggested change
| 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.

Comment thread e2e/run_all.sh Outdated
@typhoonzero
typhoonzero merged commit fddd5d1 into master Jun 10, 2026
3 checks passed
@typhoonzero
typhoonzero deleted the codex/training-guides-kueue-preemption branch June 10, 2026 06:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant