Add colocated python Dockerfile and build upload script

[ycchenzheng]

Description

This PR is to add colocated python image related Dockerfile and build&upload script.

1. src/dependencies/dockerfiles/colocated_python.Dockerfile is added to install all necessary dependencies for base image us-docker.pkg.dev/cloud-tpu-v2-images/pathways-colocated-python/sidecar
2. src/dependencies/scripts/build_upload_colocated_python_image.sh will automatically pick base docker image from us-docker.pkg.dev/cloud-tpu-v2-images/pathways-colocated-python/sidecar that matches jax version in src/dependencies/requirements/generated_requirements/tpu-requirements.txt

Example usage:

# Example 1: Use defaults for PROJECT and LOCAL_IMAGE_NAME
# bash src/dependencies/scripts/build_upload_colocated_python_image.sh PROJECT=my-tpu-dev

# Example 2: Specify a custom project and local name
# bash src/dependencies/scripts/build_upload_colocated_python_image.sh PROJECT=my-tpu-dev LOCAL_IMAGE_NAME=maxtext-cp

# Example 3: Provide a complete IMAGE_LOCATION, overriding defaults
# bash src/dependencies/scripts/build_upload_colocated_python_image.sh IMAGE_LOCATION=us-docker.pkg.dev/my-artifact-repo/images/colocated-python:stable

# Example 4: Specify a custom JAX version
# bash src/dependencies/scripts/build_upload_colocated_python_image.sh PROJECT=my-tpu-dev JAX_VERSION=0.10.0

FIXES: b/515481092

Tests

Tested colocated python images built by this change with different TPU versions & topologies
Example logs

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.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[ycchenzheng]

@bvandermoon Can you please review this PR?

[SurbhiJainUSC]

Synced with @ycchenzheng offline and here are the next steps that we discussed:
This PR restricts users to run this script only by cloning Git repository. However, MaxText is only Pip installable, so we suggest:

1. Add a new flag and call src/dependencies/scripts/build_upload_colocated_python_image.sh behind that flag here: https://github.com/AI-Hypercomputer/maxtext/blob/main/src/dependencies/scripts/build_maxtext_docker_image.py#L27. This way, users can simply run build_maxtext_docker_image command with that new flag to build the docker image.
2. Modify build_upload_colocated_python_image.sh to only build docker image and use docker_upload_runner.sh to upload docker image to GCR.
3. Update https://github.com/AI-Hypercomputer/maxtext/blob/main/src/dependencies/scripts/upload_maxtext_docker_image.py to upload colocated python image to GCR.

[ycchenzheng]

Synced with @ycchenzheng offline and here are the next steps that we discussed: This PR restricts users to run this script only by cloning Git repository. However, MaxText is only Pip installable, so we suggest:

1. Add a new flag and call src/dependencies/scripts/build_upload_colocated_python_image.sh behind that flag here: https://github.com/AI-Hypercomputer/maxtext/blob/main/src/dependencies/scripts/build_maxtext_docker_image.py#L27. This way, users can simply run build_maxtext_docker_image command with that new flag to build the docker image.
2. Modify build_upload_colocated_python_image.sh to only build docker image and use docker_upload_runner.sh to upload docker image to GCR.
3. Update https://github.com/AI-Hypercomputer/maxtext/blob/main/src/dependencies/scripts/upload_maxtext_docker_image.py to upload colocated python image to GCR.

Tracking future implementation in b/528026019