feat(runtimes): KEP-2442-jax-runtime

.github/workflows/build-and-push-images.yaml

@@ -38,6 +38,9 @@ jobs:
           - component-name: mlx-runtime
             dockerfile: cmd/runtimes/mlx/Dockerfile
             platforms: linux/amd64
+          - component-name: jax-runtime
+            dockerfile: cmd/runtimes/jax/Dockerfile
+            platforms: linux/amd64,linux/arm64
           - component-name: torchtune-trainer
             dockerfile: cmd/trainers/torchtune/Dockerfile
             platforms: linux/amd64,linux/arm64

charts/kubeflow-trainer/crds/trainer.kubeflow.org_clustertrainingruntimes.yaml

@@ -47,6 +47,9 @@ spec:
                 description: mlPolicy provides the ML-specific parameters for the
                   model training.
                 properties:
+                  jax:
+                    description: jax defines the configuration for the JAX Runtime
+                    type: object
                   mpi:
                     description: mpi defines the configuration for the MPI Runtime.
                     properties:

charts/kubeflow-trainer/crds/trainer.kubeflow.org_trainingruntimes.yaml

@@ -47,6 +47,9 @@ spec:
                 description: mlPolicy provides the ML-specific parameters for the
                   model training.
                 properties:
+                  jax:
+                    description: jax defines the configuration for the JAX Runtime
+                    type: object
                   mpi:
                     description: mpi defines the configuration for the MPI Runtime.
                     properties:

cmd/runtimes/jax/Dockerfile

@@ -0,0 +1,36 @@
+FROM ghcr.io/nvidia/jax:jax-2026-01-04 as gpu-base
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt update && apt install -y --no-install-recommends \
+        build-essential \
+        cmake \
+        git \
+        libgoogle-glog-dev \
+        libgflags-dev \
+        libprotobuf-dev \
+        protobuf-compiler \
+        python3-dev pip && rm -f /usr/bin/python &&  \
+    ln -s /usr/bin/python3 /usr/bin/python && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN pip install --no-cache-dir --upgrade pip
+
+FROM gpu-base as tpu-base
+
+RUN pip install --no-cache-dir "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html || \
+    echo "TPU support not available" && \
+    pip install --no-cache-dir libtpu-nightly || \
+    echo "libtpu-nightly not available"
+
+FROM tpu-base as gloo-base
+
+RUN git clone https://github.com/facebookincubator/gloo.git \
+    && cd gloo \
+    && git checkout 43b7acbf372cdce14075f3526e39153b7e433b53 \
+    && mkdir build \
+    && cd build \
+    && cmake ../ \
+    && make \
+    && make install
+
+FROM gloo-base as production

manifests/base/crds/trainer.kubeflow.org_clustertrainingruntimes.yaml

@@ -47,6 +47,9 @@ spec:
                 description: mlPolicy provides the ML-specific parameters for the
                   model training.
                 properties:
+                  jax:
+                    description: jax defines the configuration for the JAX Runtime
+                    type: object
                   mpi:
                     description: mpi defines the configuration for the MPI Runtime.
                     properties:

manifests/base/crds/trainer.kubeflow.org_trainingruntimes.yaml

@@ -47,6 +47,9 @@ spec:
                 description: mlPolicy provides the ML-specific parameters for the
                   model training.
                 properties:
+                  jax:
+                    description: jax defines the configuration for the JAX Runtime
+                    type: object
                   mpi:
                     description: mpi defines the configuration for the MPI Runtime.
                     properties:

manifests/base/runtimes/jax_distributed.yaml

@@ -0,0 +1,24 @@
+apiVersion: trainer.kubeflow.org/v1alpha1
+kind: ClusterTrainingRuntime
+metadata:
+  name: jax-distributed
+  labels:
+    trainer.kubeflow.org/framework: jax
+spec:
+  mlPolicy:
+    numNodes: 1
+    jax: {}
+  template:
+    spec:
+      replicatedJobs:
+        - name: node
+          template:
+            metadata:
+              labels:
+                trainer.kubeflow.org/trainjob-ancestor-step: trainer
+            spec:
+              template:
+                spec:
+                  containers:
+                    - name: node
+                      image: ghcr.io/kubeflow/trainer/jax-runtime

manifests/base/runtimes/kustomization.yaml

@@ -3,5 +3,6 @@ kind: Kustomization
 resources:
   - deepspeed_distributed.yaml
   - mlx_distributed.yaml
+  - jax_distributed.yaml
   - torch_distributed.yaml
   - torchtune

manifests/overlays/runtimes/kustomization.yaml

@@ -15,3 +15,5 @@ images:
     newTag: latest
   - name: ghcr.io/kubeflow/trainer/deepspeed-runtime
     newTag: latest
+  - name: ghcr.io/kubeflow/trainer/jax-runtime
+    newTag: latest

pkg/apis/trainer/v1alpha1/trainingruntime_types.go

@@ -194,6 +194,10 @@ type MLPolicySource struct {
 	// mpi defines the configuration for the MPI Runtime.
 	// +optional
 	MPI *MPIMLPolicySource `json:"mpi,omitempty"`
+
+	// jax defines the configuration for the JAX Runtime
+	// +optional
+	JAX *JAXMLPolicySource `json:"jax,omitempty"`
 }
 
 // TorchMLPolicySource represents a PyTorch runtime configuration.
@@ -239,6 +243,9 @@ type TorchElasticPolicy struct {
 	Metrics []autoscalingv2.MetricSpec `json:"metrics,omitempty"`
 }
 
+// JAXMLPolicySource represents a jax runtime configuration.
+type JAXMLPolicySource struct{}
+
 // MPIMLPolicySource represents a MPI runtime configuration.
 type MPIMLPolicySource struct {
 	// numProcPerNode is the number of processes per node.