* added yaml examples for trainer
Signed-off-by: narayanasabari <[email protected]>
* Examples fixed
Signed-off-by: narayanasabari <[email protected]>
* fix(examples): address maintainer feedback on YAML examples
- Remove hello-world, with-runtime, pytorch-simple examples
- Fix multi-node: remove custom image, use PET_* env var prefix
- Fix volcano: move podGroupPolicy to ClusterTrainingRuntime
- Fix podspec-overrides: move podTemplateOverrides to top-level,
use trainer.resourcesPerNode and trainer.env instead of invalid
ContainerOverride fields
- Trim all READMEs to avoid duplicating kubeflow.org docs
Signed-off-by: Sabari Narayana <[email protected]>
* feat(examples): migrate YAML examples to RuntimePatches API
Address review feedback from @andreyvelich on PR #2925:
- Rename 01-podspec-overrides.yaml -> 01-runtime-patches.yaml and rewrite
using the new spec.runtimePatches API (PodTemplateOverrides was removed
in favor of the manager-keyed RuntimePatches mechanism).
- Inline a TrainingRuntime in 04-multi-step.yaml so the dataset-initializer
step is actually defined; the default torch-distributed runtime has no
initializer replicatedJob and the example would fail validation otherwise.
- Replace nonexistent trainer.kubeflow.org/job-name selectors throughout
with jobset.sigs.k8s.io/jobset-name (for pods/podgroups/workloads) and
trainer.kubeflow.org/trainjob-ancestor-step (to target a specific step).
- Bump default PyTorch image to 2.10.0 to match upstream torch-distributed
ClusterTrainingRuntime.
- Add ManagedBy hint for MultiKueue in the Kueue example.
- Trim READMEs to avoid duplicating kubeflow.org documentation.
Signed-off-by: Sabari Narayana <[email protected]>
* docs(examples): address review feedback on YAML examples
Address review feedback from @andreyvelich on PR #2925:
- Flatten examples/yaml/: drop the basic/ and advanced/ subfolders since
there is only one "basic" example. All YAMLs now live directly under
examples/yaml/ and are renumbered 01-05.
- Remove namespace: default from all manifests (it is the implicit default
and adds noise).
- Use torchrun as the entrypoint in 01-multi-node.yaml and 02-runtime-patches.yaml,
with torch.distributed API in the script (matches the pattern in the
PyTorch user guide).
- Restructure examples/README.md: SDK section first, YAML section last;
list all framework subfolders (megatron, flux, xgboost, jax, local);
refresh SDK snippet from kubeflow/sdk README; link sdk.kubeflow.org.
Signed-off-by: Sabari Narayana <[email protected]>
---------
Signed-off-by: narayanasabari <[email protected]>
Signed-off-by: Sabari Narayana <[email protected]>
* fix(e2e): bump Megatron notebook Complete-wait timeout to 10m
The Megatron TP notebook waits for the TrainJob to reach Complete with
timeout=120. On the oracle-vm-gpu-a10-1 CI runner the happy-path time
from TrainJob creation to Complete is ~96s (measured on the last passing
GPU E2E run, 2026-04-15). Any runner slowdown, image-pull delay, or
GPU-advertisement latency on top of that pushes the test past the 120s
budget and papermill raises TimeoutError, even though the TrainJob is
still on track to finish.
Every GPU E2E run has been failing here since 2026-04-16 ~16:00 UTC on
every branch, with no functional repo change between the last pass and
first fail. Bumping to 600s gives enough headroom for cold image pulls
and scheduling variance without masking real failures (papermill's
outer --execution-timeout is 1800s).
Signed-off-by: XploY04 <[email protected]>
* fix(examples): update Megatron tokenizer library to null-text for 0.17.0
megatron-core 0.17.0 (published 2026-04-16 20:22 UTC) tightened the set
of accepted tokenizer library names in MegatronTokenizer.from_pretrained.
The bare "null" value is no longer accepted; the null-tokenizer keys are
now "null-text" and "null-multimodal":
0.16.1: if library not in ['byte-level', 'null']: assert tokenizer_path
0.17.0: if library not in ['byte-level', 'null-text', 'null-multimodal']:
assert tokenizer_path
The notebook's call with metadata_path={"library": "null"} therefore
triggers `AssertionError: Tokenizer path must be specified.` on the GPU
E2E runner, which now installs 0.17.0 by default. Renaming to
"null-text" routes through the same NullTokenizer(vocab_size) library
class the old "null" key used, so behavior is unchanged on 0.17.0.
This explains why GPU E2E ran green on 2026-04-15 (installed 0.16.1)
and started failing with the tokenizer assertion once 0.17.0 landed on
PyPI. The previous commit's 120s -> 600s wait timeout bump is what made
this assertion visible; with the old 120s budget the notebook timed out
before the TrainJob reached Failed and the real error was hidden.
Signed-off-by: XploY04 <[email protected]>
---------
Signed-off-by: XploY04 <[email protected]>
Replace DMatrix with QuantileDMatrix in distributed XGBoost training
example when using histogram tree method. This reduces memory usage
and aligns with XGBoost best practices for distributed workloads.
Fixes#3300
Signed-off-by: Suyash Nalawade <[email protected]>
* feat: add Megatron-Core GPT Tensor Parallelism example notebook
Signed-off-by: XploY04 <[email protected]>
* added megatron notebook to the e2e gpu test
Signed-off-by: XploY04 <[email protected]>
* feat: Parameterize Megatron core GPT notebook's tensor parallelism and GPU count, and update the e2e test workflow to pass these parameters.
Signed-off-by: XploY04 <[email protected]>
* change the number of gpus to 2
Signed-off-by: XploY04 <[email protected]>
* docs: update Megatron GPT tensor parallelism example notebook.
Signed-off-by: XploY04 <[email protected]>
* fix: correct minor typos and improve code readability in Megatron-Core GPT notebook
Signed-off-by: XploY04 <[email protected]>
* fix: update tensor model parallel size retrieval and improve code clarity in Megatron-Core GPT notebook
Signed-off-by: XploY04 <[email protected]>
* fix: correct wording in training function description for clarity in Megatron-Core GPT notebook
Signed-off-by: XploY04 <[email protected]>
* feat: add verification step for TrainJob completion in Megatron-Core GPT notebook
Signed-off-by: XploY04 <[email protected]>
* fix: use multi-node setup for Megatron TP to work with GPU time-slicing
GPU time-slicing (replicas=2) advertises 2 GPUs to Kubernetes but
exposes only 1 CUDA device inside the container. This caused torchrun
to launch 1 worker (WORLD_SIZE=1), failing Megatron's TP requirement
of world_size >= 2.
Switch from 1 node with 2 GPUs to 2 nodes with 1 GPU each. This
creates 2 pods, each getting 1 time-sliced GPU, giving WORLD_SIZE=2
for tensor parallelism across pods.
Signed-off-by: XploY04 <[email protected]>
* fix: use PyTorch devel image for Megatron compile_helpers() support
Megatron-Core's compile_helpers() requires `make` and `gcc` to build
C dataset helpers. The default runtime image doesn't include build
tools. Switch to the devel image which has them pre-installed.
Signed-off-by: XploY04 <[email protected]>
* fix: install megatron-core from git for compile_helpers() support
The pip package for megatron-core doesn't include the Makefile and
C source files needed by compile_helpers(). Install from the git
repository instead, which includes all source files for compiling
the C++ dataset helpers required by MockGPTDataset.
Signed-off-by: XploY04 <[email protected]>
* fix: install megatron-core from zip archive instead of git
The devel image doesn't have git installed, so pip install from
git+https fails. Use the GitHub zip archive URL instead, which
pip can download and extract directly without needing git.
Signed-off-by: XploY04 <[email protected]>
* fix: download Megatron C++ source files for compile_helpers()
The megatron-core pip package excludes the Makefile and helpers.cpp
needed by compile_helpers(). Download these two files from GitHub
at runtime before compilation. Also add pybind11 as an explicit
dependency since it's needed by the Makefile but only included as
a build-time dependency in the pip package.
Signed-off-by: XploY04 <[email protected]>
* fix: use correct null tokenizer library name for pip megatron-core
The pip release of megatron-core uses "null" as the tokenizer library
name, while the GitHub main branch uses "null-text". Change to "null"
to match the pip version.
Signed-off-by: XploY04 <[email protected]>
* fix: increase wait_for_job_status timeout to 1800s
The megatron training job takes longer than 900s due to pip install,
C++ compilation, and multi-node training setup on time-sliced GPUs.
Increase timeout to 1800s to match PAPERMILL_TIMEOUT.
Signed-off-by: XploY04 <[email protected]>
* fix: use runtime image with apt-get for build tools instead of devel image
The devel image (3.5GB) is too slow to pull, causing PAPERMILL_TIMEOUT.
Switch back to the runtime image and install only make and g++ via
apt-get (~50MB). This provides everything compile_helpers() needs
without the overhead of pulling the full devel image.
Signed-off-by: XploY04 <[email protected]>
* fix: disable NCCL shared memory to avoid /dev/shm limit in Kubernetes
Kubernetes pods get 64MB /dev/shm by default. NCCL needs more for
its communication buffers, causing checkpoint save to fail with
"Error while creating shared memory segment". Set NCCL_SHM_DISABLE=1
to use network sockets instead.
Also reduce wait_for_job_status timeout to 60s since training
completes quickly and the main wait happens in get_job_logs.
Signed-off-by: XploY04 <[email protected]>
* debug: add NCCL_DEBUG=INFO to diagnose /dev/shm failure on node-1
Signed-off-by: XploY04 <[email protected]>
* fix: mount /dev/shm as emptyDir to fix NCCL shared memory exhaustion
NCCL proxy service allocates ~33MB per communicator in /dev/shm. The
default Kubernetes /dev/shm is 64MB (Docker default), which is
insufficient for workloads that create multiple NCCL communicators.
Patch ClusterTrainingRuntimes in GPU e2e setup to mount /dev/shm as
a Memory-backed emptyDir, and remove unnecessary NCCL workaround
env vars from the Megatron notebook.
Signed-off-by: XploY04 <[email protected]>
* fix: remove dist_checkpointing from Megatron notebook
Megatron-Core dist_checkpointing uses multiprocessing.spawn internally
to create a Manager queue for async writes. The Kubeflow SDK generates
training scripts without an if __name__ == '__main__' guard, so the
spawned child re-imports the script and re-executes the training
function, causing a RuntimeError. Remove the checkpoint step since
training (Steps 1-5) is the core TP demonstration. Also remove unused
imports and fix the GPU prerequisites text.
Signed-off-by: XploY04 <[email protected]>
---------
Signed-off-by: XploY04 <[email protected]>
* feat: customize R generation to include GPUs
Flux detection of GPUs depends on hwloc plugins, and a newer
version. To get around any edge cases where the dependencies are missing,
we can easily generate the R from the expected resource spec (cores
and gpus). We can also add a shared memory mount to ensure the job
MPI gets all available shared memory of the host. The current
default is 64M (automatic from container runtime) and it can have
implications for MPI performance.
Signed-off-by: vsoch <[email protected]>
* review: use generateRange function for gpus
Also update volume to be added to container. This still
is pending testing on AWS!
Signed-off-by: vsoch <[email protected]>
* review and add gpu trainer example
Signed-off-by: vsoch <[email protected]>
---------
Signed-off-by: vsoch <[email protected]>
Co-authored-by: vsoch <[email protected]>
* Add XGBoost runtime plugin scaffold and register in framework
Signed-off-by: Krishna-kg732 <[email protected]>
Signed-off-by: krishna-kg732 <[email protected]>
* regenerated assets
Signed-off-by: Krishna-kg732 <[email protected]>
Signed-off-by: krishna-kg732 <[email protected]>
* Add XGBoost plugin unit tests and update framework test registry
Signed-off-by: Krishna-kg732 <[email protected]>
Signed-off-by: krishna-kg732 <[email protected]>
* added :Runtime image , unit tests for :(Resources are not set, Resources are set in TrainJob) ,using the ContainerTrainerPort instead of default port
Signed-off-by: krishna-kg732 <[email protected]>
* minor changes
Signed-off-by: krishna-kg732 <[email protected]>
* added plainML fallback(JAX+XGBOOST)
Signed-off-by: krishna-kg732 <[email protected]>
* test: add XGBoost runtime resource resolution unit tests
- Add test case for resources set in Runtime only
- Add test case for resources set in both Runtime and TrainJob
- Add xgboost_distributed.yaml ClusterTrainingRuntime manifest
- Update kustomization.yaml with xgboost-runtime image
Addresses feedback from PR #3200 for complete resource resolution test coverage.
Signed-off-by: krishna-kg732 <[email protected]>
* added autogenerat3ed files
Signed-off-by: krishna-kg732 <[email protected]>
* CI fixes
Signed-off-by: krishna-kg732 <[email protected]>
* chore: fix ci lint errors
Signed-off-by: krishna-kg732 <[email protected]>
* Add XGBoost integration and E2E tests, fix Validate guard clause and nil safety
Signed-off-by: krishna-kg732 <[email protected]>
* trigger ci
Signed-off-by: krishna-kg732 <[email protected]>
* fix: Add XGBoost support to test wrapper, CEL validation, and CI publish workflow
- Add XGBoostPolicy() method to MLPolicySourceWrapper for integration tests
- Add has(self.xgboost) to MLPolicy CEL mutual-exclusivity validation rule
- Regenerate CRD manifests with updated CEL rule
- Add xgboost-runtime to build-and-push-images CI workflow
- Add XGBoost image build step to E2E setup script (commented out due to Kind bug)
Signed-off-by: krishna-kg732 <[email protected]>
* docs: Add XGBoost distributed training example notebook
Add example notebook demonstrating distributed XGBoost training using
the Kubeflow Trainer SDK and the xgboost-distributed ClusterTrainingRuntime.
The notebook uses the modern XGBoost Collective API with RabitTracker
for distributed coordination and synthetic classification data.
Signed-off-by: krishna-kg732 <[email protected]>
* test: Skip XGBoost E2E test until runtime image is published
Signed-off-by: krishna-kg732 <[email protected]>
* chore : address review comments
Signed-off-by: krishna-kg732 <[email protected]>
* chore:resolve merge conflicts
Signed-off-by: krishna-kg732 <[email protected]>
---------
Signed-off-by: Krishna-kg732 <[email protected]>
Signed-off-by: krishna-kg732 <[email protected]>
* feat: support for flux framework as hpc manager
Flux supports the majority of MPI flavors/variants, and can be used
to bootstrap MPI as a plugin. It adds other features for scheduling
and topology that can be used for simulations and ai/ml jobs.
This changeset adds the plugin implementation, including the
plugin module, tests, and an example with a small README to
serve as documentation for the time being.
Signed-off-by: vsoch <[email protected]>
* review: slim down use of annotations
We still need to put the original command in an annotation
to retieve later, but others can be re-derived from the
environment
Signed-off-by: vsoch <[email protected]>
* review: clean up tasks and remove extra FluxPolicySource
Signed-off-by: vsoch <[email protected]>
* review: ensure to unset trainer args
Signed-off-by: vsoch <[email protected]>
* review: better support for nproc per node and nodes
Signed-off-by: vsoch <[email protected]>
* review: add constants types for flux
Signed-off-by: vsoch <[email protected]>
* review: add gpu specification from resources
Signed-off-by: vsoch <[email protected]>
* Fix TrainJob NumProcPerNode
Signed-off-by: Andrey Velichkevich <[email protected]>
* Remove elasticPolicy from kubebuilder:validation
Signed-off-by: Andrey Velichkevich <[email protected]>
* chore: use go embed for long string templates flux plugin
Signed-off-by: vsoch <[email protected]>
* cleanup: removed un-used annotation
Signed-off-by: vsoch <[email protected]>
---------
Signed-off-by: vsoch <[email protected]>
Signed-off-by: Andrey Velichkevich <[email protected]>
Co-authored-by: vsoch <[email protected]>
Co-authored-by: Andrey Velichkevich <[email protected]>
* feat(examples): align speech and audio examples with MNIST TrainJob flow
This commit updates the speech-recognition and audio-classification notebooks to consistently use the Kubeflow Trainer SDK for TrainJob creation, status checking, and log streaming. It also adjusts the test workflows.
Signed-off-by: Sneha Das <[email protected]>
* fix: resolve torchcodec dependency and ddp timeout
Signed-off-by: Sneha Das <[email protected]>
* fix: align installation steps with mnist example
Signed-off-by: Sneha Das <[email protected]>
* fix: align installation step in speech recognition with mnist example
Signed-off-by: Sneha Das <[email protected]>
* fix: update SDK installation and paths in audio/speech notebooks
Signed-off-by: Sneha Das <[email protected]>
* chore: fix end of file newlines in notebooks
Signed-off-by: Sneha Das <[email protected]>
* fix: updated the installation path
Signed-off-by: Sneha Das <[email protected]>
---------
Signed-off-by: Sneha Das <[email protected]>
* feat(examples): added torch.compile to examples of pytorch and local
Signed-off-by: Ishtiyaque-Alam <[email protected]>
* Clearing the output of the notebook
Signed-off-by: Ishtiyaque-Alam <[email protected]>
* Reverting a change from jobset builder
Signed-off-by: Ishtiyaque-Alam <[email protected]>
---------
Signed-off-by: Ishtiyaque-Alam <[email protected]>
Signed-off-by: Ishtiyaque-Alam <[email protected]>
* update examples to reflect func_args now being unpacked
Signed-off-by: Brian Gallagher <[email protected]>
* fix: lower case func arguments and add types
Signed-off-by: Brian Gallagher <[email protected]>
---------
Signed-off-by: Brian Gallagher <[email protected]>
* Add question-answer example
Signed-off-by: solanyn <[email protected]>
* chore: remove unused lines, add TODO comment
Signed-off-by: solanyn <[email protected]>
* chore: update example description
Co-authored-by: Andrey Velichkevich <[email protected]>
Signed-off-by: Andrew Chen <[email protected]>
* chore: update question-answering example
* run train job on CPU
* reduce batch size, dataset size and train epochs
* make upload to bucket optional
* add notebook to e2e-test
* set model name as trainjob argument
Signed-off-by: solanyn <[email protected]>
* chore: extend e2e-run-notebook timeout
* e2e tests fail if trainjobs launched by notebook do not finish in 3s
* extends the timeout to 5min to block and wait for longer trainjobs until timeout or trainjob completes
Signed-off-by: solanyn <[email protected]>
* chore: update example to wait for trainjob running status
* revert change to e2e-run-notebook.sh
Signed-off-by: solanyn <[email protected]>
---------
Signed-off-by: solanyn <[email protected]>
Signed-off-by: Andrew Chen <[email protected]>
Co-authored-by: Andrey Velichkevich <[email protected]>
