run worker in launcher pod

Author: kuizhiqing

examples/v2beta1/pi/pi-intel.yaml

@@ -27,7 +27,7 @@ spec:
             resources:
               limits:
                 cpu: 1
-                memory: 1Gi
+                memory: 200m
     Worker:
       replicas: 2
       template:

examples/v2beta1/pi/pi-mpich.yaml

@@ -27,7 +27,7 @@ spec:
             resources:
               limits:
                 cpu: 1
-                memory: 1Gi
+                memory: 200m
     Worker:
       replicas: 2
       template:

examples/v2beta1/pi/pi.yaml

@@ -27,7 +27,7 @@ spec:
             resources:
               limits:
                 cpu: 1
-                memory: 1Gi
+                memory: 200m
     Worker:
       replicas: 2
       template:

examples/v2beta1/pi/pi2.yaml

@@ -0,0 +1,48 @@
+apiVersion: kubeflow.org/v2beta1
+kind: MPIJob
+metadata:
+  name: pi2
+spec:
+  slotsPerWorker: 1
+  runPolicy:
+    cleanPodPolicy: Running
+    ttlSecondsAfterFinished: 60
+  sshAuthMountPath: /home/mpiuser/.ssh
+  mpiReplicaSpecs:
+    Launcher:
+      replicas: 1
+      template:
+        spec:
+          containers:
+          - image: mpioperator/mpi-pi:openmpi
+            name: mpi-launcher
+            securityContext:
+              runAsUser: 1000
+            command:
+            - bash
+            args:
+            - -c
+            - "/usr/sbin/sshd -f /home/mpiuser/.sshd_config && mpirun /home/mpiuser/pi"
+            resources:
+              limits:
+                cpu: 1
+                memory: 1Gi
+    Worker:
+      replicas: 2
+      template:
+        spec:
+          containers:
+          - image: mpioperator/mpi-pi:openmpi
+            name: mpi-worker
+            securityContext:
+              runAsUser: 1000
+            command:
+            - /usr/sbin/sshd
+            args:
+            - -De
+            - -f
+            - /home/mpiuser/.sshd_config
+            resources:
+              limits:
+                cpu: 1
+                memory: 1Gi

pkg/controller/mpi_job_controller.go

@@ -628,15 +628,12 @@ func (c *MPIJobController) syncHandler(key string) error {
 				return err
 			}
 		}
-		if mpiJob.Spec.MPIImplementation == kubeflow.MPIImplementationIntel ||
-			mpiJob.Spec.MPIImplementation == kubeflow.MPIImplementationMPICH {
-			// The Intel and MPICH implementations require 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)
-			}
+		// NEW: always create service for launcher for different implementations and compatible with running launcher as worker
+		// The Intel and MPICH implementations require 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.
+		if _, err = c.getOrCreateService(mpiJob, newLauncherService(mpiJob)); err != nil {
+			return fmt.Errorf("getting or creating Service to front launcher: %w", err)
 		}
 		if launcher == nil {
 			if mpiJob.Spec.LauncherCreationPolicy == kubeflow.LauncherCreationPolicyAtStartup || c.countReadyWorkerPods(worker) == len(worker) {
@@ -1246,6 +1243,26 @@ func (c *MPIJobController) doUpdateJobStatus(mpiJob *kubeflow.MPIJob) error {
 	return err
 }
 
+// enableLauncherAsWorker check whether to run worker process in launcher
+func enableLauncherAsWorker(mpiJob *kubeflow.MPIJob) bool {
+	// case 1: have no worker
+	worker := mpiJob.Spec.MPIReplicaSpecs[kubeflow.MPIReplicaTypeWorker]
+	if worker == nil || len(worker.Template.Spec.Containers) < 1 {
+		return true
+	}
+
+	// case -1: no resource declaration for launcher
+	launcher := mpiJob.Spec.MPIReplicaSpecs[kubeflow.MPIReplicaTypeLauncher]
+	launcherRes := launcher.Template.Spec.Containers[0].Resources
+	if launcherRes.Limits == nil && launcherRes.Requests == nil {
+		return false
+	}
+
+	// case 2: launcher declare the same resource as worker
+	workerRes := worker.Template.Spec.Containers[0].Resources
+	return equality.Semantic.DeepEqual(workerRes, launcherRes)
+}
+
 // newConfigMap creates a new ConfigMap containing configurations for an MPIJob
 // resource. It also sets the appropriate OwnerReferences on the resource so
 // handleObject can discover the MPIJob resource that 'owns' it.
@@ -1256,12 +1273,26 @@ func newConfigMap(mpiJob *kubeflow.MPIJob, workerReplicas int32) *corev1.ConfigM
 	if mpiJob.Spec.SlotsPerWorker != nil {
 		slots = int(*mpiJob.Spec.SlotsPerWorker)
 	}
+	// note that pod.spec.dnsConfig also affect the svc resolution
+	// ref: https://kubernetes.io/docs/concepts/services-networking/dns-pod-service/
+	// launcher can be reach with hostname or service name
+	if enableLauncherAsWorker(mpiJob) {
+		launcherService := mpiJob.Name + launcherSuffix
+		switch mpiJob.Spec.MPIImplementation {
+		case kubeflow.MPIImplementationOpenMPI:
+			buffer.WriteString(fmt.Sprintf("%s.%s.svc slots=%d\n", launcherService, mpiJob.Namespace, slots))
+		case kubeflow.MPIImplementationIntel, kubeflow.MPIImplementationMPICH:
+			buffer.WriteString(fmt.Sprintf("%s.%s.svc:%d\n", launcherService, mpiJob.Namespace, slots))
+		}
+	}
+
 	for i := 0; i < int(workerReplicas); i++ {
+		name := workerName(mpiJob, i)
 		switch mpiJob.Spec.MPIImplementation {
 		case kubeflow.MPIImplementationOpenMPI:
-			buffer.WriteString(fmt.Sprintf("%s%s-%d.%s.%s.svc slots=%d\n", mpiJob.Name, workerSuffix, i, workersService, mpiJob.Namespace, slots))
+			buffer.WriteString(fmt.Sprintf("%s.%s.%s.svc slots=%d\n", name, workersService, mpiJob.Namespace, slots))
 		case kubeflow.MPIImplementationIntel, kubeflow.MPIImplementationMPICH:
-			buffer.WriteString(fmt.Sprintf("%s%s-%d.%s.%s.svc:%d\n", mpiJob.Name, workerSuffix, i, workersService, mpiJob.Namespace, slots))
+			buffer.WriteString(fmt.Sprintf("%s.%s.%s.svc:%d\n", name, workersService, mpiJob.Namespace, slots))
 		}
 	}
 

pkg/controller/mpi_job_controller_test.go

@@ -25,6 +25,7 @@ import (
 	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
 	schedulingv1 "k8s.io/api/scheduling/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/schema"
@@ -526,10 +527,7 @@ func TestAllResourcesCreated(t *testing.T) {
 			for i := 0; i < 5; i++ {
 				f.expectCreatePodAction(fmjc.newWorker(mpiJobCopy, i))
 			}
-			if implementation == kubeflow.MPIImplementationIntel ||
-				implementation == kubeflow.MPIImplementationMPICH {
-				f.expectCreateServiceAction(newLauncherService(mpiJobCopy))
-			}
+			f.expectCreateServiceAction(newLauncherService(mpiJobCopy))
 			f.expectCreateJobAction(fmjc.newLauncherJob(mpiJobCopy))
 
 			mpiJobCopy.Status.Conditions = []kubeflow.JobCondition{newCondition(kubeflow.JobCreated, corev1.ConditionTrue, mpiJobCreatedReason, "MPIJob default/foo is created.")}
@@ -822,10 +820,7 @@ func TestCreateSuspendedMPIJob(t *testing.T) {
 				t.Fatalf("Failed creating secret")
 			}
 			f.expectCreateSecretAction(secret)
-			if implementation == kubeflow.MPIImplementationIntel ||
-				implementation == kubeflow.MPIImplementationMPICH {
-				f.expectCreateServiceAction(newLauncherService(mpiJob))
-			}
+			f.expectCreateServiceAction(newLauncherService(mpiJob))
 
 			// expect creating of the launcher
 			fmjc := f.newFakeMPIJobController()
@@ -887,6 +882,7 @@ func TestSuspendedRunningMPIJob(t *testing.T) {
 
 	// setup objects
 	scheme.Scheme.Default(mpiJob)
+	f.setUpService(newLauncherService(mpiJob))
 	f.setUpService(newWorkersService(mpiJob))
 
 	cfgMap := newConfigMap(mpiJob, replicas)
@@ -991,6 +987,7 @@ func TestResumeMPIJob(t *testing.T) {
 	// expect the launcher update to resume it
 	launcherCopy := launcher.DeepCopy()
 	launcherCopy.Spec.Suspend = pointer.Bool(false)
+	f.expectCreateServiceAction(newLauncherService(mpiJob))
 	f.expectUpdateJobAction(launcherCopy)
 
 	// expect an update to add the conditions
@@ -1047,6 +1044,7 @@ func TestLauncherActiveWorkerNotReady(t *testing.T) {
 	configMap := newConfigMap(mpiJobCopy, replicas)
 	updateDiscoverHostsInConfigMap(configMap, mpiJobCopy, nil)
 	f.setUpConfigMap(configMap)
+	f.setUpService(newLauncherService(mpiJob))
 	f.setUpService(newWorkersService(mpiJobCopy))
 	secret, err := newSSHAuthSecret(mpiJobCopy)
 	if err != nil {
@@ -1097,6 +1095,7 @@ func TestLauncherActiveWorkerReady(t *testing.T) {
 
 	mpiJobCopy := mpiJob.DeepCopy()
 	scheme.Scheme.Default(mpiJobCopy)
+	f.setUpService(newLauncherService(mpiJob))
 	f.setUpService(newWorkersService(mpiJobCopy))
 	secret, err := newSSHAuthSecret(mpiJobCopy)
 	if err != nil {
@@ -1156,6 +1155,7 @@ func TestWorkerReady(t *testing.T) {
 
 	mpiJobCopy := mpiJob.DeepCopy()
 	scheme.Scheme.Default(mpiJobCopy)
+	f.setUpService(newLauncherService(mpiJob))
 	f.setUpService(newWorkersService(mpiJobCopy))
 	secret, err := newSSHAuthSecret(mpiJobCopy)
 	if err != nil {
@@ -1532,20 +1532,71 @@ func TestNewLauncherAndWorker(t *testing.T) {
 	}
 }
 
+func newReplicaSpec(name string, cpu string) *kubeflow.ReplicaSpec {
+	return &kubeflow.ReplicaSpec{
+		Template: corev1.PodTemplateSpec{
+			Spec: corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name: name,
+						Resources: corev1.ResourceRequirements{
+							Limits: corev1.ResourceList{
+								corev1.ResourceCPU: resource.MustParse(cpu),
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+}
+
 func TestNewConfigMap(t *testing.T) {
 	testCases := map[string]struct {
 		mpiJob         *kubeflow.MPIJob
 		workerReplicas int32
 		wantCM         *corev1.ConfigMap
 	}{
-		"OpenMPI without slots": {
+		"OpenMPI without slots, enable launcher as worker": {
 			mpiJob: &kubeflow.MPIJob{
 				ObjectMeta: metav1.ObjectMeta{
 					Name:      "openmpi-without-slots",
 					Namespace: "tenant-a",
 				},
 				Spec: kubeflow.MPIJobSpec{
 					MPIImplementation: kubeflow.MPIImplementationOpenMPI,
+					MPIReplicaSpecs: map[kubeflow.MPIReplicaType]*kubeflow.ReplicaSpec{
+						kubeflow.MPIReplicaTypeLauncher: newReplicaSpec("launcher", "2"),
+						kubeflow.MPIReplicaTypeWorker:   newReplicaSpec("worker", "2"),
+					},
+				},
+			},
+			workerReplicas: 2,
+			wantCM: &corev1.ConfigMap{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "openmpi-without-slots-config",
+					Namespace: "tenant-a",
+					Labels: map[string]string{
+						"app": "openmpi-without-slots",
+					},
+				},
+				Data: map[string]string{
+					"hostfile": "openmpi-without-slots-launcher.tenant-a.svc slots=1\nopenmpi-without-slots-worker-0.openmpi-without-slots-worker.tenant-a.svc slots=1\nopenmpi-without-slots-worker-1.openmpi-without-slots-worker.tenant-a.svc slots=1\n",
+				},
+			},
+		},
+		"OpenMPI without slots, disable launcher as worker": {
+			mpiJob: &kubeflow.MPIJob{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "openmpi-without-slots",
+					Namespace: "tenant-a",
+				},
+				Spec: kubeflow.MPIJobSpec{
+					MPIImplementation: kubeflow.MPIImplementationOpenMPI,
+					MPIReplicaSpecs: map[kubeflow.MPIReplicaType]*kubeflow.ReplicaSpec{
+						kubeflow.MPIReplicaTypeLauncher: newReplicaSpec("launcher", "1"),
+						kubeflow.MPIReplicaTypeWorker:   newReplicaSpec("worker", "2"),
+					},
 				},
 			},
 			workerReplicas: 2,
@@ -1571,6 +1622,10 @@ func TestNewConfigMap(t *testing.T) {
 				Spec: kubeflow.MPIJobSpec{
 					SlotsPerWorker:    pointer.Int32(10),
 					MPIImplementation: kubeflow.MPIImplementationIntel,
+					MPIReplicaSpecs: map[kubeflow.MPIReplicaType]*kubeflow.ReplicaSpec{
+						kubeflow.MPIReplicaTypeLauncher: newReplicaSpec("launcher", "1"),
+						kubeflow.MPIReplicaTypeWorker:   newReplicaSpec("worker", "2"),
+					},
 				},
 			},
 			workerReplicas: 1,
@@ -1596,6 +1651,10 @@ func TestNewConfigMap(t *testing.T) {
 				Spec: kubeflow.MPIJobSpec{
 					SlotsPerWorker:    pointer.Int32(10),
 					MPIImplementation: kubeflow.MPIImplementationMPICH,
+					MPIReplicaSpecs: map[kubeflow.MPIReplicaType]*kubeflow.ReplicaSpec{
+						kubeflow.MPIReplicaTypeLauncher: newReplicaSpec("launcher", "1"),
+						kubeflow.MPIReplicaTypeWorker:   newReplicaSpec("worker", "2"),
+					},
 				},
 			},
 			workerReplicas: 1,

test/integration/mpi_job_controller_test.go

@@ -17,6 +17,7 @@ package integration
 import (
 	"context"
 	"fmt"
+	"strings"
 	"testing"
 	"time"
 
@@ -895,7 +896,7 @@ func validateMPIJobDependencies(
 	if err := wait.PollUntilContextTimeout(ctx, waitInterval, wait.ForeverTestTimeout, false, func(ctx context.Context) (bool, error) {
 		problems = nil
 		var err error
-		svc, err = getServiceForJob(ctx, kubeClient, job)
+		svc, err = getServiceForJob(ctx, kubeClient, job, "worker")
 		if err != nil {
 			return false, err
 		}
@@ -1026,14 +1027,16 @@ func updatePodsCondition(ctx context.Context, client kubernetes.Interface, pods
 	return nil
 }
 
-func getServiceForJob(ctx context.Context, client kubernetes.Interface, job *kubeflow.MPIJob) (*corev1.Service, error) {
+func getServiceForJob(ctx context.Context, client kubernetes.Interface, job *kubeflow.MPIJob, role string) (*corev1.Service, error) {
 	result, err := client.CoreV1().Services(job.Namespace).List(ctx, metav1.ListOptions{})
 	if err != nil {
 		return nil, err
 	}
 	for _, obj := range result.Items {
-		if metav1.IsControlledBy(&obj, job) {
-			return &obj, nil
+		if strings.HasSuffix(obj.Name, role) {
+			if metav1.IsControlledBy(&obj, job) {
+				return &obj, nil
+			}
 		}
 	}
 	return nil, nil