Add default Intel MPI env variables to MPIJob

Co-authored-by: Yuki Iwai <[email protected]>
Signed-off-by: Tuomas Katila <[email protected]>

Author: tkatila

pkg/controller.v1/mpi/mpijob.go

@@ -45,6 +45,7 @@ const (
 	initContainerCpu        = "100m"
 	initContainerEphStorage = "5Gi"
 	initContainerMem        = "512Mi"
+	iMPIDefaultBootstrap    = "rsh"
 )
 
 const (
@@ -218,6 +219,26 @@ func isGPULauncher(mpiJob *kubeflowv1.MPIJob) bool {
 	return false
 }
 
+// hasIntelMPIBootstrapValues returns the existence of I_MPI_HYDRA_BOOTSTRAP
+// and I_MPI_HYDRA_BOOTSTRAP_EXEC values.
+// There are also _EXEC_EXTRA_ARGS and _AUTOFORK under the I_MPI_HYDRA_BOOTSTRAP
+// prefix but those are not checked on purpose.
+func hasIntelMPIBootstrapValues(envs []corev1.EnvVar) (bootstrap, exec bool) {
+	for _, env := range envs {
+		if env.Name == "I_MPI_HYDRA_BOOTSTRAP" {
+			bootstrap = true
+		} else if env.Name == "I_MPI_HYDRA_BOOTSTRAP_EXEC" {
+			exec = true
+		}
+
+		if bootstrap && exec {
+			break
+		}
+	}
+
+	return bootstrap, exec
+}
+
 func defaultReplicaLabels(genericLabels map[string]string, roleLabelVal string) map[string]string {
 	replicaLabels := map[string]string{}
 	for k, v := range genericLabels {

pkg/controller.v1/mpi/mpijob_controller.go

@@ -1152,6 +1152,26 @@ func (jc *MPIJobReconciler) newLauncher(mpiJob *kubeflowv1.MPIJob, kubectlDelive
 			})
 	}
 
+	// Add default Intel MPI bootstrap variables if not provided by the user.
+	if bootstrap, exec := hasIntelMPIBootstrapValues(container.Env); !bootstrap || !exec {
+		if !bootstrap {
+			container.Env = append(container.Env,
+				corev1.EnvVar{
+					Name:  "I_MPI_HYDRA_BOOTSTRAP",
+					Value: iMPIDefaultBootstrap,
+				},
+			)
+		}
+		if !exec {
+			container.Env = append(container.Env,
+				corev1.EnvVar{
+					Name:  "I_MPI_HYDRA_BOOTSTRAP_EXEC",
+					Value: fmt.Sprintf("%s/%s", configMountPath, kubexecScriptName),
+				},
+			)
+		}
+	}
+
 	container.VolumeMounts = append(container.VolumeMounts,
 		corev1.VolumeMount{
 			Name:      kubectlVolumeName,

pkg/controller.v1/mpi/mpijob_controller_test.go

@@ -709,6 +709,93 @@ var _ = Describe("MPIJob controller", func() {
 		})
 	})
 
+	Context("Test launcher's Intel MPI handling", func() {
+		It("Should create a launcher job with Intel MPI env variables", func() {
+			getIMPIValues := func(envs []corev1.EnvVar) (bootstrap, exec string) {
+				for _, env := range envs {
+					if env.Name == "I_MPI_HYDRA_BOOTSTRAP" {
+						bootstrap = env.Value
+					}
+					if env.Name == "I_MPI_HYDRA_BOOTSTRAP_EXEC" {
+						exec = env.Value
+					}
+				}
+
+				return bootstrap, exec
+			}
+
+			By("By creating MPIJobs with and without preset env variables")
+
+			testCases := map[string]struct {
+				envVariables  map[string]string
+				expectedAdded int
+			}{
+				"withoutIMPIValues": {
+					envVariables:  map[string]string{"X_MPI_HYDRA_BOOTSTRAP": "foo"},
+					expectedAdded: 2,
+				},
+				"withIMPIBootstrap": {
+					envVariables:  map[string]string{"I_MPI_HYDRA_BOOTSTRAP": "RSH"},
+					expectedAdded: 1,
+				},
+				"withIMPIBootstrapExec": {
+					envVariables:  map[string]string{"I_MPI_HYDRA_BOOTSTRAP_EXEC": "/script.sh"},
+					expectedAdded: 1,
+				},
+				"withIMPIBootstrapAndExec": {
+					envVariables: map[string]string{
+						"I_MPI_HYDRA_BOOTSTRAP":      "RSH",
+						"I_MPI_HYDRA_BOOTSTRAP_EXEC": "/script.sh",
+					},
+					expectedAdded: 0,
+				},
+			}
+
+			for testName, testCase := range testCases {
+				ctx := context.Background()
+				startTime := metav1.Now()
+				completionTime := metav1.Now()
+
+				jobName := "test-launcher-creation-" + strings.ToLower(testName)
+
+				var replicas int32 = 1
+				mpiJob := newMPIJob(jobName, &replicas, 1, gpuResourceName, &startTime, &completionTime)
+				Expect(testK8sClient.Create(ctx, mpiJob)).Should(Succeed())
+
+				template := &mpiJob.Spec.MPIReplicaSpecs[kubeflowv1.MPIJobReplicaTypeLauncher].Template
+				Expect(len(template.Spec.Containers) == 1).To(BeTrue())
+
+				cont := &template.Spec.Containers[0]
+
+				for k, v := range testCase.envVariables {
+					cont.Env = append(cont.Env,
+						corev1.EnvVar{
+							Name:  k,
+							Value: v,
+						},
+					)
+				}
+
+				launcher := reconciler.newLauncher(mpiJob, "kubectl-delivery", false)
+
+				Expect(len(launcher.Spec.Containers) == 1).To(BeTrue())
+
+				bootstrap, exec := getIMPIValues(launcher.Spec.Containers[0].Env)
+
+				if val, found := testCase.envVariables["I_MPI_HYDRA_BOOTSTRAP"]; found {
+					Expect(bootstrap).To(BeEquivalentTo(val))
+				} else {
+					Expect(bootstrap).To(BeEquivalentTo(iMPIDefaultBootstrap))
+				}
+
+				if val, found := testCase.envVariables["I_MPI_HYDRA_BOOTSTRAP_EXEC"]; found {
+					Expect(exec).To(BeEquivalentTo(val))
+				} else {
+					Expect(exec).To(BeEquivalentTo("/etc/mpi/kubexec.sh"))
+				}
+			}
+		})
+	})
 })
 
 func ReplicaStatusMatch(replicaStatuses map[common.ReplicaType]*common.ReplicaStatus,