Repository: kubeflow/training-operator (language: python)

Issue:
Training Operator crashes when submitting PyTorchJob with elasticPolicy but without worker template defined
### What happened?

When submitting a PyTorchJob with an elasticPolicy set but only a master template (no worker template defined), the Training Operator crashes with a nil pointer exception that seems to be coming from [this line](https://github.com/kubeflow/training-operator/blob/master/pkg/apis/kubeflow.org/v1/pytorch_defaults.go#L49).

### What did you expect to happen?

We expect the Training Operator to not crash. We are interested in contributing a fix for this, but are wondering if it would make more sense to not allow users to set an elasticPolicy if they only have a master replica or to update the logic to allow this behavior. If we block the behavior, we could add this to the validating webhook so the user is informed of the error. If we allow the behavior, since [the default c10d store is set on worker 0](https://github.com/kubeflow/training-operator/blob/master/pkg/controller.v1/pytorch/elastic.go#L150), we would need to update this piece as well. What are your thoughts?

### Environment

Kubernetes version:
`v1.29.8`

Training Operator version:
`v1-855e096`, also tested a local build using the latest on master

Training Operator Python SDK version:
N/A

### Impacted by this bug?

Give it a 👍 We prioritize the issues with most 👍


Discussion / hints:
One should come with worker section in elastic mode.

Feel free to add validation logic in this case. 

Looking forward to your contribution.

Fix patch:
diff --git a/pkg/webhooks/pytorch/pytorchjob_webhook.go b/pkg/webhooks/pytorch/pytorchjob_webhook.go
index 2459815935..bc2d9fc527 100644
--- a/pkg/webhooks/pytorch/pytorchjob_webhook.go
+++ b/pkg/webhooks/pytorch/pytorchjob_webhook.go
@@ -93,13 +93,19 @@ func validatePyTorchJob(oldJob, newJob *trainingoperator.PyTorchJob) (admission.
 func validateSpec(spec trainingoperator.PyTorchJobSpec) (admission.Warnings, field.ErrorList) {
 	var allErrs field.ErrorList
 	var warnings admission.Warnings
-
-	if spec.ElasticPolicy != nil && spec.ElasticPolicy.NProcPerNode != nil {
-		elasticNProcPerNodePath := specPath.Child("elasticPolicy").Child("nProcPerNode")
-		nprocPerNodePath := specPath.Child("nprocPerNode")
-		warnings = append(warnings, fmt.Sprintf("%s is deprecated, use %s instead", elasticNProcPerNodePath.String(), nprocPerNodePath.String()))
-		if spec.NprocPerNode != nil {
-			allErrs = append(allErrs, field.Forbidden(elasticNProcPerNodePath, fmt.Sprintf("must not be used with %s", nprocPerNodePath)))
+	if spec.ElasticPolicy != nil {
+		_, ok := spec.PyTorchReplicaSpecs[trainingoperator.PyTorchJobReplicaTypeWorker]
+		workerPath := pytorchReplicaSpecPath.Key(string(trainingoperator.PyTorchJobReplicaTypeWorker))
+		if !ok {
+			allErrs = append(allErrs, field.Required(workerPath, "must be configured if elastic policy is used"))
+		}
+		if spec.ElasticPolicy.NProcPerNode != nil {
+			elasticNProcPerNodePath := specPath.Child("elasticPolicy").Child("nProcPerNode")
+			nprocPerNodePath := specPath.Child("nprocPerNode")
+			warnings = append(warnings, fmt.Sprintf("%s is deprecated, use %s instead", elasticNProcPerNodePath.String(), nprocPerNodePath.String()))
+			if spec.NprocPerNode != nil {
+				allErrs = append(allErrs, field.Forbidden(elasticNProcPerNodePath, fmt.Sprintf("must not be used with %s", nprocPerNodePath)))
+			}
 		}
 	}
 	allErrs = append(allErrs, validatePyTorchReplicaSpecs(spec.PyTorchReplicaSpecs)...)
@@ -147,6 +153,8 @@ func validatePyTorchReplicaSpecs(rSpecs map[trainingoperator.ReplicaType]*traini
 			if rSpec.Replicas == nil || int(*rSpec.Replicas) != 1 {
 				allErrs = append(allErrs, field.Forbidden(rolePath.Child("replicas"), "must be 1"))
 			}
+		} else if rSpec.Replicas != nil && int(*rSpec.Replicas) < 1 {
+			allErrs = append(allErrs, field.Forbidden(rolePath.Child("replicas"), "must be at least 1"))
 		}
 	}
 	return allErrs


Test patch:
diff --git a/pkg/webhooks/pytorch/pytorchjob_webhook_test.go b/pkg/webhooks/pytorch/pytorchjob_webhook_test.go
index 7757e36b3e..e1b7d3f69e 100644
--- a/pkg/webhooks/pytorch/pytorchjob_webhook_test.go
+++ b/pkg/webhooks/pytorch/pytorchjob_webhook_test.go
@@ -88,6 +88,9 @@ func TestValidateV1PyTorchJob(t *testing.T) {
 					RunPolicy: trainingoperator.RunPolicy{
 						ManagedBy: ptr.To(trainingoperator.KubeflowJobsController),
 					},
+					ElasticPolicy: &trainingoperator.ElasticPolicy{
+						RDZVBackend: ptr.To(trainingoperator.BackendC10D),
+					},
 					PyTorchReplicaSpecs: validPyTorchReplicaSpecs,
 				},
 			},
@@ -247,6 +250,19 @@ func TestValidateV1PyTorchJob(t *testing.T) {
 								},
 							},
 						},
+						trainingoperator.PyTorchJobReplicaTypeWorker: {
+							Replicas: ptr.To[int32](1),
+							Template: corev1.PodTemplateSpec{
+								Spec: corev1.PodSpec{
+									Containers: []corev1.Container{
+										{
+											Name:  "pytorch",
+											Image: "gcr.io/kubeflow-ci/pytorch-dist-mnist_test:1.0",
+										},
+									},
+								},
+							},
+						},
 					},
 				},
 			},
@@ -279,6 +295,19 @@ func TestValidateV1PyTorchJob(t *testing.T) {
 								},
 							},
 						},
+						trainingoperator.PyTorchJobReplicaTypeWorker: {
+							Replicas: ptr.To[int32](1),
+							Template: corev1.PodTemplateSpec{
+								Spec: corev1.PodSpec{
+									Containers: []corev1.Container{
+										{
+											Name:  "pytorch",
+											Image: "gcr.io/kubeflow-ci/pytorch-dist-mnist_test:1.0",
+										},
+									},
+								},
+							},
+						},
 					},
 				},
 			},
@@ -335,6 +364,62 @@ func TestValidateV1PyTorchJob(t *testing.T) {
 				field.Invalid(field.NewPath("spec", "runPolicy", "managedBy"), trainingoperator.MultiKueueController, apivalidation.FieldImmutableErrorMsg),
 			},
 		},
+		"attempt to configure elasticPolicy when no worker is configured": {
+			pytorchJob: &trainingoperator.PyTorchJob{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "test",
+				},
+				Spec: trainingoperator.PyTorchJobSpec{
+					ElasticPolicy: &trainingoperator.ElasticPolicy{},
+					PyTorchReplicaSpecs: map[trainingoperator.ReplicaType]*trainingoperator.ReplicaSpec{
+						trainingoperator.PyTorchJobReplicaTypeMaster: {
+							Replicas: ptr.To[int32](1),
+							Template: corev1.PodTemplateSpec{
+								Spec: corev1.PodSpec{
+									Containers: []corev1.Container{
+										{
+											Name:  "pytorch",
+											Image: "gcr.io/kubeflow-ci/pytorch-dist-mnist_test:1.0",
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			wantErr: field.ErrorList{
+				field.Required(pytorchReplicaSpecPath.Key(string(trainingoperator.PyTorchJobReplicaTypeWorker)), ""),
+			},
+		},
+		"attempt to configure worker with 0 replicas": {
+			pytorchJob: &trainingoperator.PyTorchJob{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "test",
+				},
+				Spec: trainingoperator.PyTorchJobSpec{
+					ElasticPolicy: &trainingoperator.ElasticPolicy{},
+					PyTorchReplicaSpecs: map[trainingoperator.ReplicaType]*trainingoperator.ReplicaSpec{
+						trainingoperator.PyTorchJobReplicaTypeWorker: {
+							Replicas: ptr.To[int32](0),
+							Template: corev1.PodTemplateSpec{
+								Spec: corev1.PodSpec{
+									Containers: []corev1.Container{
+										{
+											Name:  "pytorch",
+											Image: "gcr.io/kubeflow-ci/pytorch-dist-mnist_test:1.0",
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			wantErr: field.ErrorList{
+				field.Forbidden(pytorchReplicaSpecPath.Key(string(trainingoperator.PyTorchJobReplicaTypeWorker)).Child("replicas"), ""),
+			},
+		},
 	}
 
 	for name, tc := range testCases {