kubeflow / kubeflow/trainer

[BUG] Distributed environment variables are not injected when TrainJob spec.trainer is omitted in Torch, JAX, and XGBoost plugins

Open
#4,076 0 comments 0 reactions 0 assignees View on GitHub
kind/bug needs-triage
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.