Implement Suspend semantics for MPIJob

# Conflicts:
#	pkg/apis/kubeflow/v2beta1/types.go
#	pkg/controller/mpi_job_controller.go
#	pkg/controller/mpi_job_controller_status.go
#	pkg/controller/mpi_job_controller_test.go
#	test/integration/mpi_job_controller_test.go

Author: mimowo

crd/kubeflow.org_mpijobs.yaml

@@ -7866,6 +7866,18 @@ spec:
                         format: int32
                         type: integer
                     type: object
+                  suspend:
+                    default: false
+                    description: "suspend specifies whether the MPIJob controller
+                      should create Pods or not. If a MPIJob is created with suspend
+                      set to true, no Pods are created by the MPIJob controller. If
+                      a MPIJob is suspended after creation (i.e. the flag goes from
+                      false to true), the MPIJob controller will delete all active
+                      Pods and PodGroups associated with this MPIJob. Also, it will
+                      suspend the Launcher Job. Users must design their workload to
+                      gracefully handle this. Suspending a Job will reset the StartTime
+                      field of the MPIJob. \n Defaults to false."
+                    type: boolean
                   ttlSecondsAfterFinished:
                     description: TTLSecondsAfterFinished is the TTL to clean up jobs.
                       It may take extra ReconcilePeriod seconds for the cleanup, since

go.mod

@@ -17,6 +17,7 @@ require (
 	k8s.io/klog v1.0.0
 	k8s.io/kube-openapi v0.0.0-20230109183929-3758b55a6596
 	k8s.io/sample-controller v0.25.6
+	k8s.io/utils v0.0.0-20220728103510-ee6ede2d64ed
 	sigs.k8s.io/controller-runtime v0.13.1
 	volcano.sh/apis v1.7.0
 )
@@ -73,7 +74,6 @@ require (
 	k8s.io/component-base v0.25.6 // indirect
 	k8s.io/gengo v0.0.0-20211129171323-c02415ce4185 // indirect
 	k8s.io/klog/v2 v2.70.1 // indirect
-	k8s.io/utils v0.0.0-20220728103510-ee6ede2d64ed // indirect
 	sigs.k8s.io/json v0.0.0-20220713155537-f223a00ba0e2 // indirect
 	sigs.k8s.io/structured-merge-diff/v4 v4.2.3 // indirect
 	sigs.k8s.io/yaml v1.3.0 // indirect

manifests/base/crd.yaml

@@ -49,6 +49,8 @@ spec:
                     type: integer
                     minimum: 0
                     description: "Specifies the number of retries before marking the launcher Job as failed. Defaults to 6."
+                  suspend:
+                    type: boolean
               sshAuthMountPath:
                 type: string
               mpiImplementation:
@@ -94,7 +96,7 @@ spec:
                   properties:
                     type:
                       type: string
-                      enum: ["Created", "Running", "Restarting", "Succeeded", "Failed"]
+                      enum: ["Created", "Running", "Restarting", "Succeeded", "Suspended", "Failed"]
                     status:
                       type: string
                       enum: ["True", "False", "Unknown"]

pkg/apis/kubeflow/v2beta1/openapi_generated.go

@@ -399,6 +399,10 @@
           "description": "SchedulingPolicy defines the policy related to scheduling, e.g. gang-scheduling",
           "$ref": "#/definitions/v2beta1.SchedulingPolicy"
         },
+        "suspend": {
+          "description": "suspend specifies whether the MPIJob controller should create Pods or not. If a MPIJob is created with suspend set to true, no Pods are created by the MPIJob controller. If a MPIJob is suspended after creation (i.e. the flag goes from false to true), the MPIJob controller will delete all active Pods and PodGroups associated with this MPIJob. Also, it will suspend the Launcher Job. Users must design their workload to gracefully handle this. Suspending a Job will reset the StartTime field of the MPIJob.\n\nDefaults to false.",
+          "type": "boolean"
+        },
         "ttlSecondsAfterFinished": {
           "description": "TTLSecondsAfterFinished is the TTL to clean up jobs. It may take extra ReconcilePeriod seconds for the cleanup, since reconcile gets called periodically. Default to infinite.",
           "type": "integer",

pkg/apis/kubeflow/v2beta1/swagger.json

@@ -84,6 +84,18 @@ type RunPolicy struct {
 	// SchedulingPolicy defines the policy related to scheduling, e.g. gang-scheduling
 	// +optional
 	SchedulingPolicy *SchedulingPolicy `json:"schedulingPolicy,omitempty"`
+
+	// suspend specifies whether the MPIJob controller should create Pods or not.
+	// If a MPIJob is created with suspend set to true, no Pods are created by
+	// the MPIJob controller. If a MPIJob is suspended after creation (i.e. the
+	// flag goes from false to true), the MPIJob controller will delete all
+	// active Pods and PodGroups associated with this MPIJob. Also, it will suspend the
+	// Launcher Job. Users must design their workload to gracefully handle this.
+	// Suspending a Job will reset the StartTime field of the MPIJob.
+	//
+	// Defaults to false.
+	// +kubebuilder:default:=false
+	Suspend *bool `json:"suspend,omitempty"`
 }
 
 type MPIJobSpec struct {
@@ -239,6 +251,9 @@ const (
 	// The training is complete without error.
 	JobSucceeded JobConditionType = "Succeeded"
 
+	// JobSuspended means the job has been suspended.
+	JobSuspended JobConditionType = "Suspended"
+
 	// JobFailed means one or more sub-resources (e.g. services/pods) of this job
 	// reached phase failed with no restarting.
 	// The training has failed its execution.

pkg/apis/kubeflow/v2beta1/types.go

@@ -33,6 +33,7 @@ import (
 	"golang.org/x/crypto/ssh"
 	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
+	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/equality"
 	"k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -50,6 +51,7 @@ import (
 	"k8s.io/client-go/tools/record"
 	"k8s.io/client-go/util/workqueue"
 	"k8s.io/klog"
+	"k8s.io/utils/pointer"
 	podgroupv1beta1 "volcano.sh/apis/pkg/apis/scheduling/v1beta1"
 	volcanoclient "volcano.sh/apis/pkg/client/clientset/versioned"
 	podgroupsinformer "volcano.sh/apis/pkg/client/informers/externalversions/scheduling/v1beta1"
@@ -493,7 +495,7 @@ func (c *MPIJobController) syncHandler(key string) error {
 
 	if len(mpiJob.Status.Conditions) == 0 {
 		msg := fmt.Sprintf("MPIJob %s/%s is created.", mpiJob.Namespace, mpiJob.Name)
-		updateMPIJobConditions(mpiJob, kubeflow.JobCreated, mpiJobCreatedReason, msg)
+		updateMPIJobConditions(mpiJob, kubeflow.JobCreated, v1.ConditionTrue, mpiJobCreatedReason, msg)
 		c.recorder.Event(mpiJob, corev1.EventTypeNormal, "MPIJobCreated", msg)
 		mpiJobsCreatedCount.Inc()
 	}
@@ -503,24 +505,13 @@ func (c *MPIJobController) syncHandler(key string) error {
 	// cleanup and stop retrying the MPIJob.
 	if isFinished(mpiJob.Status) && mpiJob.Status.CompletionTime != nil {
 		if isCleanUpPods(mpiJob.Spec.RunPolicy.CleanPodPolicy) {
-			// set worker StatefulSet Replicas to 0.
-			if err := c.deleteWorkerPods(mpiJob); err != nil {
-				return err
-			}
-			initializeMPIJobStatuses(mpiJob, kubeflow.MPIReplicaTypeWorker)
-			if c.gangSchedulerName != "" {
-				if err := c.deletePodGroups(mpiJob); err != nil {
-					return err
-				}
-			}
-			mpiJob.Status.ReplicaStatuses[kubeflow.MPIReplicaTypeWorker].Active = 0
-			return c.updateStatusHandler(mpiJob)
+			return cleanUpWorkerPods(mpiJob, c)
 		}
 		return nil
 	}
 
 	// first set StartTime.
-	if mpiJob.Status.StartTime == nil {
+	if mpiJob.Status.StartTime == nil && !isMPIJobSuspended(mpiJob) {
 		now := metav1.Now()
 		mpiJob.Status.StartTime = &now
 	}
@@ -549,17 +540,18 @@ func (c *MPIJobController) syncHandler(key string) error {
 			return fmt.Errorf("creating SSH auth secret: %w", err)
 		}
 
-		// Get the PodGroup for this MPIJob
-		if c.gangSchedulerName != "" {
-			if podgroup, err := c.getOrCreatePodGroups(mpiJob, workerReplicas(mpiJob)+1); podgroup == nil || err != nil {
+		if !isMPIJobSuspended(mpiJob) {
+			// Get the PodGroup for this MPIJob
+			if c.gangSchedulerName != "" {
+				if podgroup, err := c.getOrCreatePodGroups(mpiJob, workerReplicas(mpiJob)+1); podgroup == nil || err != nil {
+					return err
+				}
+			}
+			worker, err = c.getOrCreateWorker(mpiJob)
+			if err != nil {
 				return err
 			}
 		}
-
-		worker, err = c.getOrCreateWorker(mpiJob)
-		if err != nil {
-			return err
-		}
 		if mpiJob.Spec.MPIImplementation == kubeflow.MPIImplementationIntel {
 			// The Intel implementation requires workers to communicate with the
 			// launcher through its hostname. For that, we create a Service which
@@ -585,9 +577,40 @@ func (c *MPIJobController) syncHandler(key string) error {
 		return err
 	}
 
+	if launcher != nil {
+		if isMPIJobSuspended(mpiJob) != isJobSuspended(launcher) {
+			// align the suspension state of launcher with the MPIJob
+			launcher.Spec.Suspend = pointer.Bool(isMPIJobSuspended(mpiJob))
+			if _, err := c.kubeClient.BatchV1().Jobs(namespace).Update(context.TODO(), launcher, metav1.UpdateOptions{}); err != nil {
+				return err
+			}
+		}
+	}
+
+	// cleanup the running worker pods if the MPI job is suspended
+	if isMPIJobSuspended(mpiJob) {
+		if err := cleanUpWorkerPods(mpiJob, c); err != nil {
+			return err
+		}
+	}
 	return nil
 }
 
+func cleanUpWorkerPods(mpiJob *kubeflow.MPIJob, c *MPIJobController) error {
+	// set worker StatefulSet Replicas to 0.
+	if err := c.deleteWorkerPods(mpiJob); err != nil {
+		return err
+	}
+	initializeMPIJobStatuses(mpiJob, kubeflow.MPIReplicaTypeWorker)
+	if c.gangSchedulerName != "" {
+		if err := c.deletePodGroups(mpiJob); err != nil {
+			return err
+		}
+	}
+	mpiJob.Status.ReplicaStatuses[kubeflow.MPIReplicaTypeWorker].Active = 0
+	return c.updateStatusHandler(mpiJob)
+}
+
 // getLauncherJob gets the launcher Job controlled by this MPIJob.
 func (c *MPIJobController) getLauncherJob(mpiJob *kubeflow.MPIJob) (*batchv1.Job, error) {
 	launcher, err := c.jobLister.Jobs(mpiJob.Namespace).Get(mpiJob.Name + launcherSuffix)
@@ -857,6 +880,14 @@ func (c *MPIJobController) getOrCreateWorker(mpiJob *kubeflow.MPIJob) ([]*corev1
 	return workerPods, nil
 }
 
+func isMPIJobSuspended(mpiJob *kubeflow.MPIJob) bool {
+	return pointer.BoolDeref(mpiJob.Spec.RunPolicy.Suspend, false)
+}
+
+func isJobSuspended(job *batchv1.Job) bool {
+	return pointer.BoolDeref(job.Spec.Suspend, false)
+}
+
 func (c *MPIJobController) deleteWorkerPods(mpiJob *kubeflow.MPIJob) error {
 	var (
 		workerPrefix       = mpiJob.Name + workerSuffix
@@ -905,6 +936,19 @@ func (c *MPIJobController) updateMPIJobStatus(mpiJob *kubeflow.MPIJob, launcher
 	if err != nil {
 		return fmt.Errorf("checking launcher pods running: %w", err)
 	}
+	if isMPIJobSuspended(mpiJob) {
+		// it is suspended now
+		if updateMPIJobConditions(mpiJob, kubeflow.JobSuspended, v1.ConditionTrue, "MPIJobSuspended", "MPIJob suspended") {
+			c.recorder.Event(mpiJob, corev1.EventTypeNormal, "MPIJobSuspended", "MPIJob suspended")
+		}
+	} else if getCondition(mpiJob.Status, kubeflow.JobSuspended) != nil {
+		// it is not suspended now, consider resumed if the condition was set before
+		if updateMPIJobConditions(mpiJob, kubeflow.JobSuspended, v1.ConditionFalse, "MPIJobResumed", "MPIJob resumed") {
+			c.recorder.Event(mpiJob, corev1.EventTypeNormal, "MPIJobResumed", "MPIJob resumed")
+			now := metav1.NewTime(time.Now())
+			mpiJob.Status.StartTime = &now
+		}
+	}
 	// Job.status.Active accounts for Pending and Running pods. Count running pods
 	// from the lister instead.
 	launcherPodsCnt := countRunningPods(launcherPods)
@@ -919,7 +963,7 @@ func (c *MPIJobController) updateMPIJobStatus(mpiJob *kubeflow.MPIJob, launcher
 			if mpiJob.Status.CompletionTime == nil {
 				mpiJob.Status.CompletionTime = launcher.Status.CompletionTime
 			}
-			updateMPIJobConditions(mpiJob, kubeflow.JobSucceeded, mpiJobSucceededReason, msg)
+			updateMPIJobConditions(mpiJob, kubeflow.JobSucceeded, v1.ConditionTrue, mpiJobSucceededReason, msg)
 			mpiJobsSuccessCount.Inc()
 		} else if isJobFailed(launcher) {
 			c.updateMPIJobFailedStatus(mpiJob, launcher, launcherPods)
@@ -953,14 +997,17 @@ func (c *MPIJobController) updateMPIJobStatus(mpiJob *kubeflow.MPIJob, launcher
 	if evict > 0 {
 		msg := fmt.Sprintf("%d/%d workers are evicted", evict, len(worker))
 		klog.Infof("MPIJob <%s/%s>: %v", mpiJob.Namespace, mpiJob.Name, msg)
-		updateMPIJobConditions(mpiJob, kubeflow.JobFailed, mpiJobEvict, msg)
+		updateMPIJobConditions(mpiJob, kubeflow.JobFailed, v1.ConditionTrue, mpiJobEvict, msg)
 		c.recorder.Event(mpiJob, corev1.EventTypeWarning, mpiJobEvict, msg)
 	}
 
 	if launcher != nil && launcherPodsCnt >= 1 && running == len(worker) {
 		msg := fmt.Sprintf("MPIJob %s/%s is running.", mpiJob.Namespace, mpiJob.Name)
-		updateMPIJobConditions(mpiJob, kubeflow.JobRunning, mpiJobRunningReason, msg)
+		updateMPIJobConditions(mpiJob, kubeflow.JobRunning, v1.ConditionTrue, mpiJobRunningReason, msg)
 		c.recorder.Eventf(mpiJob, corev1.EventTypeNormal, "MPIJobRunning", "MPIJob %s/%s is running", mpiJob.Namespace, mpiJob.Name)
+	} else if isMPIJobSuspended(mpiJob) {
+		msg := fmt.Sprintf("MPIJob %s/%s is suspended.", mpiJob.Namespace, mpiJob.Name)
+		updateMPIJobConditions(mpiJob, kubeflow.JobRunning, v1.ConditionFalse, mpiJobSuspendedReason, msg)
 	}
 
 	// no need to update the mpijob if the status hasn't changed since last time.
@@ -999,7 +1046,7 @@ func (c *MPIJobController) updateMPIJobFailedStatus(mpiJob *kubeflow.MPIJob, lau
 		now := metav1.Now()
 		mpiJob.Status.CompletionTime = &now
 	}
-	updateMPIJobConditions(mpiJob, kubeflow.JobFailed, reason, msg)
+	updateMPIJobConditions(mpiJob, kubeflow.JobFailed, v1.ConditionTrue, reason, msg)
 	mpiJobsFailureCount.Inc()
 }
 
@@ -1304,7 +1351,7 @@ func (c *MPIJobController) newWorker(mpiJob *kubeflow.MPIJob, index int) *corev1
 }
 
 func (c *MPIJobController) newLauncherJob(mpiJob *kubeflow.MPIJob) *batchv1.Job {
-	return &batchv1.Job{
+	job := &batchv1.Job{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:      mpiJob.Name + launcherSuffix,
 			Namespace: mpiJob.Namespace,
@@ -1322,6 +1369,10 @@ func (c *MPIJobController) newLauncherJob(mpiJob *kubeflow.MPIJob) *batchv1.Job
 			Template:                c.newLauncherPodTemplate(mpiJob),
 		},
 	}
+	if isMPIJobSuspended(mpiJob) {
+		job.Spec.Suspend = pointer.Bool(true)
+	}
+	return job
 }
 
 // newLauncherPodTemplate creates a new launcher Job for an MPIJob resource. It also sets

pkg/apis/kubeflow/v2beta1/zz_generated.deepcopy.go

@@ -28,6 +28,8 @@ const (
 	mpiJobSucceededReason = "MPIJobSucceeded"
 	// mpiJobRunningReason is added in a mpijob when it is running.
 	mpiJobRunningReason = "MPIJobRunning"
+	// mpiJobSuspendedReason is added in a mpijob when it is suspended.
+	mpiJobSuspendedReason = "MPIJobSuspended"
 	// mpiJobFailedReason is added in a mpijob when it is failed.
 	mpiJobFailedReason = "MPIJobFailed"
 	// mpiJobEvict
@@ -45,16 +47,16 @@ func initializeMPIJobStatuses(mpiJob *kubeflow.MPIJob, mtype kubeflow.MPIReplica
 }
 
 // updateMPIJobConditions updates the conditions of the given mpiJob.
-func updateMPIJobConditions(mpiJob *kubeflow.MPIJob, conditionType kubeflow.JobConditionType, reason, message string) {
-	condition := newCondition(conditionType, reason, message)
-	setCondition(&mpiJob.Status, condition)
+func updateMPIJobConditions(mpiJob *kubeflow.MPIJob, conditionType kubeflow.JobConditionType, status v1.ConditionStatus, reason, message string) bool {
+	condition := newCondition(conditionType, status, reason, message)
+	return setCondition(&mpiJob.Status, condition)
 }
 
 // newCondition creates a new mpiJob condition.
-func newCondition(conditionType kubeflow.JobConditionType, reason, message string) kubeflow.JobCondition {
+func newCondition(conditionType kubeflow.JobConditionType, status v1.ConditionStatus, reason, message string) kubeflow.JobCondition {
 	return kubeflow.JobCondition{
 		Type:               conditionType,
-		Status:             v1.ConditionTrue,
+		Status:             status,
 		LastUpdateTime:     metav1.Now(),
 		LastTransitionTime: metav1.Now(),
 		Reason:             reason,
@@ -96,13 +98,12 @@ func isFailed(status kubeflow.JobStatus) bool {
 // setCondition updates the mpiJob to include the provided condition.
 // If the condition that we are about to add already exists
 // and has the same status and reason then we are not going to update.
-func setCondition(status *kubeflow.JobStatus, condition kubeflow.JobCondition) {
-
+func setCondition(status *kubeflow.JobStatus, condition kubeflow.JobCondition) bool {
 	currentCond := getCondition(*status, condition.Type)
 
 	// Do nothing if condition doesn't change
 	if currentCond != nil && currentCond.Status == condition.Status && currentCond.Reason == condition.Reason {
-		return
+		return false
 	}
 
 	// Do not update lastTransitionTime if the status of the condition doesn't change.
@@ -113,6 +114,7 @@ func setCondition(status *kubeflow.JobStatus, condition kubeflow.JobCondition) {
 	// Append the updated condition
 	newConditions := filterOutCondition(status.Conditions, condition.Type)
 	status.Conditions = append(newConditions, condition)
+	return true
 }
 
 // filterOutCondition returns a new slice of mpiJob conditions without conditions with the provided type.