kubeflow / kubeflow/trainer

MPI: `spec.trainer.image` and `spec.trainer.command` cannot both be used once the worker job carries the trainer ancestor label

Open
#4,030 6 comments 1 reaction 1 assignee Claimed by @Harsh-2005d View on GitHub
area/mpi kind/bug
Dominant language
Go
Stars
2.2k
Forks
1.1k
Avg merge
3d 22h
Merged PRs (30d)
39

Description

### What happened?

Following up on #2825, where the suggested workaround for changing the MPI worker image is
to put `trainer.kubeflow.org/trainjob-ancestor-step: trainer` on the `node` replicatedJob
and not set `command` on the TrainJob.

That works. Tested on v2.1.0 with a two-replicatedJob OpenMPI runtime
(`runLauncherAsNode: true`):

```
# label on the node rJob, TrainJob sets image only
launcher image: my.registry/mpi-app:test command: ["sh","-c"]
node image: my.registry/mpi-app:test command: ["/usr/sbin/sshd"]

Complete=True Workers: 6
Rank 0,1 on launcher-0-0 | Rank 2,3 on node-0-0 | Rank 4,5 on node-0-1
```

The problem is what happens around it.

**Once the `node` rJob carries the label, setting `command` also lands on the workers and
replaces the SSH server:**

```
launcher command: ["sh","-c","echo I-AM-THE-TRAINER-COMMAND; sleep 3600"]
args: []
node command: ["sh","-c","echo I-AM-THE-TRAINER-COMMAND; sleep 3600"]
args: ["-De","-f","/home/mpiuser/.sshd_config"]
```

The workers stop running `sshd`, so the launcher has nothing to connect to. And because
`command` and `args` are applied independently, each only when non-nil, the worker ends up
with the TrainJob's `command` next to the runtime's leftover `sshd` args.

`args` behaves the same way. Setting only `args`, with no `command`:

```
node command: ["/usr/sbin/sshd"] args: ["I-AM-TRAINER-ARGS"]
launcher command: ["sh","-c"] args: ["I-AM-TRAINER-ARGS"]
```

so the worker runs `/usr/sbin/sshd I-AM-TRAINER-ARGS` and the daemon does not come up.

Without the `trainer.kubeflow.org/trainjob-ancestor-step: trainer` label, `spec.trainer.image` changes only
the launcher, with no warning:

```
$ kubectl get jobset imgtest -o jsonpath='{range .spec.replicatedJobs[*]}{.name}{" -> "}{.template.spec.template.spec.containers[*].image}{"\n"}{end}'
launcher -> busybox:1.36
node -> mpioperator/mpi-pi:openmpi
```

Both MPI runtimes in the repo keep launcher and worker on the same image
(`deepspeed-distributed`, `mlx-distributed`), which matches how MPI works. `mpirun`/`prte`
spawns daemons on the workers over SSH and they need to be compatible. A version skew
between launcher and workers does not necessarily fail loudly: in #4028 an Open MPI 5
mismatch produces a job that exits 0 with every rank on a single node.

**On the "why not create runtimes per environment" question from #2825.** You asked the
reporter there what the concrete challenge was, and I don't think the thread ever answered
it, so here is our version. We have prebuilt images that differ by dependency versions,
built ahead of time because we cannot install packages at run time. Adopting the workaround
means the entrypoint has to come from somewhere other than the TrainJob, and both options
have a cost:

- Put the command in the runtime. The runtime then pins the entrypoint, so we need a
runtime per application rather than per image. There are more applications than
dependency sets, so that is more runtimes than we started with.
- Bake the entrypoint into each image, as suggested in #2825. This does give one runtime
for many images, but the command is then fixed at build time. Changing an argument means
rebuilding and repushing, and the only per-job knob left is `env`, since `args` is gated
the same way as `command`.

### What did you expect to happen?

I expected `spec.trainer.image` and `spec.trainer.command` to be usable together on an MPI
runtime, either with the image applying to launcher and workers alike, or with an error explaining
why it cannot.

My question is whether the coupling is intentional: is there a reason `spec.trainer.image`
has to be applied under the same condition as `spec.trainer.command` / `args`, or could
`image` be applied independently?

If the coupling is not deliberate, applying `spec.trainer.image` to every container named
`node` across replicatedJobs, leaving `command`/`args` gated on the ancestor label as they
are now, would cover this without a new API field. Container name looks like a sufficient
scope: every runtime names the training container `node`, including both the MPI launcher
and worker, while initializer containers are named `dataset-initializer` and
`model-initializer`, so the torchtune runtimes are unaffected.

That would be a behaviour change for anyone relying on `trainer.image` reaching only the
launcher. I could not think of a case where different launcher and worker images would be
wanted for one MPI job, but you may know of one. In any case, I would assume this to be
the corner case which requires an ad-hoc runtime, rather the more general use case of
using the same container image for both launcher and workers/nodes.

If changing the semantics is not acceptable, an alternative would be to warn or reject at
admission when a TrainJob would leave differing images across an MPI runtime's
replicatedJobs, and when `command`/`args` are set against a runtime whose `node` rJob
carries the trainer ancestor label. That would at least make both cases visible instead of
silent.

Happy to open a PR for whichever you prefer.

### Environment

Kubernetes version:
```bash
$ kubectl version
Client Version: v1.35.1
Kustomize Version: v5.7.1
Server Version: v1.30.5
```

Kubeflow Trainer version:
```bash
$ kubectl get pods -n kubeflow-system -l app.kubernetes.io/name=kubeflow-trainer -o jsonpath="{.items[*].spec.containers[*].image}"

```
(returns nothing on v2.1.0, because the pod label is `app.kubernetes.io/name=trainer`)
```bash
$ kubectl get pods -n kubeflow-system -l app.kubernetes.io/name=trainer -o jsonpath="{.items[*].spec.containers[*].image}"
ghcr.io/kubeflow/trainer/trainer-controller-manager:v2.1.0
```
JobSet v0.10.1. The gating code in `builder.go` is unchanged from v2.0.0 through master.

Kubeflow Python SDK version: not used, TrainJobs are applied directly as YAML.

### Impacted by this bug?

Give it a 👍 We prioritize the issues with most 👍

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.