mirror of https://github.com/hanzo-ml/trainer.git synced 2026-09-24 00:31:34 +00:00

Files

T

Sabari NarayanaandGitHub ba2bc4a07a feat(examples): Add kubectl-friendly YAML examples for TrainJob and TrainingRuntime (#2925 )

* 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]>

2026-05-01 16:41:03 +00:00

2.8 KiB

Raw Blame History

Kubeflow Trainer Examples

This directory contains examples for using Kubeflow Trainer with different interfaces and frameworks.

Directory Structure

examples/
├── pytorch/       # PyTorch SDK examples (AI Practitioners)
├── deepspeed/     # DeepSpeed framework examples
├── mlx/           # MLX framework examples
├── torchtune/     # TorchTune fine-tuning examples
├── megatron/      # Megatron-LM examples
├── flux/          # Flux workload manager examples
├── xgboost/       # XGBoost examples
├── jax/           # JAX framework examples
├── local/         # Local backend examples
└── yaml/          # YAML examples for kubectl users (Platform Admins)

For AI Practitioners (Python SDK)

Use the Kubeflow Python SDK for a code-first experience. Install with pip install -U kubeflow, then submit a TrainJob:

from kubeflow.trainer import TrainerClient, CustomTrainer, TrainJobTemplate

def get_torch_dist(learning_rate: str, num_epochs: str):
    import os
    import torch
    import torch.distributed as dist

    dist.init_process_group(backend="gloo")
    print("PyTorch Distributed Environment")
    print(f"WORLD_SIZE: {dist.get_world_size()}")
    print(f"RANK: {dist.get_rank()}")
    print(f"LOCAL_RANK: {os.environ['LOCAL_RANK']}")

    lr = float(learning_rate)
    epochs = int(num_epochs)
    loss = 1.0 - (lr * 2) - (epochs * 0.01)
    if dist.get_rank() == 0:
        print(f"loss={loss}")

template = TrainJobTemplate(
    runtime="torch-distributed",
    trainer=CustomTrainer(
        func=get_torch_dist,
        func_args={"learning_rate": "0.01", "num_epochs": "5"},
        num_nodes=3,
        resources_per_node={"cpu": 2},
    ),
)

job_id = TrainerClient().train(**template)
TrainerClient().wait_for_job_status(job_id)
print("\n".join(TrainerClient().get_job_logs(name=job_id)))

Browse Python SDK Examples

For Platform Administrators (YAML + kubectl)

Ready-to-use YAML manifests that can be applied directly with kubectl:

# Multi-node distributed training
kubectl apply -f yaml/01-multi-node.yaml

# Pod customization with the runtimePatches API
kubectl apply -f yaml/02-runtime-patches.yaml

Browse YAML Examples

Documentation

Contributing

Found a bug or have a feature request? Please open an issue!

Want to contribute an example? Check out our contributing guidelines.

2.8 KiB Raw Blame History