Merge pull request #2898 from AI-Hypercomputer:tests_docker_image

PiperOrigin-RevId: 850456883

Author: Google-ML-Automation

.github/workflows/UploadDockerImages.yml

@@ -102,6 +102,7 @@ jobs:
       dockerfile: ${{ matrix.dockerfile }}
       maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }}
       image_date: ${{ needs.setup.outputs.image_date }}
+      base_image: gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:${{ needs.setup.outputs.image_date }}
       is_post_training: true
 
   gpu-pre-training:

.github/workflows/build_and_push_docker_image.yml

@@ -37,6 +37,10 @@ on:
       image_date:
         required: true
         type: string
+      base_image:
+        required: false
+        type: string
+        default: ''
       is_post_training:
         required: false
         type: boolean
@@ -114,7 +118,8 @@ jobs:
             MODE=${{ inputs.build_mode }}
             JAX_VERSION=NONE
             LIBTPU_VERSION=NONE
-            BASEIMAGE=gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:${{ inputs.image_date }}
+            INCLUDE_TEST_ASSETS=true
+            ${{ inputs.base_image != '' && format('BASEIMAGE={0}', inputs.base_image) || '' }}
 
       - name: Add tags to Docker image
         if: steps.check.outputs.should_run == 'true'

dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile

@@ -58,5 +58,14 @@ RUN --mount=type=cache,target=/root/.cache/pip bash /deps/tools/setup/setup.sh M
 # Now copy the remaining code (source files that may change frequently)
 COPY . .
 
+# Download test assets from GCS if building image with test assets
+ARG INCLUDE_TEST_ASSETS=false
+RUN if [ "$INCLUDE_TEST_ASSETS" = "true" ]; then \
+        echo "Downloading test assets from GCS..."; \
+        if ! gcloud storage cp -r gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_ROOT}"; then \
+        echo "WARNING: Failed to download test assets from GCS. These files are only used for end-to-end tests; you may not have access to the bucket."; \
+        fi; \
+    fi
+
 # Install (editable) MaxText
 RUN test -f '/tmp/venv_created' && "$(tail -n1 /tmp/venv_created)"/bin/activate ; pip install --no-dependencies -e .

dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile

@@ -55,5 +55,14 @@ RUN --mount=type=cache,target=/root/.cache/pip bash /deps/tools/setup/setup.sh M
 # Now copy the remaining code (source files that may change frequently)
 COPY . .
 
+# Download test assets from GCS if building image with test assets
+ARG INCLUDE_TEST_ASSETS=false
+RUN if [ "$INCLUDE_TEST_ASSETS" = "true" ]; then \
+        echo "Downloading test assets from GCS..."; \
+        if ! gcloud storage cp -r gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_ROOT}"; then \
+        echo "WARNING: Failed to download test assets from GCS. These files are only used for end-to-end tests; you may not have access to the bucket."; \
+        fi; \
+    fi
+
 # Install (editable) MaxText
 RUN test -f '/tmp/venv_created' && "$(tail -n1 /tmp/venv_created)"/bin/activate ; pip install --no-dependencies -e .