Include test_assets from GCS in tpu/gpu dockerfiles

[SurbhiJainUSC]

Description

• Include test_assets from GCS in TPU/GPU docker images build nightly

FIXES: b/471848644

Tests

https://github.com/AI-Hypercomputer/maxtext/actions/runs/20583769072/job/59116469871

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[RissyRan]

Shall we also apply this to nightly image? maxtext_jax_nightly

[shuningjin]

It seems both "maxtext_jax_nightly" and "maxtext_jax_stable" are using this yml

maxtext/.github/workflows/UploadDockerImages.yml

Lines 65 to 73 in 5595c85

	- device: tpu
	build_mode: stable
	image_name: maxtext_jax_stable
	dockerfile: ./dependencies/dockerfiles/maxtext_dependencies.Dockerfile
	- device: tpu
	build_mode: nightly
	image_name: maxtext_jax_nightly
	dockerfile: ./dependencies/dockerfiles/maxtext_dependencies.Dockerfile
	uses: ./.github/workflows/build_and_push_docker_image.yml

where include_test_assests=true and base is "maxtext_jax_stable" from a certain date

maxtext/.github/workflows/build_and_push_docker_image.yml

Lines 117 to 118 in 5595c85

	INCLUDE_TEST_ASSETS=true
	BASEIMAGE=gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:${{ inputs.image_date }}

[SurbhiJainUSC]

Yes, that's right Shuning.

[shuningjin]

Thank you for the quick fix!