[BUG] Distributed environment variables are not injected when TrainJob spec.trainer is omitted in Torch, JAX, and XGBoost plugins
- Dominant language
- Go
- Stars
- 2.2k
- Forks
- 1.1k
- Avg merge
- 3d 22h
- Merged PRs (30d)
- 39
Description
### What happened?
In the TrainJob API (`pkg/apis/trainer/v1alpha1/trainjob_types.go`), `spec.trainer` is marked as optional (`+optional`, `omitzero`). When a user submits a minimal `TrainJob` that relies entirely on a pre-configured `ClusterTrainingRuntime` / `TrainingRuntime` template (without specifying `spec.trainer`), the controller fails to inject critical distributed environment variables across multiple ML plugins (**Torch**, **JAX**, and **XGBoost**).
This causes distributed jobs using runtime defaults to fail because essential rendezvous and rank variables are missing.
---
### Root Cause Analysis
Across `pkg/runtime/framework/plugins/torch/torch.go`, `jax/jax.go`, and `xgboost/xgboost.go`, locating the `trainerContainer` is mistakenly guarded by `if trainJob.Spec.Trainer != nil`:
#### 1. Torch Plugin (`pkg/runtime/framework/plugins/torch/torch.go` lines 138–172)
```go
// Update envs for Info object.
var trainerContainer *runtime.Container
if trainJob.Spec.Trainer != nil {
if trainerContainer = info.FindContainerByPodSetAncestorContainerName(constants.AncestorTrainer, constants.Node); trainerContainer != nil {
apply.UpsertEnvVars(&trainerContainer.Env, apply.EnvVars(trainJob.Spec.Trainer.Env...)...)
}
}
...
// Inject PET_* envs into trainer main container (always).
if trainerContainer != nil {
apply.UpsertEnvVars(&trainerContainer.Env, petEnvs...)
apply.UpsertEnvVars(&trainerContainer.Env, masterEnvVars...)
}
```
When `trainJob.Spec.Trainer == nil`, `trainerContainer` remains `nil`. Consequently, `petEnvs` and `masterEnvVars` are never injected.
#### 2. JAX Plugin (`pkg/runtime/framework/plugins/jax/jax.go` lines 69–95)
```go
var trainerContainer *runtime.Container
if trainJob.Spec.Trainer != nil {
if trainerContainer = info.FindContainerByPodSetAncestorContainerName(constants.AncestorTrainer, constants.Node); trainerContainer != nil {
// Set JAX distributed environment variables
apply.UpsertEnvVars(&trainerContainer.Env,
*corev1ac.EnvVar().WithName("JAX_NUM_PROCESSES")...
```
If `trainJob.Spec.Trainer` is omitted, the entire JAX distributed environment block is skipped.
#### 3. XGBoost Plugin (`pkg/runtime/framework/plugins/xgboost/xgboost.go` lines 93–130)
```go
var trainerContainer *runtime.Container
if trainJob.Spec.Trainer != nil {
if trainerContainer = info.FindContainerByPodSetAncestorContainerName(
constants.AncestorTrainer, constants.Node,
); trainerContainer != nil {
// Inject DMLC_* environment variables
apply.UpsertEnvVars(&trainerContainer.Env,
*corev1ac.EnvVar().WithName(constants.XGBoostEnvTrackerURI)...
```
Similarly, all `DMLC_*` environment variables (`DMLC_TRACKER_URI`, `DMLC_TRACKER_PORT`, `DMLC_TASK_ID`, `DMLC_NUM_WORKER`) are skipped.
---
### What did you expect to happen?
The `trainerContainer` originates from the `runtime.Info` template spec (derived from the `ClusterTrainingRuntime`), not `trainJob.Spec.Trainer`.
The plugins should look up `trainerContainer` independently of `trainJob.Spec.Trainer`:
```go
trainerContainer := info.FindContainerByPodSetAncestorContainerName(constants.AncestorTrainer, constants.Node)
if trainerContainer != nil && trainJob.Spec.Trainer != nil {
apply.UpsertEnvVars(&trainerContainer.Env, apply.EnvVars(trainJob.Spec.Trainer.Env...)...)
}
if trainerContainer != nil {
// Inject framework-specific distributed envs (PET_*, JAX_*, DMLC_*)
}
```
---
### Environment
Kubernetes version:
```bash
$ kubectl version
```
Kubeflow Trainer version:
```bash
$ kubectl get pods -n kubeflow-system -l app.kubernetes.io/name=kubeflow-trainer -o jsonpath="{.items[*].spec.containers[*].image}"
```
Kubeflow Python SDK version:
```bash
$ pip show kubeflow
```
### Impacted by this bug?
Give it a 👍 We prioritize the issues with most 👍
Contributor guide
Research direction
Start with pkg/apis/trainer/v1alpha1/trainjob_types.go, then inspect the trainerContainer handling in pkg/runtime/framework/plugins/torch/torch.go, jax/jax.go, and xgboost/xgboost.go. Compare behavior for TrainJobs with and without spec.trainer while using runtime templates. Done means the Torch, JAX, and XGBoost distributed environment variables are injected when the trainer is defined only by the runtime template.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- go, kubernetes
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 78/100