Build MaxText docker image from PyPI

SurbhiJainUSC · SurbhiJainUSC · commit 8328cbdbdb21 · 2026-04-30T19:04:18.000Z
diff --git a/.github/workflows/UploadDockerImages.yml b/.github/workflows/UploadDockerImages.yml
@@ -1,4 +1,4 @@
-# Copyright 2023–2025 Google LLC
+# Copyright 2023–2026 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -42,27 +42,16 @@ permissions:
   contents: read
 
 jobs:
-  setup:
-    runs-on: ubuntu-latest
-    outputs:
-      maxtext_sha: ${{ steps.vars.outputs.maxtext_sha }}
-      image_date: ${{ steps.vars.outputs.image_date }}
-    steps:
-     - name: Checkout MaxText
-       uses: actions/checkout@v5
-
-     - name: Get metadata
-       id: vars
-       run: |
-        # MaxText SHA
-        echo "maxtext_sha=$(git rev-parse HEAD)" >> $GITHUB_OUTPUT
-
-        # Image date
-        echo "image_date=$(date +%Y-%m-%d)" >> $GITHUB_OUTPUT
+  build_and_upload_maxtext_package:
+    uses: ./.github/workflows/build_package.yml
+    with:
+      device_type: tpu
+      device_name: v4-8
+      cloud_runner: linux-x86-n2-16-buildkit
 
   build-and-push:
     name: ${{ matrix.image_name }}
-    needs: setup
+    needs: build_and_upload_maxtext_package
     strategy:
       fail-fast: false
       matrix:
@@ -71,36 +60,34 @@ jobs:
             build_mode: stable
             workflow: pre-training
             image_name: maxtext_jax_stable
-            dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
+            include_test_assets: true
           - device: tpu
             build_mode: nightly
             workflow: pre-training
             image_name: maxtext_jax_nightly
-            dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
+            include_test_assets: true
           - device: tpu
             build_mode: nightly
             workflow: post-training
             image_name: maxtext_post_training_nightly
-            dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
+            include_test_assets: true
           - device: gpu
             build_mode: stable
             workflow: pre-training
             image_name: maxtext_gpu_jax_stable
-            dockerfile: ./src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
+            include_test_assets: true
           - device: gpu
             build_mode: nightly
             workflow: pre-training
             image_name: maxtext_gpu_jax_nightly
-            dockerfile: ./src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
+            include_test_assets: true
     uses: ./.github/workflows/build_and_push_docker_image.yml
     with:
       image_name: ${{ matrix.image_name }}${{ inputs.image_suffix  }}
       device: ${{ matrix.device }}
       build_mode: ${{ matrix.build_mode }}
       workflow: ${{ matrix.workflow }}
-      dockerfile: ${{ matrix.dockerfile }}
-      maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }}
-      image_date: ${{ needs.setup.outputs.image_date }}
+      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
     secrets:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
 
@@ -131,7 +118,6 @@ jobs:
             build_mode: nightly
             workflow: pre-training
             image_name: maxtext_gpu_jax_nightly
-
     uses: ./.github/workflows/promote_docker_image.yml
     with:
       image_name: ${{ matrix.image_name }}${{ inputs.image_suffix  }}
diff --git a/.github/workflows/build_and_push_docker_image.yml b/.github/workflows/build_and_push_docker_image.yml
@@ -1,4 +1,4 @@
-# Copyright 2025 Google LLC
+# Copyright 2023-2026 Google LLC
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -28,15 +28,9 @@ on:
       build_mode:
         required: true
         type: string
-      dockerfile:
-        required: true
-        type: string
       maxtext_sha:
         required: true
         type: string
-      image_date:
-        required: false
-        type: string
       workflow:
         required: false
         type: string
@@ -45,6 +39,10 @@ on:
         required: false
         type: string
         default: ''
+      include_test_assets:
+        required: false
+        type: boolean
+        default: false
     secrets:
       HF_TOKEN:
         required: true
@@ -89,48 +87,57 @@ jobs:
           echo "workflow: ${{ inputs.workflow }}"
           echo "build_mode: ${{ inputs.build_mode }}"
           echo "image_name: ${{ inputs.image_name }}"
-          echo "dockerfile: ${{ inputs.dockerfile }}"
 
-      - name: Checkout MaxText
-        uses: actions/checkout@v5
+      - name: Configure Docker
+        run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io -q
         if: steps.check.outputs.should_run == 'true'
-        with:
-          # This ensures that every job clones the exact same commit as "setup" job
-          ref: ${{ inputs.maxtext_sha }}
 
-      - name: Mark git repositories as safe
-        run: git config --global --add safe.directory ${GITHUB_WORKSPACE}
-        if: steps.check.outputs.should_run == 'true'
+      #- 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: Configure Docker
-        run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io -q
+      - name: Download MaxText wheel
+        uses: actions/download-artifact@v4
         if: steps.check.outputs.should_run == 'true'
+        with:
+          name: maxtext-wheel
 
-      - name: Set up Docker BuildX
-        uses: docker/setup-buildx-action@v3.11.1
+      - name: Install uv and set the Python version
+        uses: astral-sh/setup-uv@v7
         if: steps.check.outputs.should_run == 'true'
         with:
-          driver: remote
-          endpoint: tcp://localhost:1234
+          python-version: '3.12'
+          enable-cache: true
+
+      - name: Install MaxText wheel
+        if: steps.check.outputs.should_run == 'true'
+        shell: bash
+        run: |
+          uv venv --seed
+          source .venv/bin/activate
+          maxtext_wheel=$(ls maxtext-*-py3-none-any.whl 2>/dev/null)
+          uv pip install ${maxtext_wheel}[runner] --resolution=lowest
 
       - name: Build and push Docker image
-        uses: docker/build-push-action@v6
         if: steps.check.outputs.should_run == 'true'
-        with:
-          push: true
-          context: .
-          file: ${{ inputs.dockerfile }}
-          tags: gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}:${{ github.run_id }}
-          cache-from: type=gha
-          outputs: type=image,compression=zstd,force-compression=true
-          build-args: |
-            DEVICE=${{ inputs.device }}
-            MODE=${{ inputs.build_mode }}
+        shell: bash
+        env:
+          IMAGE_TAG: "gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}:${{ github.run_id }}"
+          INCLUDE_TEST_ASSETS: ${{ inputs.include_test_assets }}
+        run: |
+          source .venv/bin/activate
+          echo "Building MaxText Docker image..."
+          sudo .venv/bin/python3 -m dependencies.scripts.build_maxtext_docker_image \
+            DEVICE=${{ inputs.device }} \
+            MODE=${{ inputs.build_mode }} \
             WORKFLOW=${{ inputs.workflow }}
-            PACKAGE_DIR=./src
-            JAX_VERSION=NONE
-            LIBTPU_VERSION=NONE
-            INCLUDE_TEST_ASSETS=true
+
+          echo "Pushing MaxText Docker image..."
+          docker push "${IMAGE_TAG}"
+          echo "Successfully pushed ${IMAGE_TAG}"
 
       - name: Add tags to Docker image
         if: steps.check.outputs.should_run == 'true'
@@ -146,16 +153,16 @@ jobs:
             echo "Tagging docker images corresponding to nightly release..."
 
             # Add date tag
-            gcloud container images add-tag "${TEMP_IMG}" "$SOURCE_IMAGE:${INPUTS_IMAGE_DATE}" --quiet
+            IMAGE_DATE="$(date +%Y-%m-%d)"
+            gcloud container images add-tag "${TEMP_IMG}" "$SOURCE_IMAGE:$IMAGE_DATE" --quiet
 
             # Convert date to YYYYMMDD format
-            clean_date=$(echo "${INPUTS_IMAGE_DATE}" | sed 's/[-:]//g' | cut -c1-8)
+            clean_date=$(echo "$IMAGE_DATE" | sed 's/[-:]//g' | cut -c1-8)
 
             # Add MaxText tag
             gcloud container images add-tag "${TEMP_IMG}" "${SOURCE_IMAGE}:maxtext_${MAXTEXT_SHA}_${clean_date}" --quiet
           fi
         env:
           INPUTS_IMAGE_NAME: ${{ inputs.image_name }}
-          INPUTS_IMAGE_DATE: ${{ inputs.image_date }}
           INPUTS_VERSION_NAME: ${{ inputs.version_name }}
           MAXTEXT_SHA: ${{ inputs.maxtext_sha }}
diff --git a/.github/workflows/pypi_release.yml b/.github/workflows/pypi_release.yml
@@ -1,4 +1,4 @@
-# Copyright 2025 Google LLC
+# Copyright 2023-2026 Google LLC
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -103,24 +103,20 @@ jobs:
             build_mode: stable
             image_name: maxtext_jax_stable
             workflow: pre-training
-            dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
           - device: gpu
             build_mode: stable
             image_name: maxtext_gpu_jax_stable
             workflow: pre-training
-            dockerfile: ./src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
           - device: tpu
             build_mode: stable
             image_name: maxtext_post_training_stable
             workflow: post-training
-            dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
     uses: ./.github/workflows/build_and_push_docker_image.yml
     with:
       image_name: ${{ matrix.image_name }}
       device: ${{ matrix.device }}
       build_mode: ${{ matrix.build_mode }}
       workflow: ${{ matrix.workflow }}
-      dockerfile: ${{ matrix.dockerfile }}
       maxtext_sha: ${{ github.sha }}
       version_name: ${{ needs.get_latest_maxtext_pypi_version.outputs.latest_pypi_version }}
     secrets:
