Skip to content

TrainJob accepts zero and negative numNodes / numProcPerNode (no lower-bound validation) #3805

Description

@sanskar-singh-2403

What happened?

What happened?

TrainJob.Spec.Trainer.NumNodes and TrainJob.Spec.Trainer.NumProcPerNode have no lower-bound validation. A TrainJob with numNodes: 0, numNodes: -3, or numProcPerNode: -5 is accepted at every layer (CRD schema, validating webhook, and runtime plugins), and the invalid values propagate into the pod spec.

1. The CRD schema has no minimum. In pkg/apis/trainer/v1alpha1/trainjob_types.go (lines 275-288) both fields are declared without a +kubebuilder:validation:Minimum marker:

// numNodes is the number of training nodes.
// +optional
NumNodes *int32 `json:"numNodes,omitempty"`

// numProcPerNode is the number of processes/workers/slots on every training node.
// For the MPI runtime only int value can be set to represent number of slots per node.
// For the Torch runtime the value defaults to `auto` and can be overridden with an int.
// +optional
NumProcPerNode *int32 `json:"numProcPerNode,omitempty"`

The generated CRD (manifests/base/crds/trainer.kubeflow.org_trainjobs.yaml) therefore has type: integer with no minimum, so the API server accepts any int32.

2. The validating webhook does not check them either. pkg/webhooks/trainjob_webhook.go ValidateCreate/ValidateUpdate only delegate to runtime.ValidateObjects, and none of the runtime plugins (torch.go, mpi.go, etc.) bound-check these fields.

3. The bad values propagate into the pod spec. For the Torch runtime:

  • pkg/runtime/framework/plugins/torch/torch.go:122-123 does numProcPerNode = intstr.FromInt32(*trainJob.Spec.Trainer.NumProcPerNode) and renders it directly into the container env PET_NPROC_PER_NODE (torch.go:150-151). A negative value becomes PET_NPROC_PER_NODE=-5, so torchrun fails at pod runtime instead of at admission.
  • torch.go:118 (and plainml/plainml.go:58, jax/jax.go:66, xgboost/xgboost.go:89) do *trainerPS.Count = *trainJob.Spec.Trainer.NumNodes, overwriting the otherwise-guaranteed-positive PodSet count (built as ptr.To(max(count, 1)) in runtime.go:147) with 0 or a negative number, which becomes the JobSet completions/parallelism.

The runtime-side sibling field is already guarded, which shows >= 1 is the intended contract:

// pkg/apis/trainer/v1alpha1/trainingruntime_types.go:305
// +kubebuilder:validation:XValidation:rule="self >= 1",message="NumProcPerNode in fluxPolicy must be >= 1"

Only the user-facing TrainJob-level overrides are missing the guard.

Reproduction

Apply a TrainJob that sets non-positive values against any Torch runtime:

apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainJob
metadata:
  name: bad-numproc
spec:
  runtimeRef:
    name: torch-distributed
    kind: ClusterTrainingRuntime
  trainer:
    numNodes: 0
    numProcPerNode: -5

kubectl apply succeeds. The trainer container receives PET_NPROC_PER_NODE=-5 and the trainer PodSet count is set to 0.

I also confirmed this at the unit level against master by driving the Torch plugin directly (Torch.Validate returns 0 errors, then Torch.EnforceMLPolicy renders the bad env and count):

Validate() returned 0 errors: []
Resulting PET_NPROC_PER_NODE="-5", trainer PodSet Count=0

Proposed fix

Add lower-bound markers on both fields in pkg/apis/trainer/v1alpha1/trainjob_types.go and regenerate the CRDs with make manifests:

// numNodes is the number of training nodes.
// +kubebuilder:validation:Minimum=1
// +optional
NumNodes *int32 `json:"numNodes,omitempty"`

// numProcPerNode is the number of processes/workers/slots on every training node.
// +kubebuilder:validation:Minimum=1
// +optional
NumProcPerNode *int32 `json:"numProcPerNode,omitempty"`

I verified this marker generates minimum: 1 into the TrainJob CRD for both fields, so the API server rejects non-positive values at admission. Fixing it at the CRD level covers all runtimes at once (Torch, PlainML, JAX, XGBoost, MPI) since they read the same two fields. Note: numProcPerNode defaults to auto for Torch by being left unset (nil), which Minimum=1 does not affect; it only bounds the explicit int override.

I am happy to open a PR with the marker change, the make manifests / make generate regeneration, and unit tests covering the zero and negative cases.

What did you expect to happen?

What did you expect to happen?

Admission should reject numNodes and numProcPerNode values below 1 with a clear validation error, instead of deferring the failure to pod runtime (torchrun crash) or silently producing a JobSet with a zero/negative pod count.

Environment

Kubernetes version:

$ kubectl version

Kubeflow Trainer version:

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

Kubeflow Python SDK version:

$ pip show kubeflow

Impacted by this bug?

Give it a 👍 We prioritize the issues with most 👍

Metadata

Metadata

Labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions