MPICH support

.github/workflows/main.yml

@@ -1,5 +1,5 @@
 name: build
-on: [push, pull_request]
+on: [push, pull_request, workflow_dispatch]
 jobs:
   validate:
     name: Validate 

Makefile

@@ -101,6 +101,9 @@ test_images:
 	${IMG_BUILDER} build -t mpioperator/intel build/base -f build/base/intel.Dockerfile
 	${IMG_BUILDER} build -t mpioperator/intel-builder build/base -f build/base/intel-builder.Dockerfile
 	${IMG_BUILDER} build -t mpioperator/mpi-pi:intel examples/v2beta1/pi -f examples/v2beta1/pi/intel.Dockerfile
+	${IMG_BUILDER} build -t mpioperator/mpich build/base -f build/base/mpich.Dockerfile
+	${IMG_BUILDER} build -t mpioperator/mpich-builder build/base -f build/base/mpich-builder.Dockerfile
+	${IMG_BUILDER} build -t mpioperator/mpi-pi:mpich examples/v2beta1/pi -f examples/v2beta1/pi/mpich.Dockerfile
 
 .PHONY: tidy
 tidy:

README.md

@@ -209,7 +209,13 @@ total images/sec: 308.27
 For a sample that uses Intel MPI, see:
 
 ```bash
-cat examples/pi/pi-intel.yaml
+cat examples/v2beta1/pi/pi-intel.yaml
+```
+
+For a sample that uses MPICH, see:
+
+```bash
+cat examples/v2beta1/pi/pi-mpich.yaml
 ```
 
 ## Exposed Metrics

build/base/mpich-builder.Dockerfile

@@ -0,0 +1,7 @@
+FROM debian:buster as builder
+
+RUN apt update \
+    && apt install -y --no-install-recommends \
+        g++ \
+        libmpich-dev \
+    && rm -rf /var/lib/apt/lists/*

build/base/mpich-entrypoint.sh

@@ -0,0 +1,31 @@
+#!/usr/bin/env bash
+
+function resolve_host() {
+  host="$1"
+  check="nslookup $host"
+  max_retry=10
+  counter=0
+  backoff=0.1
+  until $check > /dev/null
+  do
+    if [ $counter -eq $max_retry ]; then
+      echo "Couldn't resolve $host"
+      return
+    fi
+    sleep $backoff
+    echo "Couldn't resolve $host... Retrying"
+    ((counter++))
+    backoff=$(echo - | awk "{print $backoff + $backoff}")
+  done
+  echo "Resolved $host"
+}
+
+if [ "$K_MPI_JOB_ROLE" == "launcher" ]; then
+  resolve_host "$HOSTNAME"
+  cat /etc/mpi/hostfile | while read host
+  do
+    resolve_host $host
+  done
+fi
+
+exec "$@"

build/base/mpich.Dockerfile

@@ -0,0 +1,8 @@
+FROM mpioperator/base
+
+RUN apt update \
+    && apt install -y --no-install-recommends mpich dnsutils \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY mpich-entrypoint.sh /entrypoint.sh
+ENTRYPOINT ["/entrypoint.sh"]

deploy/v2beta1/mpi-operator.yaml

@@ -121,6 +121,7 @@ spec:
                 enum:
                 - OpenMPI
                 - Intel
+                - MPICH
                 type: string
               mpiReplicaSpecs:
                 properties:

examples/v2beta1/pi/README.md

@@ -19,9 +19,15 @@ For Intel MPI:
 docker build -t mpi-pi . -f intel.Dockerfile
 ```
 
+For MPICH:
+
+```bash
+docker build -t mpi-pi . -f mpich.Dockerfile
+```
+
 ## Create MPIJob
 
-Modify `pi.yaml` (for OpenMPI) or `pi-intel.yaml` (for Intel MPI) to set up the
+Modify `pi.yaml` (for OpenMPI), `pi-intel.yaml` (for Intel MPI) or `pi-mpich.yaml` (for MPICH) to set up the
 image name from your own registry.
 
 Then, run:

examples/v2beta1/pi/mpich.Dockerfile

@@ -0,0 +1,8 @@
+FROM mpioperator/mpich-builder as builder
+
+COPY pi.cc /src/pi.cc
+RUN mpic++ /src/pi.cc -o /pi
+
+FROM mpioperator/mpich
+
+COPY --from=builder /pi /home/mpiuser/pi

examples/v2beta1/pi/pi-mpich.yaml

@@ -0,0 +1,54 @@
+apiVersion: kubeflow.org/v2beta1
+kind: MPIJob
+metadata:
+  name: pi
+spec:
+  slotsPerWorker: 1
+  runPolicy:
+    cleanPodPolicy: Running
+  sshAuthMountPath: /home/mpiuser/.ssh
+  mpiImplementation: MPICH
+  mpiReplicaSpecs:
+    Launcher:
+      replicas: 1
+      template:
+        spec:
+          containers:
+          - image: mpioperator/mpi-pi:mpich
+            imagePullPolicy: Always
+            name: mpi-launcher
+            securityContext:
+              runAsUser: 1000
+            args:
+            - mpirun
+            - -n
+            - "2"
+            - /home/mpiuser/pi
+            resources:
+              limits:
+                cpu: 1
+                memory: 1Gi
+    Worker:
+      replicas: 2
+      template:
+        spec:
+          containers:
+          - image: mpioperator/mpi-pi:mpich
+            imagePullPolicy: Always
+            name: mpi-worker
+            securityContext:
+              runAsUser: 1000
+            command:
+            args:
+            - /usr/sbin/sshd
+            - -De
+            - -f
+            - /home/mpiuser/.sshd_config
+            readinessProbe:
+              tcpSocket:
+                port: 2222
+              initialDelaySeconds: 2
+            resources:
+              limits:
+                cpu: 1
+                memory: 1Gi

manifests/base/crd.yaml

@@ -127,7 +127,7 @@ spec:
                 type: string
               mpiImplementation:
                 type: string
-                enum: ["OpenMPI", "Intel"]
+                enum: ["OpenMPI", "Intel", "MPICH"]
               mpiReplicaSpecs:
                 type: object
                 properties:

proposals/scalable-robust-operator.md

@@ -148,7 +148,7 @@ following changes:
       doesn’t support changes to the completions field. This can be supported
       starting from 1.23. In the meantime, we can replicate the behavior by
       creating a new Job and doing Pod adoption.
-  - For Intel MPI, we also need a headless Service to front the launcher,
+  - For Intel MPI and MPICH, we also need a headless Service to front the launcher,
     because workers communicate back to the launcher using its hostname.
 - **Revert the use of the Job API for the launcher.**
   - The Job controller handles retries when the launcher or any of the workers fail.

v2/crd/kubeflow.org_mpijobs.yaml

@@ -37,10 +37,11 @@ spec:
               mpiImplementation:
                 default: OpenMPI
                 description: MPIImplementation is the MPI implementation. Options
-                  are "OpenMPI" (default) and "Intel".
+                  are "OpenMPI" (default), "Intel" and "MPICH".
                 enum:
                 - OpenMPI
                 - Intel
+                - MPICH
                 type: string
               mpiReplicaSpecs:
                 additionalProperties:

v2/pkg/apis/kubeflow/v2beta1/swagger.json

@@ -223,7 +223,7 @@
       ],
       "properties": {
         "mpiImplementation": {
-          "description": "MPIImplementation is the MPI implementation. Options are \"OpenMPI\" (default) and \"Intel\".",
+          "description": "MPIImplementation is the MPI implementation. Options are \"OpenMPI\" (default) \"Intel\" and \"MPICH\".",
           "type": "string"
         },
         "mpiReplicaSpecs": {

v2/pkg/apis/kubeflow/v2beta1/types.go

@@ -56,8 +56,8 @@ type MPIJobSpec struct {
 	SSHAuthMountPath string `json:"sshAuthMountPath,omitempty"`
 
 	// MPIImplementation is the MPI implementation.
-	// Options are "OpenMPI" (default) and "Intel".
-	// +kubebuilder:validation:Enum:=OpenMPI;Intel
+	// Options are "OpenMPI" (default), "Intel" and "MPICH".
+	// +kubebuilder:validation:Enum:=OpenMPI;Intel;MPICH
 	// +kubebuilder:default:=OpenMPI
 	MPIImplementation MPIImplementation `json:"mpiImplementation,omitempty"`
 }
@@ -78,4 +78,5 @@ type MPIImplementation string
 const (
 	MPIImplementationOpenMPI MPIImplementation = "OpenMPI"
 	MPIImplementationIntel   MPIImplementation = "Intel"
+	MPIImplementationMPICH   MPIImplementation = "MPICH"
 )

v2/pkg/apis/kubeflow/validation/validation.go

@@ -35,7 +35,8 @@ var (
 
 	validMPIImplementations = sets.NewString(
 		string(kubeflow.MPIImplementationOpenMPI),
-		string(kubeflow.MPIImplementationIntel))
+		string(kubeflow.MPIImplementationIntel),
+		string(kubeflow.MPIImplementationMPICH))
 
 	validRestartPolicies = sets.NewString(
 		string(common.RestartPolicyNever),

v2/pkg/controller/mpi_job_controller.go

@@ -115,6 +115,7 @@ const (
 
 	openMPISlotsEnv  = "OMPI_MCA_orte_set_default_slots"
 	intelMPISlotsEnv = "I_MPI_PERHOST"
+	mpichSlotsEnv    = "TODO"
 )
 
 var (
@@ -199,6 +200,16 @@ var (
 			Value: "-o ConnectionAttempts=10",
 		},
 	}
+	mpichEnvVars = []corev1.EnvVar{
+		{
+			Name:  "HYDRA_HOST_FILE",
+			Value: fmt.Sprintf("%s/%s", configMountPath, hostfileName),
+		},
+		{
+			Name:  "HYDRA_LAUNCH_EXTRA_ARGS",
+			Value: "-o ConnectionAttempts=10",
+		},
+	}
 	nvidiaDisableEnvVars = []corev1.EnvVar{
 		{Name: "NVIDIA_VISIBLE_DEVICES"},
 		{Name: "NVIDIA_DRIVER_CAPABILITIES"},
@@ -560,8 +571,9 @@ func (c *MPIJobController) syncHandler(key string) error {
 		if err != nil {
 			return err
 		}
-		if mpiJob.Spec.MPIImplementation == kubeflow.MPIImplementationIntel {
-			// The Intel implementation requires workers to communicate with the
+		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))
@@ -1374,6 +1386,12 @@ func (c *MPIJobController) newLauncherPodTemplate(mpiJob *kubeflow.MPIJob) corev
 			Name:  intelMPISlotsEnv,
 			Value: slotsStr,
 		})
+	case kubeflow.MPIImplementationMPICH:
+		container.Env = append(container.Env, mpichEnvVars...)
+		container.Env = append(container.Env, corev1.EnvVar{
+			Name:  mpichSlotsEnv,
+			Value: slotsStr,
+		})
 	}
 
 	container.Env = append(container.Env,

v2/pkg/controller/mpi_job_controller_test.go

@@ -457,7 +457,7 @@ func TestDoNothingWithInvalidMPIJob(t *testing.T) {
 }
 
 func TestAllResourcesCreated(t *testing.T) {
-	impls := []kubeflow.MPIImplementation{kubeflow.MPIImplementationOpenMPI, kubeflow.MPIImplementationIntel}
+	impls := []kubeflow.MPIImplementation{kubeflow.MPIImplementationOpenMPI, kubeflow.MPIImplementationIntel, kubeflow.MPIImplementationMPICH}
 	for _, implementation := range impls {
 		t.Run(string(implementation), func(t *testing.T) {
 			f := newFixture(t)
@@ -481,7 +481,8 @@ func TestAllResourcesCreated(t *testing.T) {
 			for i := 0; i < 5; i++ {
 				f.expectCreatePodAction(fmjc.newWorker(mpiJobCopy, i))
 			}
-			if implementation == kubeflow.MPIImplementationIntel {
+			if implementation == kubeflow.MPIImplementationIntel ||
+			   implementation == kubeflow.MPIImplementationMPICH {
 				f.expectCreateServiceAction(newLauncherService(mpiJobCopy))
 			}
 			f.expectCreateJobAction(fmjc.newLauncherJob(mpiJobCopy))

v2/test/e2e/e2e_suite_test.go

@@ -45,6 +45,7 @@ const (
 	defaultKindImage        = "kindest/node:v1.21.2"
 	openMPIImage            = "mpioperator/mpi-pi:openmpi"
 	intelMPIImage           = "mpioperator/mpi-pi:intel"
+	mpichMPIImage           = "mpioperator/mpi-pi:mpich"
 	rootPath                = "../../.."
 	kubectlPath             = rootPath + "/bin/kubectl"
 	operatorManifestsPath   = rootPath + "/manifests/overlays/dev"