kubeflow / kubeflow/sdk

Specialized Trainer Abstractions and RuntimeConfig for the Kubeflow SDK

Open
#285 3 comments 0 reactions 0 assignees View on GitHub
lifecycle/stale
Dominant language
Python
Stars
148
Forks
262
Avg merge
1d 2h
Merged PRs (30d)
1

Description

## Problem

The current SDK offers two trainer abstractions that leave a usability gap for the majority of distributed training workloads:

- **`CustomTrainer`** is generic but requires users to manually look up runtimes by name, provides no framework-specific validation, and mixes runtime-environment settings (`packages_to_install`, `pip_index_urls`, `env`) with scaling/training concerns (`num_nodes`, `resources_per_node`, `func`) in a single flat dataclass.

- **`BuiltinTrainer`** is narrow, hardcoded to `TorchTuneConfig`, and cannot be extended to other config-driven frameworks without modifying the class itself.

For standard use cases like "run this PyTorch DDP function on N nodes" or "run this MPI script across a cluster," neither abstraction fits cleanly.

## Proposal

1. **Specialized, framework-aware trainers** (`TorchTrainer`, `MPITrainer`, `JAXTrainer`, `XGBoostTrainer`) that:
- Auto-discover the correct `ClusterTrainingRuntime` using the existing `trainer.kubeflow.org/framework` label
- Validate trainer/runtime compatibility at submission time (not at execution time)
- Expose typed, framework-specific arguments (e.g., `max_restarts` for Torch, `num_proc_per_node` for MPI)
- Share a common `BaseTrainer` abstract interface enabling community extensions

2. **`RuntimeConfig` dataclass** that cleanly separates per-job runtime environment settings (packages, pip config, env vars) from training logic, passed as a dedicated parameter to `TrainerClient.train()`.

Both changes are additive and 100% backward-compatible. Existing `CustomTrainer` and `BuiltinTrainer` usage is unaffected.

## Example

**Before:**
```python
runtime = client.get_runtime("torch-distributed")
job_id = client.train(
runtime=runtime,
trainer=CustomTrainer(
func=train_fn,
func_args={"lr": 1e-4},
packages_to_install=["transformers"],
env={"NCCL_DEBUG": "INFO"},
num_nodes=4,
resources_per_node={"gpu": 1},
),
)
```

**After:**
```python
job_id = client.train(
trainer=TorchTrainer(
func=train_fn,
func_args={"lr": 1e-4},
num_nodes=4,
resources_per_node={"gpu": 1},
max_restarts=3,
),
runtime_config=RuntimeConfig(
packages=["transformers"],
env={"NCCL_DEBUG": "INFO"},
),
)
```

## References

- KEP-2170: [Kubeflow Trainer V2 API](https://github.com/kubeflow/trainer/blob/master/docs/proposals/2170-kubeflow-trainer-v2/README.md)
- Runtime framework label: [Runtime Guide](https://www.kubeflow.org/docs/components/trainer/operator-guides/runtime/)
- Full design proposal will be submitted as a PR to this repo

/cc @kubeflow/kubeflow-sdk-team

Contributor guide

Open the contributing guide

Research direction

Start with the existing CustomTrainer, BuiltinTrainer, and TrainerClient.train entry points, then read KEP-2170 and the Runtime Guide linked in the issue. Define the specialized trainer interface and RuntimeConfig boundaries before implementation; done means framework/runtime validation, typed trainer options, and backward-compatible existing usage.

Written by the indexing model from the issue text.

Assessment

Tech stack
kubernetes, python, pytorch
Domain
backend-api-design, distributed-systems, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.