feat: Migrate RHAI notebooks to RuntimePatches API - #836
Conversation
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
|
Warning Review limit reached
Next review available in: 46 minutes You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Central YAML (base), Organization UI (inherited) Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughNine trainer resource notebooks replace Estimated code review effort: 3 (Moderate) | ~25 minutes 🚥 Pre-merge checks | ✅ 10✅ Passed checks (10 passed)
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
tests/trainer/resources/rhai_features_fsdp_shared_state.ipynb (1)
167-168:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winFSDP mode mismatch makes this notebook functionally identical to
rhai_features_fsdp_full_state.ipynb.Line 167 has the correct
shard_grad_opmode commented out, and line 168 appliesfull_shard— the same setting used by the full-state notebook. This means the shared-state notebook tests nothing distinct. The markdown title at line 7 also still reads "Full State" rather than "Shared State".🐛 Fix: restore the intended FSDP mode
- # training_args_kwargs["fsdp"] = "shard_grad_op auto_wrap" - training_args_kwargs["fsdp"] = "full_shard auto_wrap" + training_args_kwargs["fsdp"] = "shard_grad_op auto_wrap"And fix the notebook title cell:
-"# RHAI Features Test - TrainJob Submission with FSDP (Full State)\n", +"# RHAI Features Test - TrainJob Submission with FSDP (Shared State)\n",🤖 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 `@tests/trainer/resources/rhai_features_fsdp_shared_state.ipynb` around lines 167 - 168, The notebook uses the wrong FSDP mode and title: update the cell that sets training_args_kwargs["fsdp"] to restore the intended shared-state mode by setting it to "shard_grad_op auto_wrap" (undo the current "full_shard auto_wrap"), and edit the markdown title cell (currently "Full State") to read "Shared State" so the notebook and its configuration reflect the shared-state test scenario.tests/trainer/resources/lora.ipynb (1)
568-573:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winSame stale
pod_logsbug as inosft.ipynb.Line 568 fetches a new generator, but line 571 discards it and uses the already-exhausted
pod_logsgenerator from the cell at line 524. The log output will always be empty.🐛 Fix
-logs = client.get_job_logs(name=job_name, follow=False) - -# Collect all log lines from the generator into a list -logs = list(pod_logs) +logs = list(client.get_job_logs(name=job_name, follow=False)) log_text = "\n".join(str(line) for line in logs) print(log_text)🤖 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 `@tests/trainer/resources/lora.ipynb` around lines 568 - 573, The code fetches a new generator into "logs" via client.get_job_logs(name=job_name, follow=False) but then erroneously iterates the previously-exhausted "pod_logs" generator; replace uses of "pod_logs" with the newly fetched "logs" (or rename the fetched generator to pod_logs and remove the old exhausted variable) so that the list construction and join use the fresh generator returned by client.get_job_logs.tests/trainer/resources/osft.ipynb (1)
494-501:⚠️ Potential issue | 🟡 Minor | ⚡ Quick win
pod_logsis already exhausted when this cell executes — the log print is always empty.Line 496 fetches a new generator into
logs, but line 499 immediately overwrites it withlist(pod_logs)— wherepod_logswas assigned and fully consumed (list(pod_logs)) back at lines 452 and 455 in the previous cell. The result is that this cell always prints empty output.🐛 Fix: use the generator just fetched
-logs = client.get_job_logs(name=job_name, follow=False) - -# Collect all log lines from the generator into a list -logs = list(pod_logs) +# Collect all log lines from the generator into a list +logs = list(client.get_job_logs(name=job_name, follow=False)) log_text = "\n".join(str(line) for line in logs) print(log_text)🤖 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 `@tests/trainer/resources/osft.ipynb` around lines 494 - 501, The cell is printing an empty log because it ignores the newly fetched generator: it assigns logs = client.get_job_logs(name=job_name, follow=False) but then overwrites logs with list(pod_logs), where pod_logs is an already-exhausted generator; replace uses of pod_logs here with the newly returned generator (logs) so you consume the fresh generator (e.g., convert logs to a list or join its lines) and then build log_text from that list; update any references to pod_logs in this cell to use logs (or rename for clarity) and ensure you don't reuse an exhausted generator from a previous cell.
🤖 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 `@tests/trainer/resources/mnist.ipynb`:
- Line 437: The PVC name pvc_name (assigned from os.getenv("SHARED_PVC_NAME",
"")) is injected unguarded into the k8s manifest as claimName which causes an
opaque API error when empty; update the manifest assembly in mnist.ipynb to only
include the "persistentVolumeClaim": {"claimName": pvc_name} entry when pvc_name
is non-empty (truthy) or omit the entire volumeMount/volume entry
otherwise—locate the code that constructs the dict containing
"persistentVolumeClaim" and wrap or conditionally add that key based on pvc_name
to mirror the validation used in lora.ipynb.
In `@tests/trainer/resources/torchrun_failure.ipynb`:
- Around line 350-352: The test uses hard-coded patch targets
ReplicatedJobPatch(name="node") and ContainerPatch(name="node") which silently
no-op if the loaded TrainingRuntime uses different replica or container names;
update the test to derive the replica and container names from the runtime
metadata (th_runtime) or assert/validate the expected names before creating
patches, then construct ReplicatedJobPatch and ContainerPatch using those
derived names so PVC injection applies reliably (reference identifiers:
ReplicatedJobPatch, ContainerPatch, th_runtime).
---
Outside diff comments:
In `@tests/trainer/resources/lora.ipynb`:
- Around line 568-573: The code fetches a new generator into "logs" via
client.get_job_logs(name=job_name, follow=False) but then erroneously iterates
the previously-exhausted "pod_logs" generator; replace uses of "pod_logs" with
the newly fetched "logs" (or rename the fetched generator to pod_logs and remove
the old exhausted variable) so that the list construction and join use the fresh
generator returned by client.get_job_logs.
In `@tests/trainer/resources/osft.ipynb`:
- Around line 494-501: The cell is printing an empty log because it ignores the
newly fetched generator: it assigns logs = client.get_job_logs(name=job_name,
follow=False) but then overwrites logs with list(pod_logs), where pod_logs is an
already-exhausted generator; replace uses of pod_logs here with the newly
returned generator (logs) so you consume the fresh generator (e.g., convert logs
to a list or join its lines) and then build log_text from that list; update any
references to pod_logs in this cell to use logs (or rename for clarity) and
ensure you don't reuse an exhausted generator from a previous cell.
In `@tests/trainer/resources/rhai_features_fsdp_shared_state.ipynb`:
- Around line 167-168: The notebook uses the wrong FSDP mode and title: update
the cell that sets training_args_kwargs["fsdp"] to restore the intended
shared-state mode by setting it to "shard_grad_op auto_wrap" (undo the current
"full_shard auto_wrap"), and edit the markdown title cell (currently "Full
State") to read "Shared State" so the notebook and its configuration reflect the
shared-state test scenario.
🪄 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: Central YAML (base), Organization UI (inherited)
Review profile: CHILL
Plan: Enterprise
Run ID: 59113999-7d72-4863-847b-835add079f0a
📒 Files selected for processing (9)
tests/trainer/resources/lora.ipynbtests/trainer/resources/mnist.ipynbtests/trainer/resources/osft.ipynbtests/trainer/resources/rhai_features.ipynbtests/trainer/resources/rhai_features_deepspeed_stage0.ipynbtests/trainer/resources/rhai_features_fsdp_full_state.ipynbtests/trainer/resources/rhai_features_fsdp_shared_state.ipynbtests/trainer/resources/sft.ipynbtests/trainer/resources/torchrun_failure.ipynb
| " volumes=[\n", | ||
| " {\n", | ||
| " \"name\": \"work\",\n", | ||
| " \"persistentVolumeClaim\": {\"claimName\": pvc_name},\n", |
There was a problem hiding this comment.
pvc_name is used without an emptiness guard; an empty value produces an opaque Kubernetes API error.
pvc_name = os.getenv("SHARED_PVC_NAME", "") (line 400) defaults to "". The changed code at line 437 injects it directly as claimName. Other notebooks in the same repo (e.g., lora.ipynb) validate the PVC name before use.
pvc_name = os.getenv("SHARED_PVC_NAME", "")
print(f"[notebook] Using PVC: {pvc_name}")
+if not pvc_name:
+ raise RuntimeError("SHARED_PVC_NAME environment variable is required")🤖 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 `@tests/trainer/resources/mnist.ipynb` at line 437, The PVC name pvc_name
(assigned from os.getenv("SHARED_PVC_NAME", "")) is injected unguarded into the
k8s manifest as claimName which causes an opaque API error when empty; update
the manifest assembly in mnist.ipynb to only include the
"persistentVolumeClaim": {"claimName": pvc_name} entry when pvc_name is
non-empty (truthy) or omit the entire volumeMount/volume entry otherwise—locate
the code that constructs the dict containing "persistentVolumeClaim" and wrap or
conditionally add that key based on pvc_name to mirror the validation used in
lora.ipynb.
| " ReplicatedJobPatch(\n", | ||
| " name=\"node\",\n", | ||
| " template=JobTemplatePatch(\n", |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify whether runtime fixtures and notebooks consistently use "node" as replicated job/container names.
set -euo pipefail
echo "== Notebooks using ReplicatedJobPatch/ContainerPatch names =="
rg -n -C2 'ReplicatedJobPatch\(|ContainerPatch\(' --type-add 'ipynb:*.ipynb' -t ipynb
echo
echo "== Runtime manifests likely defining container/replicated job names =="
rg -n -C3 'kind:\s*TrainingRuntime|replicatedJobs:|containers:|name:\s*node' --glob '*.y*ml'
echo
echo "Expected result:"
echo "- Every TrainingRuntime used by these tests defines replicated job and container name 'node'."
echo "- If not, patch target names must be made runtime-aware or explicitly validated in notebook code."Repository: opendatahub-io/distributed-workloads
Length of output: 25132
🏁 Script executed:
#!/bin/bash
echo "== Search for all TrainingRuntime kind: definitions =="
rg -n 'kind:\s*TrainingRuntime' --glob '*.y*ml' -A 50 | head -200
echo
echo "== Find fixture setup or runtime references in torchrun_failure.ipynb =="
rg -n 'TrainingRuntime\|th_runtime\|runtime.*=' tests/trainer/resources/torchrun_failure.ipynb -B 2 -A 2 | head -100
echo
echo "== Check if there's a conftest or setup that defines runtime fixtures =="
find tests -name conftest.py -o -name '*fixture*' -o -name '*setup*' | head -20Repository: opendatahub-io/distributed-workloads
Length of output: 283
🏁 Script executed:
#!/bin/bash
echo "== Extract imports and patch usage from torchrun_failure.ipynb =="
rg -n 'from kubeflow\|import.*patch\|ReplicatedJobPatch\|ContainerPatch' tests/trainer/resources/torchrun_failure.ipynb -B 1 -A 1 | head -150
echo
echo "== Search for kubeflow trainer options module in codebase =="
find . -path '*/kubeflow/trainer/options*' -o -path '*/kubeflow/trainer*patch*' 2>/dev/null | head -20
echo
echo "== Look for any documentation or test files explaining patch matching =="
rg -l 'ReplicatedJobPatch\|patch.*name' --type py | head -10
echo
echo "== Check structure of torchrun_failure.ipynb around patch definitions =="
sed -n '340,380p' tests/trainer/resources/torchrun_failure.ipynb | head -50Repository: opendatahub-io/distributed-workloads
Length of output: 2506
🏁 Script executed:
#!/bin/bash
echo "== Find th_runtime definition and context in torchrun_failure.ipynb =="
rg -n 'th_runtime' tests/trainer/resources/torchrun_failure.ipynb -B 3 -A 3 | head -200Repository: opendatahub-io/distributed-workloads
Length of output: 930
🏁 Script executed:
#!/bin/bash
echo "== Check if there's any test validation or assertion that patch was applied =="
rg -n 'assert\|volume_mounts\|pvc\|PVC' tests/trainer/resources/torchrun_failure.ipynb -B 2 -A 2 | head -100
echo
echo "== Search for kubeflow.v1 or similar imports and patch logic =="
rg -n 'from kubeflow\|import.*kubeflow\|class.*Patch' --type py | head -20Repository: opendatahub-io/distributed-workloads
Length of output: 224
Hard-coded patch target names silently skip PVC injection if runtime structure differs
ReplicatedJobPatch(name="node") and ContainerPatch(name="node") assume runtime internals match a fixed naming convention. If a loaded TrainingRuntime uses different replicated job or container names, the patches silently fail to apply (Kubernetes strategic merge behavior), leaving the PVC mount uninjected and causing cryptic test failures. Derive target names from th_runtime metadata or validate expected names before submission.
Affected locations
- Line 350:
ReplicatedJobPatch(name="node") - Line 360:
ContainerPatch(name="node")
Also applies to: 360-366
🤖 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 `@tests/trainer/resources/torchrun_failure.ipynb` around lines 350 - 352, The
test uses hard-coded patch targets ReplicatedJobPatch(name="node") and
ContainerPatch(name="node") which silently no-op if the loaded TrainingRuntime
uses different replica or container names; update the test to derive the replica
and container names from the runtime metadata (th_runtime) or assert/validate
the expected names before creating patches, then construct ReplicatedJobPatch
and ContainerPatch using those derived names so PVC injection applies reliably
(reference identifiers: ReplicatedJobPatch, ContainerPatch, th_runtime).
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tests/trainer/resources/lora.ipynb`:
- Around line 444-454: Pin a compatible kubeflow-trainer-api version in the
resource requirements and update the RuntimePatch/TemplatePatch notebook
examples across tests/trainer/resources/lora.ipynb (lines 444-454 and 477-506),
mnist.ipynb (386-398 and 423-454), osft.ipynb (372-382 and 405-434),
rhai_features_deepspeed_stage0.ipynb (375-386 and 465-495),
rhai_features_fsdp_full_state.ipynb (357-368 and 447-469),
rhai_features_fsdp_shared_state.ipynb (365-376 and 455-485), sft.ipynb (458-468
and 491-522), and torchrun_failure.ipynb (296-306 and 345-373) to assert
generated TrainJob output preserves the patched volumeMounts,
persistentVolumeClaim, and required container fields.
🪄 Autofix
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: Central YAML (base), Organization UI (inherited)
Review profile: CHILL
Plan: Pro Plus
Run ID: ab8c8573-7046-4c89-8589-91a7362d9bb6
📒 Files selected for processing (9)
tests/trainer/resources/lora.ipynbtests/trainer/resources/mnist.ipynbtests/trainer/resources/osft.ipynbtests/trainer/resources/rhai_features.ipynbtests/trainer/resources/rhai_features_deepspeed_stage0.ipynbtests/trainer/resources/rhai_features_fsdp_full_state.ipynbtests/trainer/resources/rhai_features_fsdp_shared_state.ipynbtests/trainer/resources/sft.ipynbtests/trainer/resources/torchrun_failure.ipynb
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/trainer/resources/rhai_features.ipynb
| "from kubeflow.trainer.options import (\n", | ||
| " RuntimePatch,\n", | ||
| " TrainingRuntimeSpecPatch,\n", | ||
| " JobSetTemplatePatch,\n", | ||
| " JobSetSpecPatch,\n", | ||
| " ReplicatedJobPatch,\n", | ||
| " JobTemplatePatch,\n", | ||
| " JobSpecPatch,\n", | ||
| " PodTemplatePatch,\n", | ||
| " PodSpecPatch,\n", | ||
| " ContainerPatch\n", |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== Kubeflow Trainer dependency declarations =="
rg -n -C2 'kubeflow.*trainer|kubeflow-trainer' \
--glob 'pyproject.toml' --glob 'requirements*.txt' --glob 'Pipfile*' \
--glob 'poetry.lock' --glob 'uv.lock'
echo "== RuntimePatch call sites =="
rg -n -C3 'RuntimePatch\(|TrainingRuntimeSpecPatch\(|replicated_jobs=|containers=|volume_mounts=' \
tests/trainer/resources
echo "== Existing generated TrainJob assertions =="
rg -n -C3 'volumeMounts|persistentVolumeClaim|replicatedJobs|containers|TrainJob' \
testsRepository: opendatahub-io/distributed-workloads
Length of output: 221
🏁 Script executed:
#!/bin/bash
set +e
echo "== repo files likely to declare dependencies =="
git ls-files | rg '(^|/)(pyproject\.toml|requirements.*\.txt|Pipfile.*|poetry\.lock|uv\.lock|setup\.(py|cfg)|setup\.py)$' || true
echo
echo "== dependency mentions across tracked files =="
rg -n -C2 'kubeflow.*trainer|kubeflow-trainer|RuntimePatch|TrainingRuntimeSpecPatch|JobSetTemplatePatch|JobSetSpecPatch|ReplicatedJobPatch|JobTemplatePatch|JobSpecPatch|PodTemplatePatch|PodSpecPatch|ContainerPatch' . --glob '!**/.git/**' --glob '!**/build/**' --glob '!**/dist/**' || true
echo
echo "== trainer resource notebook snippets (first matching snippets around RuntimePatch) =="
python3 - <<'PY'
from pathlib import Path
import re
for p in sorted(Path('tests/trainer/resources').glob('*.ipynb')):
text = p.read_text(encoding='utf-8', errors='replace')
if 'RuntimePatch(' in text or 'kubeflow.trainer.options' in text:
print(f'--- {p}')
for i, line in enumerate(text.splitlines(), 1):
if 'RuntimePatch(' in line or 'kubeflow.trainer.options' in line:
start=max(1,i-5); end=min(len(text.splitlines()),i+80)
for n,l in enumerate(text.splitlines()[start-1:end], start):
print(f'{n}: {l}')
print()
PYRepository: opendatahub-io/distributed-workloads
Length of output: 50395
🏁 Script executed:
#!/bin/bash
set +e
echo "== trainer resource requirements =="
for f in tests/trainer/resources/requirements-cpu.txt tests/trainer/resources/requirements-cuda.txt tests/trainer/resources/requirements-rocm.txt; do
echo "--- $f"
if [ -f "$f" ]; then
rg -n 'kubeflow|trainer|version|==' "$f" || true
fi
done
echo
echo "== trainer README dependency/version references =="
for f in tests/trainer/README.md tests/trainer/resources/disconnected_env/README.md; do
if [ -f "$f" ]; then
echo "--- $f"
rg -n -C2 'kubeflow-trainer|kubeflow\.trainer|2\.0\.0|version|pip install|RuntimePatch' "$f" || true
fi
done
echo
echo "== RuntimePatch merged patch classes in trainer notebooks =="
python3 - <<'PY'
from pathlib import Path
import json, re
imports = {
"RuntimePatch", "TrainingRuntimeSpecPatch", "JobSetTemplatePatch",
"JobSetSpecPatch", "ReplicatedJobPatch", "JobTemplatePatch", "JobSpecPatch",
"PodTemplatePatch", "PodSpecPatch", "ContainerPatch", "Labels", "JobTemplate",
"Job", "WorkerArgs", "TrainJob"
}
for p in sorted(Path("tests/trainer/resources").glob("*.ipynb")):
nb = json.loads(p.read_text(encoding="utf-8"))
for cell in nb.get("cells", []):
src = "".join(cell.get("source", []))
cells_imports = [x for x in imports if x in src]
if "RuntimePatch(" in src or ("kubeflow.trainer.options" in src and any(f"x in src for x in imports" and x in src for x in imports)):
out = {str(e.get("name")): e.get("status") for e in cell.get("outputs", [])}
print(f"--- {p}, cells with RuntimePatch/options API: {len(cells_imports)}")
for i, line in enumerate(src.splitlines(), 1):
if any(x in line for x in imports) or "RuntimePatch(" in line:
print(f" L{i}: {line.strip()}")
if out:
print(f" outputs= {out}")
PY
echo
echo "== RuntimePatch implementation/source references =="
rg -n -C3 'class RuntimePatch|class TrainingRuntimeSpecPatch|class .*TemplatePatch|class ContainerPatch|def .*runtime_patch|RuntimePatches' . --glob '!**/.git/**' --glob '!**/build/**' --glob '!**/dist/**' || trueRepository: opendatahub-io/distributed-workloads
Length of output: 10697
🏁 Script executed:
#!/bin/bash
set +e
echo "== changed files/stat =="
git status --short || true
git diff --stat || true
echo
echo "== any tests/assertions mentioning generated runtime fields =="
rg -n -C2 'generated.*trainjob|trainjob.*generated|runtime_patch|runtimepatches|RuntimePatches|persistent_volume_claim|persistentVolumeClaim|volumeMounts|volume_mounts|replicated_jobs|replicatedJobs|containers|container' tests/trainer --glob '!**/.git/**' || true
echo
echo "== trainer requirements files, exact tracked versions around kubeflow =="
rg -n -C2 'kubeflow|trainer' . --glob '*requirements*.txt' --glob '*Pipfile*' --glob 'pyproject.toml' --glob 'uv.lock' --glob 'poetry.lock' || trueRepository: opendatahub-io/distributed-workloads
Length of output: 50393
Pin kubeflow-trainer-api for the nested RuntimePatch API and assert the TrainJob output.
The notebooks import RuntimePatch/*TemplatePatch from kubeflow.trainer.options, but the resource requirements only pin unrelated CUDA/cp/minio packages while the Trainer API is version-agnostic. If this SDK replaces spec.replicatedJobs/containers instead of merging by name, TrainJob pods lose required runtime container fields and PVC shares. Add a pinned compatible kubeflow-trainer-api version and a generated-TrainJob assertion for the patched volumeMounts, persistentVolumeClaim, and containers.
📍 Affects 8 files
tests/trainer/resources/lora.ipynb#L444-L454(this comment)tests/trainer/resources/lora.ipynb#L477-L506tests/trainer/resources/mnist.ipynb#L386-L398tests/trainer/resources/mnist.ipynb#L423-L454tests/trainer/resources/osft.ipynb#L372-L382tests/trainer/resources/osft.ipynb#L405-L434tests/trainer/resources/rhai_features_deepspeed_stage0.ipynb#L375-L386tests/trainer/resources/rhai_features_deepspeed_stage0.ipynb#L465-L495tests/trainer/resources/rhai_features_fsdp_full_state.ipynb#L357-L368tests/trainer/resources/rhai_features_fsdp_full_state.ipynb#L447-L469tests/trainer/resources/rhai_features_fsdp_shared_state.ipynb#L365-L376tests/trainer/resources/rhai_features_fsdp_shared_state.ipynb#L455-L485tests/trainer/resources/sft.ipynb#L458-L468tests/trainer/resources/sft.ipynb#L491-L522tests/trainer/resources/torchrun_failure.ipynb#L296-L306tests/trainer/resources/torchrun_failure.ipynb#L345-L373
🤖 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 `@tests/trainer/resources/lora.ipynb` around lines 444 - 454, Pin a compatible
kubeflow-trainer-api version in the resource requirements and update the
RuntimePatch/TemplatePatch notebook examples across
tests/trainer/resources/lora.ipynb (lines 444-454 and 477-506), mnist.ipynb
(386-398 and 423-454), osft.ipynb (372-382 and 405-434),
rhai_features_deepspeed_stage0.ipynb (375-386 and 465-495),
rhai_features_fsdp_full_state.ipynb (357-368 and 447-469),
rhai_features_fsdp_shared_state.ipynb (365-376 and 455-485), sft.ipynb (458-468
and 491-522), and torchrun_failure.ipynb (296-306 and 345-373) to assert
generated TrainJob output preserves the patched volumeMounts,
persistentVolumeClaim, and required container fields.
1 participant
Description
Migrate all RHAI test notebooks from the deprecated
PodTemplateOverridesAPI to the newRuntimePatchesAPI introduced in Kubeflow Trainer SDK 2.3.Changes
PodTemplateOverrides/PodTemplateOverride/PodSpecOverride/ContainerOverrideimports withRuntimePatch/TrainingRuntimeSpecPatch/JobSetTemplatePatch/JobSetSpecPatch/ReplicatedJobPatch/JobTemplatePatch/JobSpecPatch/PodTemplatePatch/PodSpecPatch/ContainerPatchoptions=[PodTemplateOverrides(...)]calls to use the newRuntimePatch(training_runtime_spec=...)structure"workspace"to"initializer"in Go test filesNotebooks updated
grpo.ipynblora.ipynbmnist.ipynbosft.ipynbsft.ipynbrhai_features.ipynbrhai_features_fsdp_full_state.ipynbrhai_features_fsdp_shared_state.ipynbrhai_features_deepspeed_stage0.ipynbtorchrun_failure.ipynbGo files updated
jobset_workflow_test.gotrainer_fashion_mnist_training_test.goHow Has This Been Tested?
RuntimePatchAPI structure matching the SDKsync/upstream-mainbranchPodTemplateOverridesconfiguration