Build MaxText docker image from PyPI

SurbhiJainUSC · SurbhiJainUSC · commit a8bc6aee279e · 2026-04-30T19:45:19.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  }}
+      image_name: ${{ inputs.image_suffix != '' && format('{0}_{1}', matrix.image_name, inputs.image_suffix) || matrix.image_name }}
       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
@@ -54,7 +52,7 @@ permissions:
 
 jobs:
   build_and_push:
-    runs-on: linux-x86-n2-16-buildkit
+    runs-on: ubuntu-latest
     container: google/cloud-sdk:524.0.0
     if: >
       github.event_name == 'release' ||
@@ -89,18 +87,6 @@ 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
-        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: Configure Docker
         run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io -q
@@ -113,24 +99,51 @@ jobs:
           driver: remote
           endpoint: tcp://localhost:1234
 
-      - name: Build and push Docker image
-        uses: docker/build-push-action@v6
+      - name: Download MaxText wheel
+        uses: actions/download-artifact@v4
         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 }}
+          name: maxtext-wheel
+
+      - name: Install uv and set the Python version
+        uses: astral-sh/setup-uv@v7
+        if: steps.check.outputs.should_run == 'true'
+        with:
+          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
+        if: steps.check.outputs.should_run == 'true'
+        shell: bash
+        env:
+          IMAGE_TAG: "gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}:${{ github.run_id }}"
+          INCLUDE_TEST_ASSETS: ${{ inputs.include_test_assets }}
+          PROJECT: tpu-prod-env-multipod
+        run: |
+          source .venv/bin/activate
+          echo "Building MaxText Docker image..."
+          #apt-get update && apt-get install -y sudo
+          #sudo .venv/bin/python3 -m dependencies.scripts.build_maxtext_docker_image \
+          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..."
+          #sudo .venv/bin/python3 -m dependencies.scripts.upload_maxtext_docker_image \
+          upload_maxtext_docker_image \
+            CLOUD_IMAGE_NAME=${{ inputs.image_name }}
+          docker push "${IMAGE_TAG}"
+          echo "Successfully pushed ${IMAGE_TAG}"
 
       - name: Add tags to Docker image
         if: steps.check.outputs.should_run == 'true'
@@ -146,16 +159,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:
