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

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

@@ -44,6 +44,9 @@ jobs:
           - component-name: data-cache
             dockerfile: cmd/data_cache/Dockerfile
             platforms: linux/amd64,linux/arm64
+          - component-name: jax-runtime
+            dockerfile: cmd/runtimes/jax/Dockerfile
+            platforms: linux/amd64,linux/arm64
     steps:
       - name: Checkout
         uses: actions/checkout@v6

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,54 @@
+FROM ghcr.io/nvidia/jax:jax as gpu-base
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+        build-essential \
+        cmake \
+        git \
+        curl \
+        libgoogle-glog-dev \
+        libgflags-dev \
+        libprotobuf-dev \
+        protobuf-compiler \
+        python3-dev \
+        python3-pip \
+        python3-setuptools && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir \
+        numpy \
+        jax \
+        jaxlib \
+        "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+
+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
+
+RUN pip install --no-cache-dir absl-py kubernetes
+
+FROM gloo-base as production
+
+WORKDIR /workspace
+
+COPY cmd/runtimes/jax/version.sh /version.sh
+RUN chmod +x /version.sh
+
+ENTRYPOINT ["/bin/bash", "-c", "/version.sh && exec \"$@\"", "--"]
+CMD ["/bin/bash"]

cmd/runtimes/jax/version.sh

@@ -0,0 +1,51 @@
+#!/bin/bash
+
+echo "=== VERSION INFORMATION ==="
+echo "Python version: $(python3 --version)"
+echo "Pip version: $(pip --version)"
+echo ""
+
+echo "=== JAX Information ==="
+python3 -c "import jax; print(f'JAX version: {jax.__version__}')"
+python3 -c "import jaxlib; print(f'JAXLib version: {jaxlib.__version__}')"
+echo ""
+
+echo "=== NCCL Information ==="
+if dpkg -l | grep -q libnccl2; then
+    echo "NCCL version: $(dpkg -l | grep libnccl2 | awk '{print $3}')"
+else
+    echo "NCCL: Not available via package manager"
+fi
+if command -v nvcc &> /dev/null; then
+    echo "CUDA version: $(nvcc --version | grep 'release' | awk '{print $5}')"
+fi
+echo ""
+
+echo "=== TPU Information ==="
+python3 -c "
+try:
+    import jax
+    devices = jax.devices()
+    tpu_devices = [d for d in devices if d.platform == 'tpu']
+    if tpu_devices:
+        print(f'TPU devices found: {len(tpu_devices)}')
+        for d in tpu_devices:
+            print(f'  - {d}')
+    else:
+        print('No TPU devices found')
+except Exception as e:
+    print(f'Error checking TPU devices: {e}')
+"
+echo ""
+
+echo "=== Gloo Information ==="
+echo "Gloo: Built from source (commit 43b7acbf372cdce14075f3526e39153b7e433b53)"
+if [ -f /usr/local/lib/libgloo.a ]; then
+    echo "Gloo library: Found at /usr/local/lib/libgloo.so"
+else
+    echo "Gloo library: Not found"
+fi
+echo ""
+
+echo "=== Installed Python Packages ==="
+pip list

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.