[FEATURE] Integration with Kubeflow Trainer
- Dominant language
- No language data
- Stars
- 84
- Forks
- 149
- Avg merge
- 5d 15h
- Merged PRs (30d)
- 29
Description
### Certification
- [x] I certify I am an Epic Owner for Kubeflow Notebooks 2.0 and expected to create planning-related issues.
### Description
With per-workspace ServiceAccounts landed (see PR #1311), a Workspace can now bind namespace-scoped Roles to allow `kubectl exec` into TrainJob pods. The completion of this task aims for validation and documentation for the integration of this feature with Kubeflow Trainer, using an interactive MPI TrainJob as the concrete example.
There has been a previous discussion about this on Slack:
https://cloud-native.slack.com/archives/C073W562HFY/p1783469950269199
```
$ kubectl exec -it mpi-interactive-launcher-0-0-67q9c -- \
mpirun python -c "from mpi4py import MPI; import socket; c = MPI.COMM_WORLD; print(f'host={socket.gethostname()} | RANK: {c.Get_rank()}')"
Warning: Permanently added '[mpi-interactive-node-0-0.mpi-interactive]:2222' (ECDSA) to the list of known hosts.
Warning: Permanently added '[mpi-interactive-node-0-1.mpi-interactive]:2222' (ECDSA) to the list of known hosts.
Warning: Permanently added '[mpi-interactive-node-0-2.mpi-interactive]:2222' (ECDSA) to the list of known hosts.
host=mpi-interactive-launcher-0-0 | RANK: 0
host=mpi-interactive-node-0-2 | RANK: 3
host=mpi-interactive-node-0-1 | RANK: 2
host=mpi-interactive-node-0-0 | RANK: 1
```
Click to see sample manifests
```yaml
apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
name: mpi-runtime
labels:
trainer.kubeflow.org/framework: mpi
spec:
mlPolicy:
numNodes: 1
mpi:
numProcPerNode: 1
mpiImplementation: OpenMPI
sshAuthMountPath: /home/mpiuser/.ssh
runLauncherAsNode: true
template:
spec:
network:
publishNotReadyAddresses: true
successPolicy:
operator: All
targetReplicatedJobs:
- launcher
replicatedJobs:
- name: node
template:
spec:
template:
spec:
containers:
- name: node
image: ghcr.io/kubeflow/trainer/deepspeed-runtime
securityContext:
runAsUser: 1000
command:
- /usr/sbin/sshd
args:
- -De
- -f
- /home/mpiuser/.sshd_config
readinessProbe:
tcpSocket:
port: 2222
initialDelaySeconds: 5
- name: launcher
dependsOn:
- name: node
status: Ready
template:
metadata:
labels:
trainer.kubeflow.org/trainjob-ancestor-step: trainer
spec:
template:
spec:
containers:
- name: node
image: ghcr.io/kubeflow/trainer/deepspeed-runtime
securityContext:
runAsUser: 1000
---
apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainJob
metadata:
name: mpi-interactive
spec:
runtimeRef:
name: mpi-runtime
trainer:
numNodes: 4
command: ["sleep", "infinity"]
```
## Integration test (`kubeflow/community-distribution`):
Following the pattern of the existing pipeline test (`tests/workspaces_pipeline_run_test.sh` / `.github/workflows/workspaces_pipeline_run_test.yaml`), add a new test script and GitHub Actions workflow that:
1. Deploys a `ClusterTrainingRuntime` for MPI (using the manifest from Andrey's working example with `ghcr.io/kubeflow/trainer/deepspeed-runtime`)
2. Creates a `TrainJob` with `sleep infinity` to keep the MPI mesh alive
3. Creates a WorkspaceKind and Workspace configured with the appropriate `clusterRoles` to permit `exec` into TrainJob pods
4. From the Workspace pod, runs `mpirun python -c "from mpi4py import MPI; ..."` and verifies output from all ranks
The Kubeflow Trainer installation is already present in the community distribution (`applications/trainer/`) and is out of scope for this task.
## Operator guide (`kubeflow/website`):
Add a page under the [Workspaces operator guides](https://www.kubeflow.org/docs/components/workspaces/operator-guides/deployment-guide/) documenting how to:
- Set up the MPI `ClusterTrainingRuntime` for interactive use with Kubeflow Trainer
- Create a `TrainJob` for interactive use
- Configure a `WorkspaceKind` with `spec.podTemplate.serviceAccount.clusterRoles` to grant exec permissions on the mpi pod
- Connect from a Workspace and run `mpirun`
### Acceptance Criteria
- [ ] A test script (e.g. `tests/workspaces_trainer_test.sh`) exists in `kubeflow/community-distribution` that deploys the MPI ClusterTrainingRuntime, creates a TrainJob, creates a Workspace with the appropriate RBAC, and successfully runs `mpirun` from inside the Workspace pod
- [ ] A GitHub Actions workflow (e.g. `.github/workflows/workspaces_trainer_test.yaml`) runs the test on PRs that touch relevant paths (workspaces, trainer, istio, common infra)
- [ ] The workflow follows the structure of `workspaces_pipeline_run_test.yaml`
- [ ] An operator guide page is added to `kubeflow/website` under the Workspaces operator guides section, documenting the Kubeflow Trainer integration with Workspaces
- [ ] The guide includes the complete ClusterTrainingRuntime and TrainJob manifests, WorkspaceKind configuration with `clusterRoles`, and example `mpirun` command with expected output
Contributor guide
Assessment
This issue has not been assessed yet.