Implement support for suspend semantics for MPIJob

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 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

@@ -295,6 +295,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 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 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 {
@@ -130,3 +142,8 @@ const (
 	MPIImplementationOpenMPI MPIImplementation = "OpenMPI"
 	MPIImplementationIntel   MPIImplementation = "Intel"
 )
+
+const (
+	// JobSuspended means the job has been suspended.
+	JobSuspended common.JobConditionType = "Suspended"
+)

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, common.JobCreated, mpiJobCreatedReason, msg)
+		updateMPIJobConditions(mpiJob, common.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[common.ReplicaType(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
 	}
@@ -532,9 +523,10 @@ func (c *MPIJobController) syncHandler(key string) error {
 	}
 
 	var worker []*corev1.Pod
-	// We're done if the launcher either succeeded or failed.
-	done := launcher != nil && isJobFinished(launcher)
-	if !done {
+	// We're done if the launcher either succeeded or failed. We also skip
+	// creation of the auxiliary objects if the MPIJob is suspended.
+	running := !(launcher != nil && isJobFinished(launcher)) && !isMPIJobSuspended(mpiJob)
+	if running {
 		_, err := c.getOrCreateService(mpiJob, newWorkersService(mpiJob))
 		if err != nil {
 			return fmt.Errorf("getting or creating Service to front workers: %w", err)
@@ -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[common.ReplicaType(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, "Suspended", "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.ConditionTrue, "MPIJobResumed", "MPIJob resumed") {
+			c.recorder.Event(mpiJob, corev1.EventTypeNormal, "Resumed", "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, common.JobSucceeded, mpiJobSucceededReason, msg)
+			updateMPIJobConditions(mpiJob, common.JobSucceeded, v1.ConditionTrue, mpiJobSucceededReason, msg)
 			mpiJobsSuccessCount.Inc()
 		} else if isJobFailed(launcher) {
 			c.updateMPIJobFailedStatus(mpiJob, launcher, launcherPods)
@@ -953,13 +997,13 @@ 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, common.JobFailed, mpiJobEvict, msg)
+		updateMPIJobConditions(mpiJob, common.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, common.JobRunning, mpiJobRunningReason, msg)
+		updateMPIJobConditions(mpiJob, common.JobRunning, v1.ConditionTrue, mpiJobRunningReason, msg)
 		c.recorder.Eventf(mpiJob, corev1.EventTypeNormal, "MPIJobRunning", "MPIJob %s/%s is running", mpiJob.Namespace, mpiJob.Name)
 	}
 
@@ -999,7 +1043,7 @@ func (c *MPIJobController) updateMPIJobFailedStatus(mpiJob *kubeflow.MPIJob, lau
 		now := metav1.Now()
 		mpiJob.Status.CompletionTime = &now
 	}
-	updateMPIJobConditions(mpiJob, common.JobFailed, reason, msg)
+	updateMPIJobConditions(mpiJob, common.JobFailed, v1.ConditionTrue, reason, msg)
 	mpiJobsFailureCount.Inc()
 }
 
@@ -1304,7 +1348,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 +1366,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

@@ -46,16 +46,16 @@ func initializeMPIJobStatuses(mpiJob *kubeflow.MPIJob, mtype kubeflow.MPIReplica
 }
 
 // updateMPIJobConditions updates the conditions of the given mpiJob.
-func updateMPIJobConditions(mpiJob *kubeflow.MPIJob, conditionType common.JobConditionType, reason, message string) {
-	condition := newCondition(conditionType, reason, message)
-	setCondition(&mpiJob.Status, condition)
+func updateMPIJobConditions(mpiJob *kubeflow.MPIJob, conditionType common.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 common.JobConditionType, reason, message string) common.JobCondition {
+func newCondition(conditionType common.JobConditionType, status v1.ConditionStatus, reason, message string) common.JobCondition {
 	return common.JobCondition{
 		Type:               conditionType,
-		Status:             v1.ConditionTrue,
+		Status:             status,
 		LastUpdateTime:     metav1.Now(),
 		LastTransitionTime: metav1.Now(),
 		Reason:             reason,
@@ -97,13 +97,14 @@ func isFailed(status common.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 *common.JobStatus, condition common.JobCondition) {
+// Returns true if the condition is actually updated. False otherwise.
+func setCondition(status *common.JobStatus, condition common.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.
@@ -114,6 +115,7 @@ func setCondition(status *common.JobStatus, condition common.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.