Merge branch 'master' into support-for-gpu-cluster-using-oci-runner

Signed-off-by: Akash Jaiswal <akashjaiswal3846@gmail.com>

Author: jaiakash

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

@@ -4,9 +4,9 @@ on:
   push:
     branches:
       - master
-      - 'release-*'
+      - "release-*"
     tags:
-      - 'v*'
+      - "v*"
   pull_request:
 
 jobs:
@@ -34,9 +34,10 @@ jobs:
           - component-name: deepspeed-runtime
             dockerfile: cmd/runtimes/deepspeed/Dockerfile
             platforms: linux/amd64,linux/arm64
+          # TODO (andreyvelich): mlx[cuda] doesn't support arm at the moment: https://github.com/ml-explore/mlx/issues/2469
           - component-name: mlx-runtime
             dockerfile: cmd/runtimes/mlx/Dockerfile
-            platforms: linux/arm64
+            platforms: linux/amd64
           - component-name: torchtune-trainer
             dockerfile: cmd/trainers/torchtune/Dockerfile
             platforms: linux/amd64,linux/arm64

.github/workflows/test-go.yaml

@@ -8,6 +8,7 @@ jobs:
   generate:
     name: Generate
     runs-on: ubuntu-latest
+    if: ${{ github.repository == 'kubeflow/trainer' }}
     env:
       GOPATH: ${{ github.workspace }}/go
     defaults:
@@ -47,6 +48,7 @@ jobs:
   test:
     name: Test
     runs-on: ubuntu-latest
+    if: ${{ github.repository == 'kubeflow/trainer' }}
     env:
       GOPATH: ${{ github.workspace }}/go
     defaults:

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

@@ -3081,6 +3081,11 @@ spec:
                 x-kubernetes-list-type: map
             type: object
         type: object
+        x-kubernetes-validations:
+        - message: metadata.name must match RFC 1035 DNS label format
+          rule: self.metadata.name.matches('^[a-z]([-a-z0-9]*[a-z0-9])?$')
+        - message: metadata.name must be no more than 63 characters
+          rule: size(self.metadata.name) <= 63
     served: true
     storage: true
     subresources:

charts/kubeflow-trainer/templates/rbac/clusterrole.yaml

@@ -42,6 +42,14 @@ rules:
   - patch
   - update
   - watch
+- apiGroups:
+  - ""
+  resources:
+  - limitranges
+  verbs:
+  - get
+  - list
+  - watch
 - apiGroups:
   - admissionregistration.k8s.io
   resources:
@@ -62,6 +70,14 @@ rules:
   - patch
   - update
   - watch
+- apiGroups:
+  - node.k8s.io
+  resources:
+  - runtimeclasses
+  verbs:
+  - get
+  - list
+  - watch
 - apiGroups:
   - scheduling.x-k8s.io
   resources:

cmd/runtimes/deepspeed/Dockerfile

@@ -1,15 +1,11 @@
 FROM mpioperator/base:v0.6.0 AS mpi
 FROM nvidia/cuda:12.8.1-devel-ubuntu22.04
 
-# Disable interactive dialog from apt.
-ENV DEBIAN_FRONTEND noninteractive
-
-# Install libraries required for OpenMPI to work.
-RUN apt-get update && apt install -y --no-install-recommends \
-    cmake g++ gcc \
-    wget vim \
-    openssh-client openssh-server libcap2-bin \
-    libopenmpi-dev openmpi-bin
+# Install libraries required for OpenMPI to work. Image installs OpenMPI 5.0.7
+RUN apt update && apt install -y --no-install-recommends \
+    openssh-server openssh-client libcap2-bin \
+    g++ libopenmpi-dev \
+    python3-dev pip && rm -f /usr/bin/python && ln -s /usr/bin/python3 /usr/bin/python && rm -rf /var/lib/apt/lists/*
 
 # Add capability to run sshd as non-root.
 RUN setcap CAP_NET_BIND_SERVICE=+eip /usr/sbin/sshd
@@ -24,10 +20,7 @@ COPY --from=mpi /etc/ssh/ssh_config /etc/ssh/ssh_config
 COPY --from=mpi /etc/ssh/sshd_config /etc/ssh/sshd_config
 COPY --from=mpi /home/mpiuser/.sshd_config /home/mpiuser/.sshd_config
 
-# Install the required Python packages.
-RUN apt install -y python3-dev pip && rm -f /usr/bin/python && ln -s /usr/bin/python3 /usr/bin/python
-
-ENV LD_LIBRARY_PATH=/usr/local/lib:/usr/local/mpi/lib:/usr/local/mpi/lib64:${LD_LIBRARY_PATH}
+# Set home directory for mpiuser.
 ENV HOME=/home/mpiuser
 ENV PATH=$HOME/.local/bin:$PATH
 

cmd/runtimes/mlx/Dockerfile

@@ -1,14 +1,15 @@
 FROM mpioperator/base:v0.6.0 AS mpi
-FROM debian:trixie
+FROM nvidia/cuda:12.8.1-devel-ubuntu22.04
 
-# Install libraries required for OpenMPI and MLX. This image installs OpenMPI 5.0.7
+# Install libraries required for OpenMPI to work. Image installs OpenMPI 5.0.7
 RUN apt update && apt install -y --no-install-recommends \
     openssh-server openssh-client libcap2-bin \
-    libopenmpi-dev \
-    git g++ libblas-dev liblapack-dev liblapacke-dev
+    g++ libopenmpi-dev libblas-dev liblapack-dev liblapacke-dev \
+    python3-dev pip && rm -f /usr/bin/python && ln -s /usr/bin/python3 /usr/bin/python && rm -rf /var/lib/apt/lists/*
 
 # Add capability to run sshd as non-root.
 RUN setcap CAP_NET_BIND_SERVICE=+eip /usr/sbin/sshd
+RUN apt remove libcap2-bin -y
 
 # Configure mpiuser and home directory.
 RUN useradd -m mpiuser
@@ -19,12 +20,14 @@ COPY --from=mpi /etc/ssh/ssh_config /etc/ssh/ssh_config
 COPY --from=mpi /etc/ssh/sshd_config /etc/ssh/sshd_config
 COPY --from=mpi /home/mpiuser/.sshd_config /home/mpiuser/.sshd_config
 
-# Install the required Python packages. This image has Python 3.13
-RUN apt update && apt install -y python3 python3-pip && ln -s /usr/bin/python3 /usr/bin/python && apt clean
+# Set home directory for mpiuser.
+ENV HOME=/home/mpiuser
+ENV PATH=$HOME/.local/bin:$PATH
 
-# We have to build MLX and MLX Data from source.
-RUN git clone https://github.com/ml-explore/mlx.git
-RUN cd mlx && git checkout f018e248cd75dbb65668f418d6afb67842ea28b7 && CMAKE_BUILD_PARALLEL_LEVEL=8 pip install -v --break-system-packages .
+COPY cmd/runtimes/mlx/requirements.txt .
+RUN pip install --user -r requirements.txt
 
-RUN git clone https://github.com/ml-explore/mlx-data.git
-RUN cd mlx-data && git checkout 79516daa75aa3e9fd72fc5e3fb5e9e629912feac && CMAKE_BUILD_PARALLEL_LEVEL=8 pip install -v --break-system-packages .
+# Give mpiuser permission to download packages and HF models.
+# .cache directory is used by ML frameworks to download models.
+RUN chown -R mpiuser:mpiuser /home/mpiuser/.local
+RUN mkdir -p /home/mpiuser/.cache && chown -R mpiuser:mpiuser /home/mpiuser/.cache

cmd/runtimes/mlx/requirements.txt

@@ -0,0 +1,6 @@
+# MLX libraries.
+mlx[cuda]==0.28.0
+mlx-data==0.1.0
+mlx-lm==0.26.3
+# HuggingFace libraries.
+datasets==4.0.0

examples/deepspeed/text-summarization/T5-Fine-Tuning.ipynb

@@ -14,9 +14,7 @@
     "\n",
     "Wikihow dataset: https://huggingface.co/datasets/sentence-transformers/wikihow\n",
     "\n",
-    "This Notebook will use **4 x A100 NVIDIA GPUs**, to fine-tune T5 model on 2 nodes (every node has 2 GPUs).\n",
-    "\n",
-    "**TODO (andreyvelich)**: Currently, to run this Notebook you have to manualy update the container resources in the ClusterTrainingRuntime, since we don't propogate TrainJob's `resources_per_node` to the JobSet"
+    "This Notebook will use **4 x A100 NVIDIA GPUs**, to fine-tune T5 model on 2 nodes (every node has 2 GPUs)."
    ]
   },
   {
@@ -295,7 +293,7 @@
    "outputs": [],
    "source": [
     "MODEL_NAME = \"t5-base\"\n",
-    "BUCKET_NAME = \"TODO: add your bucket here\""
+    "# BUCKET_NAME = \"TODO: add your bucket here\""
    ]
   },
   {