Merge pull request #2857 from AI-Hypercomputer:ci_image

Google-ML-Automation · Google-ML-Automation · commit 4e927f50607b · 2025-12-19T17:12:32.000-08:00
PiperOrigin-RevId: 846918748
diff --git a/.github/workflows/UploadDockerImages.yml b/.github/workflows/UploadDockerImages.yml
@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# This workflow will install Python dependencies, run tests and lint with a variety of Python versions
-# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python
+# This workflow builds and pushes MaxText images for both TPU and GPU devices.
+# It runs automatically daily at 12am UTC, on Pull Requests, or manually via Workflow Dispatch.
 
 name: Build Images
 
@@ -34,14 +34,8 @@ on:
           - gpu
 
 jobs:
-  build-tpu:
-    # This job will only run for 'tpu', 'all', schedule, or PR triggers.
-    if: >
-      github.event_name == 'schedule' ||
-      github.event_name == 'pull_request' ||
-      github.event.inputs.target_device == 'all' ||
-      github.event.inputs.target_device == 'tpu'
-    
+  build:
+    name: Build ${{ matrix.device }}-${{ matrix.build_mode }} Image
     runs-on: linux-x86-n2-16-buildkit
     container: google/cloud-sdk:524.0.0
 
@@ -51,127 +45,91 @@ jobs:
       matrix:
         include:
           # TPU Image Builds
-          - image_name: maxtext_jax_stable
+          - device: tpu
+            build_mode: stable
+            image_name: maxtext_jax_stable
             dockerfile: ./dependencies/dockerfiles/maxtext_dependencies.Dockerfile
-            build_args: |
-              MODE=stable
-              JAX_VERSION=NONE
-              LIBTPU_GCS_PATH=NONE
-          - image_name: maxtext_jax_nightly
+          - device: tpu
+            build_mode: nightly
+            image_name: maxtext_jax_nightly
             dockerfile: ./dependencies/dockerfiles/maxtext_dependencies.Dockerfile
-            build_args: |
-              MODE=nightly
-              JAX_VERSION=NONE
-              LIBTPU_GCS_PATH=NONE
-          # TPU Image builds using JAX AI Image
-          - image_name: maxtext_jax_stable_stack
-            dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile
-            base_image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
-          - image_name: maxtext_stable_stack_nightly_jax
-            dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile
-            base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/tpu/jax_nightly:latest
-          - image_name: maxtext_stable_stack_candidate
-            dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile
-            base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/tpu:latest
-
-    # Setup for GKE runners per b/412986220#comment82 and b/412986220#comment90
-    steps:
-      - uses: actions/checkout@v5
-      - name: Mark git repository as safe
-        run: git config --global --add safe.directory ${GITHUB_WORKSPACE}
-      - name: Configure Docker
-        run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io -q
-      - name: Set up Docker BuildX
-        uses: docker/setup-buildx-action@v3.11.1
-        with:
-          driver: remote
-          endpoint: tcp://localhost:1234
-      # Env variables to be passed to Dockerfile
-      - name: Get short commit hash
-        id: vars
-        run: echo "sha_short=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT
-      - name: Get current date
-        id: date
-        run: echo "image_date=$(date +%Y-%m-%d)" >> $GITHUB_OUTPUT
-      # Docker BuildX command config
-      - name: Build and Push Docker Image
-        uses: docker/build-push-action@v6
-        with:
-          push: true
-          context: .
-          file: ${{ matrix.dockerfile }}
-          tags: |
-            gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:${{ steps.date.outputs.image_date }}
-            gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:latest
-          cache-from: type=gha
-          cache-to: type=gha,mode=max
-          provenance: false
-          build-args: |
-            ${{ matrix.build_args }}
-            JAX_AI_IMAGE_BASEIMAGE=${{ matrix.base_image }}
-            COMMIT_HASH=${{ steps.vars.outputs.sha_short }}
-            DEVICE=tpu
-            TEST_TYPE=xlml
+          # GPU Image Builds
+          - device: gpu
+            build_mode: stable
+            image_name: maxtext_gpu_jax_stable
+            dockerfile: ./dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
+          - device: gpu
+            build_mode: nightly
+            image_name: maxtext_gpu_jax_nightly
+            dockerfile: ./dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
 
-  # Same as tpu-build step but mirrored for GPUs
-  build-gpu:
     if: >
       github.event_name == 'schedule' ||
       github.event_name == 'pull_request' ||
-      github.event.inputs.target_device == 'all' ||
-      github.event.inputs.target_device == 'gpu'
+      github.event_name == 'workflow_dispatch' && (
+        github.event.inputs.target_device == 'all' ||
+        github.event.inputs.target_device == 'tpu' ||
+        github.event.inputs.target_device == 'gpu'
+      )
 
-    runs-on: linux-x86-n2-16-buildkit
-    container: google/cloud-sdk:524.0.0
+    # Setup for GKE runners per b/412986220#comment82 and b/412986220#comment90
+    steps:
+      - name: Check if build should run
+        id: check
+        shell: bash
+        run: |
+          if [[ "${{ github.event_name }}" == "workflow_dispatch" && "${{ github.event.inputs.target_device }}" != "all" && "${{ github.event.inputs.target_device }}" != "${{ matrix.device }}" ]]; then
+            echo "should_run=false" >> $GITHUB_OUTPUT
+            echo "Skipping build for device: ${{ matrix.device }} in ${{ matrix.build_mode }} mode."
+          else
+            echo "should_run=true" >> $GITHUB_OUTPUT
+            echo "Building for device: ${{ matrix.device }} in ${{ matrix.build_mode }} mode."
+          fi
 
-    strategy:
-      fail-fast: false
-      matrix:
-        # GPU Image Builds using JAX AI Image
-        include:
-          - image_name: maxtext_gpu_jax_stable_stack
-            dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile
-            base_image: us-central1-docker.pkg.dev/deeplearning-images/jax-ai-image/gpu:latest
-          - image_name: maxtext_gpu_stable_stack_nightly_jax
-            dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile
-            base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/gpu/jax_nightly:latest
-          - image_name: maxtext_stable_stack_candidate_gpu
-            dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile
-            base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/gpu:latest
+      - name: Checkout git repository
+        uses: actions/checkout@v5
+        if: steps.check.outputs.should_run == 'true'
 
-    steps:
-      - uses: actions/checkout@v5
       - name: Mark git repository as safe
+        if: steps.check.outputs.should_run == 'true'
         run: git config --global --add safe.directory ${GITHUB_WORKSPACE}
+
       - name: Configure Docker
-        run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io,us-central1-docker.pkg.dev -q
+        if: steps.check.outputs.should_run == 'true'
+        run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io -q
+
       - name: Set up Docker BuildX
         uses: docker/setup-buildx-action@v3.11.1
+        if: steps.check.outputs.should_run == 'true'
         with:
           driver: remote
           endpoint: tcp://localhost:1234
-      - name: Get short commit hash
+
+      # Env variables to be passed to Dockerfile
+      - name: Get metadata
         id: vars
-        run: echo "sha_short=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT
-      - name: Get current date
-        id: date
-        run: echo "image_date=$(date +%Y-%m-%d)" >> $GITHUB_OUTPUT
+        if: steps.check.outputs.should_run == 'true'
+        run: |
+          echo "commit_hash=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT
+          echo "image_date=$(date +%Y-%m-%d)" >> $GITHUB_OUTPUT
+
+      # Docker BuildX command config
       - name: Build and Push Docker Image
         uses: docker/build-push-action@v6
+        if: steps.check.outputs.should_run == 'true'
         with:
           push: true
           context: .
           file: ${{ matrix.dockerfile }}
           tags: |
-            gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:maxtext_${{ steps.vars.outputs.sha_short }}
-            gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:${{ steps.date.outputs.image_date }}
+            gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:maxtext_${{ steps.vars.outputs.commit_hash }}
+            gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:${{ steps.vars.outputs.image_date }}
             gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:latest
           cache-from: type=gha
           cache-to: type=gha,mode=max
           provenance: false
           build-args: |
-            ${{ matrix.build_args }}
-            JAX_AI_IMAGE_BASEIMAGE=${{ matrix.base_image }}
-            COMMIT_HASH=${{ steps.vars.outputs.sha_short }}
-            DEVICE=gpu
-            TEST_TYPE=xlml
+            DEVICE=${{ matrix.device }}
+            MODE=${{ matrix.build_mode }}
+            JAX_VERSION=NONE
+            LIBTPU_GCS_PATH=NONE
diff --git a/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile b/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
@@ -2,6 +2,11 @@
 ARG BASEIMAGE=ghcr.io/nvidia/jax:base
 FROM $BASEIMAGE
 
+# Move the 'EXTERNALLY-MANAGED' file to allow system-wide pip installs
+RUN if [ -f /usr/lib/python3.12/EXTERNALLY-MANAGED ]; then \
+    mv /usr/lib/python3.12/EXTERNALLY-MANAGED /usr/lib/python3.12/EXTERNALLY-MANAGED.old; \
+fi
+
 # Stopgaps measure to circumvent gpg key setup issue.
 RUN echo "deb [trusted=yes] https://developer.download.nvidia.com/devtools/repos/ubuntu2204/amd64/ /" > /etc/apt/sources.list.d/devtools-ubuntu2204-amd64.list
 
