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
 	}
@@ -535,38 +526,40 @@ func (c *MPIJobController) syncHandler(key string) error {
 	// We're done if the launcher either succeeded or failed.
 	done := launcher != nil && isJobFinished(launcher)
 	if !done {
-		_, err := c.getOrCreateService(mpiJob, newWorkersService(mpiJob))
-		if err != nil {
-			return fmt.Errorf("getting or creating Service to front workers: %w", err)
-		}
+		if !isMPIJobSuspended(mpiJob) {
+			_, err := c.getOrCreateService(mpiJob, newWorkersService(mpiJob))
+			if err != nil {
+				return fmt.Errorf("getting or creating Service to front workers: %w", err)
+			}
 
-		if config, err := c.getOrCreateConfigMap(mpiJob); config == nil || err != nil {
-			return fmt.Errorf("getting or creating ConfigMap: %w", err)
-		}
+			if config, err := c.getOrCreateConfigMap(mpiJob); config == nil || err != nil {
+				return fmt.Errorf("getting or creating ConfigMap: %w", err)
+			}
 
-		_, err = c.getOrCreateSSHAuthSecret(mpiJob)
-		if err != nil {
-			return fmt.Errorf("creating SSH auth secret: %w", err)
-		}
+			_, err = c.getOrCreateSSHAuthSecret(mpiJob)
+			if err != nil {
+				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 {
-				return err
+			// 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
-		}
-		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
-			// has the same name as the launcher's hostname.
-			_, err := c.getOrCreateService(mpiJob, newLauncherService(mpiJob))
+			worker, err = c.getOrCreateWorker(mpiJob)
 			if err != nil {
-				return fmt.Errorf("getting or creating Service to front launcher: %w", err)
+				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
+				// has the same name as the launcher's hostname.
+				_, err := c.getOrCreateService(mpiJob, newLauncherService(mpiJob))
+				if err != nil {
+					return fmt.Errorf("getting or creating Service to front launcher: %w", err)
+				}
 			}
 		}
 		if launcher == nil {
@@ -585,9 +578,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 +881,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 +937,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 +964,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 +998,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 +1044,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 +1349,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 +1367,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