kubeflow / kubeflow/trainer

JAX plugin Validate() is a no-op - missing reserved env validation

Open
#3,669 1 comment 0 reactions 1 assignee Claimed by @AdeshDeshmukh View on GitHub
Dominant language
Go
Stars
2.2k
Forks
1.1k
Avg merge
3d 22h
Merged PRs (30d)
39

Description

**What happened**:
The JAX plugin's `Validate()` method at `pkg/runtime/framework/plugins/jax/jax.go:52-54` is currently a no-op (`return nil, nil`).

Because of this, users can manually set `JAX_NUM_PROCESSES`, `JAX_PROCESS_ID`, or `JAX_COORDINATOR_ADDRESS` in their TrainJob's `trainer.env`. When the controller runs `EnforceMLPolicy`, it silently overwrites these user-provided values via `apply.UpsertEnvVars`. This leads to hard-to-debug training failures because the user's configuration is ignored without any warning or error.

All other ML policy plugins (Torch, MPI, XGBoost) validate and reject TrainJobs that set reserved environment variables. The JAX plugin is missing this validation, and the `constants.JAXReservedEnvNames` set was never declared in `pkg/constants/constants.go`.

**What did you expect to happen**:
The JAX plugin should reject TrainJobs that set reserved JAX environment variables, just like the Torch plugin rejects `PET_NNODES`, MPI rejects `OMPI_MCA_*`, etc. This ensures users get an immediate, clear error during API admission rather than silent misconfigurations at runtime.

**Environment**:
Kubernetes version: N/A (codebase issue)
Kubeflow Trainer version: master (latest)
Kubeflow Python SDK version: N/A

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.