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

Issue:
PytorchJob restartPolicy: ExitCode leaves job in a "Restarting" state after a non-retryable error
### **Steps to reproduce:**
1. Set the PyTorchJob restartPolicy: ExitCode
2. Set backoffLimit > 1
3. Have a container exit with a non-zero exit code less than 128

### **Observed Behavior:**
The pod correctly goes into an error state without any retries, however the PyTorchJob itself stays in a `Restarting` state and never goes into a `Failed` state.

Logs:
```
INFO[0289] Pod: default.kfto-01-master-0 exited with code 1  job=default.kfto-01 replica-type=master uid=4acaf99e-c0e0-424e-b6fe-9cbb2e303f3b
INFO[0289] PyTorchJob=kfto-01, ReplicaType=Master expected=1, running=0, succeeded=0, failed=1, Replicas=1
2024-04-05T11:12:25-04:00	DEBUG	events	Pod: default.kfto-01-master-0 exited with code 1	{"type": "Normal", "object": {"kind":"PyTorchJob","namespace":"default","name":"kfto-01","uid":"4acaf99e-c0e0-424e-b6fe-9cbb2e303f3b","apiVersion":"kubeflow.org/v1","resourceVersion":"1800574"}, "reason": "ExitedWithCode"}
2024-04-05T11:12:25-04:00	DEBUG	events	PyTorchJob kfto-01 is restarting because 1 Master replica(s) failed.	{"type": "Warning", "object": {"kind":"PyTorchJob","namespace":"default","name":"kfto-01","uid":"4acaf99e-c0e0-424e-b6fe-9cbb2e303f3b","apiVersion":"kubeflow.org/v1","resourceVersion":"1800574"}, "reason": "PyTorchJobRestarting"}
```

### **Proposed solution:**
One possible solution is to have the pod controller logic [check here also for the exit codes of the pods](https://github.com/kubeflow/training-operator/blob/5b2c6c8943fe6a1f8803f268f71ca714316fa6bc/pkg/controller.v1/pytorch/pytorchjob_controller.go#L437). If at least one pod has a non-retryable exit code - fail the job. Thoughts?


Discussion / hints:
Thank you for creating this @kellyaa!

> Have a container exit with a non-zero exit code less than 128

Can we make sure that if Pod has exit code less than 128 it can't be restarted ?
What if user writes training container that can be restarted even with ExitCode = 1.


What do you think @kubeflow/wg-training-leads @kuizhiqing ?
> Can we make sure that if Pod has exit code less than 128 it can't be restarted ?

Correct, if the container's exist code is less than 128, the pod is not restarted.

> What if user writes training container that can be restarted even with ExitCode = 1.

If I set the failure policy to "OnFailure" and use an exit code of 1, the pod is restarted.

@andreyvelich  Did I understand your question properly?
Thank you for creating this issue.
This seems a bug, but we need to confirm that fixing this bug would not affect on existing other behaviors.
/kind bug
Sorry for the late reply, you are right @kellyaa, user can just use `OnFailure` restart policy in that case.
Ideally, our Jobs should support podFailurePolicy, like in Batch Jobs: https://kubernetes.io/docs/concepts/workloads/controllers/job/#pod-failure-policy, but we don't have it right now in our APIs.
Temporarily, we can just check the Pods exit codes in the controller.

@kellyaa Please can you submit PR to fix it, so we can include it in our next release ?
Hi @kellyaa @anishasthana, as we discussed on [the recent AutoML and Training WG Community call](https://youtu.be/wH0Jb4BUS5o), it would be nice to include this bug fix to the Training Operator 1.8 RC.0 which we are going to cut this Friday.
Please feel free to submit PR and we can review it.

Fix patch:
diff --git a/pkg/controller.v1/common/pod.go b/pkg/controller.v1/common/pod.go
index ef5fe1dcf9..fda1a2e528 100644
--- a/pkg/controller.v1/common/pod.go
+++ b/pkg/controller.v1/common/pod.go
@@ -39,6 +39,7 @@ import (
 	"k8s.io/apimachinery/pkg/runtime"
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
 	"k8s.io/client-go/tools/cache"
+	"k8s.io/klog/v2"
 )
 
 const (
@@ -346,23 +347,30 @@ func (jc *JobController) ReconcilePods(
 				}
 			}
 			// Check if the pod is retryable.
-			if pod.Status.Phase == v1.PodFailed &&
-				(spec.RestartPolicy == apiv1.RestartPolicyExitCode && trainutil.IsRetryableExitCode(exitCode) ||
-					spec.RestartPolicy == apiv1.RestartPolicyOnFailure ||
-					spec.RestartPolicy == apiv1.RestartPolicyAlways) {
+			if pod.Status.Phase == v1.PodFailed {
 				failedPodsCount.Inc()
-				logger.Infof("Need to restart the pod: %v.%v", pod.Namespace, pod.Name)
-				if err := jc.PodControl.DeletePod(pod.Namespace, pod.Name, runtimeObject); err != nil {
-					return err
+				if spec.RestartPolicy == apiv1.RestartPolicyExitCode && trainutil.IsRetryableExitCode(exitCode) ||
+					spec.RestartPolicy == apiv1.RestartPolicyOnFailure ||
+					spec.RestartPolicy == apiv1.RestartPolicyAlways {
+					logger.Infof("Need to restart the pod: %v.%v", pod.Namespace, pod.Name)
+					if err := jc.PodControl.DeletePod(pod.Namespace, pod.Name, runtimeObject); err != nil {
+						return err
+					}
+					// Deletion is expected
+					jc.Expectations.RaiseExpectations(expectationPodsKey, 0, 1)
+
+					msg := fmt.Sprintf("job %s is restarting because %s replica(s) failed.",
+						metaObject.GetName(), rType)
+					jc.Recorder.Event(runtimeObject, v1.EventTypeWarning, commonutil.NewReason(jobKind, commonutil.JobRestartingReason), msg)
+					commonutil.UpdateJobConditions(jobStatus, apiv1.JobRestarting, v1.ConditionTrue, commonutil.NewReason(jobKind, commonutil.JobRestartingReason), msg)
+					trainingoperatorcommon.RestartedJobsCounterInc(metaObject.GetNamespace(), jc.Controller.GetFrameworkName())
+				} else if spec.RestartPolicy == apiv1.RestartPolicyExitCode && !trainutil.IsRetryableExitCode(exitCode) {
+					logger.Infof("Pod %q has a non-retryable exit code. Failing job.", klog.KObj(pod))
+					msg := fmt.Sprintf("job %q is failing because %q replica(s) failed.",
+						metaObject.GetName(), rType)
+					jc.Recorder.Event(runtimeObject, v1.EventTypeWarning, commonutil.NewReason(jobKind, commonutil.JobFailedReason), msg)
+					commonutil.UpdateJobConditions(jobStatus, apiv1.JobFailed, v1.ConditionTrue, commonutil.NewReason(jobKind, commonutil.JobFailedReason), msg)
 				}
-				// Deletion is expected
-				jc.Expectations.RaiseExpectations(expectationPodsKey, 0, 1)
-
-				msg := fmt.Sprintf("job %s is restarting because %s replica(s) failed.",
-					metaObject.GetName(), rType)
-				jc.Recorder.Event(runtimeObject, v1.EventTypeWarning, commonutil.NewReason(jobKind, commonutil.JobRestartingReason), msg)
-				commonutil.UpdateJobConditions(jobStatus, apiv1.JobRestarting, v1.ConditionTrue, commonutil.NewReason(jobKind, commonutil.JobRestartingReason), msg)
-				trainingoperatorcommon.RestartedJobsCounterInc(metaObject.GetNamespace(), jc.Controller.GetFrameworkName())
 			}
 
 			updateJobReplicaStatuses(jobStatus, rType, pod)


Test patch:
diff --git a/pkg/controller.v1/tensorflow/pod_test.go b/pkg/controller.v1/tensorflow/pod_test.go
index 188e0fcab7..1663e721b9 100644
--- a/pkg/controller.v1/tensorflow/pod_test.go
+++ b/pkg/controller.v1/tensorflow/pod_test.go
@@ -21,6 +21,7 @@ import (
 
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
+	. "github.com/onsi/gomega/gstruct"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -274,6 +275,51 @@ var _ = Describe("TFJob controller", func() {
 		})
 	})
 
+	Context("Test Unretryable Exit Code", func() {
+		It("should set the job status to Failed", func() {
+			By("Creating TFJob \"test-noretry-exit-code\" with 1 worker only")
+			ctx := context.Background()
+
+			tfJob := tftestutil.NewTFJob(1, 0)
+			tfJob.SetName("test-noretry-exit-code")
+			tfJob.SetUID(uuid.NewUUID())
+			tfJob.Spec.TFReplicaSpecs[kubeflowv1.TFJobReplicaTypeWorker].RestartPolicy = kubeflowv1.RestartPolicyExitCode
+			Expect(testK8sClient.Create(ctx, tfJob)).Should(Succeed())
+
+			_ = reconciler.ReconcileJobs(tfJob, tfJob.Spec.TFReplicaSpecs, tfJob.Status, &tfJob.Spec.RunPolicy)
+
+			created := &corev1.Pod{}
+			key := types.NamespacedName{Namespace: metav1.NamespaceDefault, Name: "test-noretry-exit-code-worker-0"}
+			Expect(testK8sClient.Get(ctx, key, created)).Should(Succeed())
+			created.Status.Phase = corev1.PodFailed
+			created.Status.ContainerStatuses = append(created.Status.ContainerStatuses, corev1.ContainerStatus{
+				Name: kubeflowv1.TFJobDefaultContainerName,
+				State: corev1.ContainerState{
+					Terminated: &corev1.ContainerStateTerminated{
+						ExitCode: 1,
+					},
+				},
+			})
+			Expect(testK8sClient.Status().Update(ctx, created)).Should(Succeed())
+
+			_ = reconciler.ReconcileJobs(tfJob, tfJob.Spec.TFReplicaSpecs, tfJob.Status, &tfJob.Spec.RunPolicy)
+
+			Eventually(func(g Gomega) {
+				updatedJob := &kubeflowv1.TFJob{}
+				g.Expect(testK8sClient.Get(ctx, types.NamespacedName{Name: tfJob.GetName(), Namespace: metav1.NamespaceDefault}, updatedJob)).Should(Succeed())
+				expectedCondition := kubeflowv1.JobCondition{ // Make sure this type is correct
+					Type:   "Failed",
+					Status: "True",
+				}
+				g.Expect(updatedJob.Status.Conditions).Should(ConsistOf(MatchFields(IgnoreExtras, Fields{
+					"Type":   Equal(expectedCondition.Type),
+					"Status": Equal(expectedCondition.Status),
+				})), "TFJob should be in Failed state")
+			}, testutil.Timeout, testutil.Interval).Should(Succeed())
+
+		})
+	})
+
 	Describe("Test Scale Down", func() {
 		It("should delete redundant Pods", func() {
 			ctx := context.Background()