Add JAX controller

.github/workflows/publish-example-images.yaml

@@ -69,3 +69,7 @@ jobs:
             platforms: linux/amd64,linux/arm64
             dockerfile: examples/pytorch/mnist/Dockerfile-mpi
             context: examples/pytorch/mnist
+          - component-name: jaxjob-simple
+            platforms: linux/amd64,linux/arm64
+            dockerfile: examples/jax/cpu-demo/Dockerfile
+            context: examples/jax/cpu-demo

PROJECT

@@ -27,4 +27,12 @@ resources:
   kind: TFJob
   path: github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1
   version: v1
+- api:
+    crdVersion: v1
+    namespaced: true
+  controller: true
+  group: kubeflow.org
+  kind: JAXJob
+  path: github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1
+  version: v1
 version: "3"

cmd/training-operator.v1/main.go

@@ -87,7 +87,7 @@ func main() {
 			"Enabling this will ensure there is only one active controller manager.")
 	flag.StringVar(&leaderElectionID, "leader-election-id", "1ca428e5.training-operator.kubeflow.org", "The ID for leader election.")
 	flag.Var(&enabledSchemes, "enable-scheme", "Enable scheme(s) as --enable-scheme=tfjob --enable-scheme=pytorchjob, case insensitive."+
-		" Now supporting TFJob, PyTorchJob, XGBoostJob, PaddleJob. By default, all supported schemes will be enabled.")
+		" Now supporting TFJob, PyTorchJob, XGBoostJob, PaddleJob, JAXJob. By default, all supported schemes will be enabled.")
 	flag.StringVar(&gangSchedulerName, "gang-scheduler-name", "", "Now Supporting volcano and scheduler-plugins."+
 		" Note: If you set another scheduler name, the training-operator assumes it's the scheduler-plugins.")
 	flag.StringVar(&namespace, "namespace", os.Getenv(EnvKubeflowNamespace), "The namespace to monitor kubeflow jobs. If unset, it monitors all namespaces cluster-wide."+

examples/jax/cpu-demo/Dockerfile

@@ -0,0 +1,25 @@
+FROM python:3.12
+
+RUN pip install jax absl-py kubernetes
+
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    cmake \
+    git \
+    libgoogle-glog-dev \
+    libgflags-dev \
+    libprotobuf-dev \
+    protobuf-compiler \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN git clone https://github.com/facebookincubator/gloo.git \
+    && cd gloo \
+    && mkdir build \
+    && cd build \
+    && cmake ../ \
+    && make \
+    && make install
+
+WORKDIR /app
+
+ADD train.py /app

examples/jax/cpu-demo/demo.yaml

@@ -0,0 +1,19 @@
+apiVersion: "kubeflow.org/v1"
+kind: JAXJob
+metadata:
+  name: jaxjob-simple
+  namespace: kubeflow
+spec:
+  jaxReplicaSpecs:
+    Worker:
+      replicas: 2
+      restartPolicy: OnFailure
+      template:
+        spec:
+          containers:
+            - name: jax
+              image: docker.io/sandipanify/jaxgoogle:latest
+              command:
+                - "python3"
+                - "train.py"
+              imagePullPolicy: Always

examples/jax/cpu-demo/train.py

@@ -0,0 +1,43 @@
+# example ref:
+# https://jax.readthedocs.io/en/latest/multi_process.html#running-multi-process-computations
+# https://github.com/GoogleCloudPlatform/ai-on-gke/blob/main/tutorials-and-examples/gpu-examples/a100-jax/train.py # noqa
+
+import os
+import socket
+
+import jax
+from absl import app
+
+jax.config.update("jax_cpu_collectives_implementation", "gloo")
+
+
+def _main(argv):
+
+    process_id = int(os.getenv("PROCESS_ID"))
+    num_processes = int(os.getenv("NUM_PROCESSES"))
+    coordinator_address = os.getenv("COORDINATOR_ADDRESS")
+    coordinator_port = int(os.getenv("COORDINATOR_PORT"))
+    coordinator_address = f"{coordinator_address}:{coordinator_port}"
+
+    jax.distributed.initialize(
+        coordinator_address=coordinator_address,
+        num_processes=num_processes,
+        process_id=process_id,
+    )
+
+    print(
+        f"JAX process {jax.process_index()}/{jax.process_count() - 1} initialized on "
+        f"{socket.gethostname()}"
+    )
+    print(f"JAX global devices:{jax.devices()}")
+    print(f"JAX local devices:{jax.local_devices()}")
+
+    print(f"JAX device count:{jax.device_count()}")
+    print(f"JAX local device count:{jax.local_device_count()}")
+
+    xs = jax.numpy.ones(jax.local_device_count())
+    print(jax.pmap(lambda x: jax.lax.psum(x, "i"), axis_name="i")(xs))
+
+
+if __name__ == "__main__":
+    app.run(_main)

manifests/base/crds/kustomization.yaml

@@ -6,3 +6,4 @@ resources:
   - kubeflow.org_xgboostjobs.yaml
   - kubeflow.org_mpijobs.yaml
   - kubeflow.org_paddlejobs.yaml
+  - kubeflow.org_jaxjobs.yaml

manifests/base/rbac/role.yaml

@@ -92,6 +92,32 @@ rules:
   - patch
   - update
   - watch
+- apiGroups:
+  - kubeflow.org
+  resources:
+  - jaxjobs
+  verbs:
+  - create
+  - delete
+  - get
+  - list
+  - patch
+  - update
+  - watch
+- apiGroups:
+  - kubeflow.org
+  resources:
+  - jaxjobs/finalizers
+  verbs:
+  - update
+- apiGroups:
+  - kubeflow.org
+  resources:
+  - jaxjobs/status
+  verbs:
+  - get
+  - patch
+  - update
 - apiGroups:
   - kubeflow.org
   resources:

manifests/base/webhook/manifests.yaml

@@ -4,6 +4,26 @@ kind: ValidatingWebhookConfiguration
 metadata:
   name: validating-webhook-configuration
 webhooks:
+- admissionReviewVersions:
+  - v1
+  clientConfig:
+    service:
+      name: webhook-service
+      namespace: system
+      path: /validate-kubeflow-org-v1-jaxjob
+  failurePolicy: Fail
+  name: validator.jaxjob.training-operator.kubeflow.org
+  rules:
+  - apiGroups:
+    - kubeflow.org
+    apiVersions:
+    - v1
+    operations:
+    - CREATE
+    - UPDATE
+    resources:
+    - jaxjobs
+  sideEffects: None
 - admissionReviewVersions:
   - v1
   clientConfig:

manifests/base/webhook/patch.yaml

@@ -10,6 +10,9 @@
 - op: replace
   path: /webhooks/3/clientConfig/service/name
   value: training-operator
+- op: replace
+  path: /webhooks/4/clientConfig/service/name
+  value: training-operator
 - op: replace
   path: /metadata/name
   value: validator.training-operator.kubeflow.org

pkg/controller.v1/jax/envvar.go

@@ -0,0 +1,107 @@
+// Copyright 2024 The Kubeflow Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License
+
+package jax
+
+import (
+	"errors"
+	"strconv"
+	"strings"
+
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/utils/ptr"
+
+	kubeflowv1 "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1"
+)
+
+var (
+	errorDefaulContainerPortNotExposed = errors.New("default container port is not exposed")
+	errorFailedToRecognizeRank         = errors.New("failed to recognize the JAXJob Rank")
+)
+
+type EnvVarGenerator interface {
+	Generate(job *kubeflowv1.JAXJob) ([]corev1.EnvVar, error)
+}
+
+func setPodEnv(jaxjob *kubeflowv1.JAXJob, podTemplateSpec *corev1.PodTemplateSpec, rtype, index string) error {
+
+	coordinatorAddr := replicaName(jaxjob.Name, kubeflowv1.JAXJobReplicaTypeWorker, 0)
+
+	coordinatorPort, err := getPortFromJAXJob(jaxjob, kubeflowv1.JAXJobReplicaTypeWorker)
+	if err != nil {
+		return err
+	}
+
+	totalReplicas := getTotalReplicas(jaxjob)
+
+	for i := range podTemplateSpec.Spec.Containers {
+
+		rank, err := strconv.Atoi(index)
+		if err != nil {
+			return errorFailedToRecognizeRank
+		}
+		// Set PYTHONUNBUFFERED to true, to disable output buffering.
+		// Ref https://stackoverflow.com/questions/59812009/what-is-the-use-of-pythonunbuffered-in-docker-file.
+		podTemplateSpec.Spec.Containers[i].Env = append(podTemplateSpec.Spec.Containers[i].Env, corev1.EnvVar{
+			Name:  "PYTHONUNBUFFERED",
+			Value: "1",
+		})
+		podTemplateSpec.Spec.Containers[i].Env = append(podTemplateSpec.Spec.Containers[i].Env, corev1.EnvVar{
+			Name:  "COORDINATOR_PORT",
+			Value: strconv.Itoa(int(coordinatorPort)),
+		})
+		podTemplateSpec.Spec.Containers[i].Env = append(podTemplateSpec.Spec.Containers[i].Env, corev1.EnvVar{
+			Name:  "COORDINATOR_ADDRESS",
+			Value: coordinatorAddr,
+		})
+		podTemplateSpec.Spec.Containers[i].Env = append(podTemplateSpec.Spec.Containers[i].Env, corev1.EnvVar{
+			Name:  "NUM_PROCESSES",
+			Value: strconv.Itoa(int(totalReplicas)),
+		})
+		podTemplateSpec.Spec.Containers[i].Env = append(podTemplateSpec.Spec.Containers[i].Env, corev1.EnvVar{
+			Name:  "PROCESS_ID",
+			Value: strconv.Itoa(rank),
+		})
+	}
+
+	return nil
+}
+
+func getTotalReplicas(job *kubeflowv1.JAXJob) int {
+	jobReplicas := 0
+	for _, r := range job.Spec.JAXReplicaSpecs {
+		jobReplicas += int(ptr.Deref[int32](r.Replicas, 0))
+	}
+	return jobReplicas
+}
+
+func replicaName(jobName string, rtype kubeflowv1.ReplicaType, index int) string {
+	n := jobName + "-" + strings.ToLower(string(rtype)) + "-" + strconv.Itoa(index)
+	return strings.Replace(n, "/", "-", -1)
+}
+
+func getPortFromJAXJob(job *kubeflowv1.JAXJob, rtype kubeflowv1.ReplicaType) (int32, error) {
+	containers := job.Spec.JAXReplicaSpecs[rtype].Template.Spec.Containers
+	for _, container := range containers {
+		if container.Name == kubeflowv1.JAXJobDefaultContainerName {
+			ports := container.Ports
+			for _, port := range ports {
+				if port.Name == kubeflowv1.JAXJobDefaultPortName {
+					return port.ContainerPort, nil
+				}
+			}
+		}
+	}
+	return -1, errorDefaulContainerPortNotExposed
+}