PR #3359: Make Docker build and upload scripts runnable through pip

[bvandermoon]

PR #3359: Make Docker build and upload scripts runnable through pip

Imported from GitHub PR #3359

Description

• Move the /dependencies directory to src/dependencies. This allows the docker build/upload scripts to be included in the pip package
• Create Python wrappers over docker_build_dependency_image.sh and docker_upload_runner.sh
• Create new scripts in pyproject.toml for these Python wrappers. This lets users run these scripts directly after pip installing MaxText

Next steps:

• Add new installation option to pyproject.toml so users can run these scripts without installing unecessary dependencies
• Update tutorials/docs that pip install maxtext to use the new script
• We should probably convert the shell scripts to Python scripts to get rid of the Python wrappers in the long-term. This is lower priority right now

Note: This changes the location the script is run from. We will update existing references we find and notify as many users as we can

FIXES: b/490520594

Tests

Setup with uv pip install -e .[tpu] --resolution=lowest

Built stable image through pip script + bash:

docker_build_dependency_image
bash src/dependencies/scripts/docker_build_dependency_image.sh

Built nightly image:

docker_build_dependency_image MODE=nightly

Ran successful workload for both builds with:

python3 -m benchmarks.benchmark_runner xpk \
    --project=$PROJECT \
    --zone=$ZONE \
    --device_type=v6e-256 \
    --num_slices=1 \
    --cluster_name=${CLUSTER_NAME} \
    --base_output_directory=${OUTPUT_DIR} \
    --model_name="llama3_1_70b_8192" \
    --base_docker_image=maxtext_base_image \
    --priority=high

docker_upload_runner CLOUD_IMAGE_NAME=<cloud_image_name>

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.

Copybara import of the project:

--
6ad79a9 by Branden Vandermoon bvandermoon@google.com:

Make Docker build and upload scripts runnable through pip

Merging this change closes #3359

FUTURE_COPYBARA_INTEGRATE_REVIEW=#3359 from AI-Hypercomputer:bvandermoon-docker-builds 6ad79a9

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!