From 6ad79a9bad9bf7112d2f7732cbafcdb55292ffbd Mon Sep 17 00:00:00 2001 From: Branden Vandermoon Date: Tue, 10 Mar 2026 01:04:29 +0000 Subject: [PATCH] Make Docker build and upload scripts runnable through pip --- .github/workflows/UploadDockerImages.yml | 12 ++++----- .github/workflows/build_and_upload_images.sh | 4 +-- .github/workflows/check_docs_build.yml | 2 +- .readthedocs.yml | 2 +- PREFLIGHT.md | 2 +- build_hooks.py | 2 +- docs/development.md | 2 +- docs/install_maxtext.md | 4 +-- docs/run_maxtext/run_maxtext_via_pathways.md | 4 +-- docs/run_maxtext/run_maxtext_via_xpk.md | 4 +-- docs/tutorials/first_run.md | 2 +- .../posttraining/rl_on_multi_host.md | 8 +++--- .../posttraining/sft_on_multi_host.md | 4 +-- pyproject.toml | 12 +++++---- .../dockerfiles/clean_py_env.Dockerfile | 0 .../dockerfiles/clean_py_env_gpu.Dockerfile | 0 .../maxtext_gpu_dependencies.Dockerfile | 2 +- ...ost_training_local_dependencies.Dockerfile | 0 .../dockerfiles/maxtext_runner.Dockerfile | 0 .../maxtext_tpu_dependencies.Dockerfile | 2 +- .../gpu-base-requirements.txt | 0 .../base_requirements/requirements.txt | 0 .../tpu-base-requirements.txt | 0 .../cuda12-requirements.txt | 0 .../tpu-post-train-requirements.txt | 0 .../tpu-requirements.txt | 0 .../requirements/requirements.txt | 0 .../requirements_decoupled_jax_0_7.1.txt | 0 .../requirements/requirements_docs.txt | 0 .../scripts/docker_build_dependency_image.py | 27 +++++++++++++++++++ .../scripts/docker_build_dependency_image.sh | 24 ++++++++--------- .../scripts/docker_upload_runner.py | 27 +++++++++++++++++++ .../scripts/docker_upload_runner.sh | 4 +-- tools/setup/setup.sh | 8 +++--- 34 files changed, 107 insertions(+), 51 deletions(-) rename {dependencies => src/dependencies}/dockerfiles/clean_py_env.Dockerfile (100%) rename {dependencies => src/dependencies}/dockerfiles/clean_py_env_gpu.Dockerfile (100%) rename {dependencies => src/dependencies}/dockerfiles/maxtext_gpu_dependencies.Dockerfile (97%) rename {dependencies => src/dependencies}/dockerfiles/maxtext_post_training_local_dependencies.Dockerfile (100%) rename {dependencies => src/dependencies}/dockerfiles/maxtext_runner.Dockerfile (100%) rename {dependencies => src/dependencies}/dockerfiles/maxtext_tpu_dependencies.Dockerfile (97%) rename {dependencies => src/dependencies}/requirements/base_requirements/gpu-base-requirements.txt (100%) rename {dependencies => src/dependencies}/requirements/base_requirements/requirements.txt (100%) rename {dependencies => src/dependencies}/requirements/base_requirements/tpu-base-requirements.txt (100%) rename {dependencies => src/dependencies}/requirements/generated_requirements/cuda12-requirements.txt (100%) rename {dependencies => src/dependencies}/requirements/generated_requirements/tpu-post-train-requirements.txt (100%) rename {dependencies => src/dependencies}/requirements/generated_requirements/tpu-requirements.txt (100%) rename {dependencies => src/dependencies}/requirements/requirements.txt (100%) rename {dependencies => src/dependencies}/requirements/requirements_decoupled_jax_0_7.1.txt (100%) rename {dependencies => src/dependencies}/requirements/requirements_docs.txt (100%) create mode 100644 src/dependencies/scripts/docker_build_dependency_image.py rename {dependencies => src/dependencies}/scripts/docker_build_dependency_image.sh (83%) create mode 100644 src/dependencies/scripts/docker_upload_runner.py rename {dependencies => src/dependencies}/scripts/docker_upload_runner.sh (97%) diff --git a/.github/workflows/UploadDockerImages.yml b/.github/workflows/UploadDockerImages.yml index 98e8cda22e..389510bcf9 100644 --- a/.github/workflows/UploadDockerImages.yml +++ b/.github/workflows/UploadDockerImages.yml @@ -65,11 +65,11 @@ jobs: - device: tpu build_mode: stable image_name: maxtext_jax_stable - dockerfile: ./dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile + dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile - device: tpu build_mode: nightly image_name: maxtext_jax_nightly - dockerfile: ./dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile + dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile uses: ./.github/workflows/build_and_push_docker_image.yml with: image_name: ${{ matrix.image_name }} @@ -88,7 +88,7 @@ jobs: device: tpu build_mode: stable workflow: post-training - dockerfile: ./dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile + dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }} image_date: ${{ needs.setup.outputs.image_date }} @@ -101,7 +101,7 @@ jobs: device: tpu build_mode: nightly workflow: post-training - dockerfile: ./dependencies/dockerfiles/maxtext_post_training_local_dependencies.Dockerfile + dockerfile: ./src/dependencies/dockerfiles/maxtext_post_training_local_dependencies.Dockerfile maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }} image_date: ${{ needs.setup.outputs.image_date }} base_image: gcr.io/tpu-prod-env-multipod/maxtext_post_training_stable:${{ needs.setup.outputs.image_date }} @@ -116,11 +116,11 @@ jobs: - device: gpu build_mode: stable image_name: maxtext_gpu_jax_stable - dockerfile: ./dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile + dockerfile: ./src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile - device: gpu build_mode: nightly image_name: maxtext_gpu_jax_nightly - dockerfile: ./dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile + dockerfile: ./src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile uses: ./.github/workflows/build_and_push_docker_image.yml with: image_name: ${{ matrix.image_name }} diff --git a/.github/workflows/build_and_upload_images.sh b/.github/workflows/build_and_upload_images.sh index da478f6efc..b03badb772 100644 --- a/.github/workflows/build_and_upload_images.sh +++ b/.github/workflows/build_and_upload_images.sh @@ -49,7 +49,7 @@ if [[ ! -v CLOUD_IMAGE_NAME ]] || [[ ! -v PROJECT ]] || [[ ! -v MODE ]] || [[ ! fi gcloud auth configure-docker us-docker.pkg.dev --quiet -bash "$MAXTEXT_REPO_ROOT"'/dependencies/scripts/docker_build_dependency_image.sh' LOCAL_IMAGE_NAME=$LOCAL_IMAGE_NAME MODE="$MODE" DEVICE="$DEVICE" +bash "$MAXTEXT_REPO_ROOT"'/src/dependencies/scripts/docker_build_dependency_image.sh' LOCAL_IMAGE_NAME=$LOCAL_IMAGE_NAME MODE="$MODE" DEVICE="$DEVICE" image_date=$(date +%Y-%m-%d) # Upload only dependencies image @@ -65,7 +65,7 @@ if ! gcloud storage cp gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_ROOT:-$ fi # Build then upload "dependencies + code" image -docker build --build-arg BASEIMAGE=${LOCAL_IMAGE_NAME} -f "$MAXTEXT_REPO_ROOT"'/dependencies/dockerfiles/maxtext_runner.Dockerfile' -t ${LOCAL_IMAGE_NAME}_runner . +docker build --build-arg BASEIMAGE=${LOCAL_IMAGE_NAME} -f "$MAXTEXT_REPO_ROOT"'/src/dependencies/dockerfiles/maxtext_runner.Dockerfile' -t ${LOCAL_IMAGE_NAME}_runner . docker tag ${LOCAL_IMAGE_NAME}_runner gcr.io/$PROJECT/${CLOUD_IMAGE_NAME}:latest docker push gcr.io/$PROJECT/${CLOUD_IMAGE_NAME}:latest docker tag ${LOCAL_IMAGE_NAME}_runner gcr.io/$PROJECT/${CLOUD_IMAGE_NAME}:${image_date} diff --git a/.github/workflows/check_docs_build.yml b/.github/workflows/check_docs_build.yml index a8d6098350..59d18110b0 100644 --- a/.github/workflows/check_docs_build.yml +++ b/.github/workflows/check_docs_build.yml @@ -27,7 +27,7 @@ jobs: run: uv venv --python 3.12 $GITHUB_WORKSPACE/venv - name: Install dependencies - run: . $GITHUB_WORKSPACE/venv/bin/activate && uv pip install -r dependencies/requirements/requirements_docs.txt + run: . $GITHUB_WORKSPACE/venv/bin/activate && uv pip install -r src/dependencies/requirements/requirements_docs.txt - name: Build documentation run: | diff --git a/.readthedocs.yml b/.readthedocs.yml index f8c5c17004..175e43698c 100644 --- a/.readthedocs.yml +++ b/.readthedocs.yml @@ -21,4 +21,4 @@ sphinx: # See https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html python: install: - - requirements: dependencies/requirements/requirements_docs.txt + - requirements: src/dependencies/requirements/requirements_docs.txt diff --git a/PREFLIGHT.md b/PREFLIGHT.md index d713f10f96..0ec1875eec 100644 --- a/PREFLIGHT.md +++ b/PREFLIGHT.md @@ -26,7 +26,7 @@ bash preflight.sh PLATFORM=GCE && numactl --membind 0 --cpunodebind=0 python3 -m ``` For GKE, -`numactl` should be built into your docker image from [maxtext_tpu_dependencies.Dockerfile](https://github.com/google/maxtext/blob/main/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile), so you can use it directly if you built the maxtext docker image. Here is an example +`numactl` should be built into your docker image from [maxtext_tpu_dependencies.Dockerfile](https://github.com/google/maxtext/blob/main/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile), so you can use it directly if you built the maxtext docker image. Here is an example ``` bash preflight.sh PLATFORM=GKE && numactl --membind 0 --cpunodebind=0 python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml run_name=${YOUR_JOB_NAME?} diff --git a/build_hooks.py b/build_hooks.py index a4ec768e5c..616a27aff7 100644 --- a/build_hooks.py +++ b/build_hooks.py @@ -17,7 +17,7 @@ import os from hatchling.builders.hooks.plugin.interface import BuildHookInterface -TPU_REQUIREMENTS_PATH = "dependencies/requirements/generated_requirements/tpu-requirements.txt" +TPU_REQUIREMENTS_PATH = "src/dependencies/requirements/generated_requirements/tpu-requirements.txt" def get_tpu_dependencies(): diff --git a/docs/development.md b/docs/development.md index 0302bcf9df..5b45dc8525 100644 --- a/docs/development.md +++ b/docs/development.md @@ -12,7 +12,7 @@ If you are writing documentation for MaxText, you may want to preview the docume First, make sure you install the necessary dependencies. You can do this by navigating to your local clone of the MaxText repo and running: ```bash -pip install -r dependencies/requirements/requirements_docs.txt +pip install -r src/dependencies/requirements/requirements_docs.txt ``` Once the dependencies are installed, you can navigate to the `docs/` folder and run: diff --git a/docs/install_maxtext.md b/docs/install_maxtext.md index 3b6a18fcdf..bf9579f53c 100644 --- a/docs/install_maxtext.md +++ b/docs/install_maxtext.md @@ -127,7 +127,7 @@ Run the following command, replacing `` with the hash you ```bash seed-env \ - --local-requirements=dependencies/requirements/base_requirements/tpu-base-requirements.txt \ + --local-requirements=src/dependencies/requirements/base_requirements/tpu-base-requirements.txt \ --host-name=MaxText \ --seed-commit= \ --python-version=3.12 \ @@ -141,7 +141,7 @@ Similarly, run the command for the GPU requirements. ```bash seed-env \ - --local-requirements=dependencies/requirements/base_requirements/cuda12-base-requirements.txt \ + --local-requirements=src/dependencies/requirements/base_requirements/cuda12-base-requirements.txt \ --host-name=MaxText \ --seed-commit= \ --python-version=3.12 \ diff --git a/docs/run_maxtext/run_maxtext_via_pathways.md b/docs/run_maxtext/run_maxtext_via_pathways.md index 6da385ad98..8e78c92d46 100644 --- a/docs/run_maxtext/run_maxtext_via_pathways.md +++ b/docs/run_maxtext/run_maxtext_via_pathways.md @@ -40,7 +40,7 @@ Before you can run a MaxText workload, you must complete the following setup ste Step 1: Build the Docker image for a TPU device. This image contains MaxText and its dependencies. ```shell - bash dependencies/scripts/docker_build_dependency_image.sh DEVICE=tpu MODE=stable + bash src/dependencies/scripts/docker_build_dependency_image.sh DEVICE=tpu MODE=stable ``` Step 2: Configure Docker to authenticate with Google Cloud @@ -52,7 +52,7 @@ Before you can run a MaxText workload, you must complete the following setup ste Step 3: Upload the image to your project's registry. Replace `$USER_runner` with your desired image name. ```shell - bash dependencies/scripts/docker_upload_runner.sh CLOUD_IMAGE_NAME=$USER_runner + bash src/dependencies/scripts/docker_upload_runner.sh CLOUD_IMAGE_NAME=$USER_runner ``` ## 2. Environment configuration diff --git a/docs/run_maxtext/run_maxtext_via_xpk.md b/docs/run_maxtext/run_maxtext_via_xpk.md index 9ad05000a9..1f2f6b04d0 100644 --- a/docs/run_maxtext/run_maxtext_via_xpk.md +++ b/docs/run_maxtext/run_maxtext_via_xpk.md @@ -130,13 +130,13 @@ ______________________________________________________________________ - **For TPUs:** ``` - bash dependencies/scripts/docker_build_dependency_image.sh DEVICE=tpu MODE=stable + bash src/dependencies/scripts/docker_build_dependency_image.sh DEVICE=tpu MODE=stable ``` - **For GPUs:** ``` - bash dependencies/scripts/docker_build_dependency_image.sh DEVICE=gpu MODE=stable + bash src/dependencies/scripts/docker_build_dependency_image.sh DEVICE=gpu MODE=stable ``` ______________________________________________________________________ diff --git a/docs/tutorials/first_run.md b/docs/tutorials/first_run.md index 7fee3673e1..776b408c91 100644 --- a/docs/tutorials/first_run.md +++ b/docs/tutorials/first_run.md @@ -79,7 +79,7 @@ You can use [demo_decoding.ipynb](https://github.com/AI-Hypercomputer/maxtext/bl ### Run MaxText on NVIDIA GPUs -1. Use `bash dependencies/scripts/docker_build_dependency_image.sh DEVICE=gpu` to build a container with the required dependencies. +1. Use `bash src/dependencies/scripts/docker_build_dependency_image.sh DEVICE=gpu` to build a container with the required dependencies. 2. After installation is complete, run training with the following command on synthetic data: ```sh diff --git a/docs/tutorials/posttraining/rl_on_multi_host.md b/docs/tutorials/posttraining/rl_on_multi_host.md index 4cb2f1de1d..ff9819a923 100644 --- a/docs/tutorials/posttraining/rl_on_multi_host.md +++ b/docs/tutorials/posttraining/rl_on_multi_host.md @@ -142,13 +142,13 @@ Run the following script to create a Docker image with stable releases of MaxText, and its post-training dependencies. The build process takes approximately 10-15 minutes. ```bash -bash dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-training +bash src/dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-training ``` For experimental features (such as improved pathwaysutils resharding API), use: ```bash -bash dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-training-experimental +bash src/dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-training-experimental ``` ### Option 2: From Github @@ -159,7 +159,7 @@ For using a version newer than the latest PyPI release, you could also build the git clone https://github.com/AI-Hypercomputer/maxtext.git cd maxtext -bash dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-training +bash src/dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-training ``` ### Upload the Docker Image @@ -170,7 +170,7 @@ bash dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-trainin > project administrator if you don't have this permission. ```bash -bash dependencies/scripts/docker_upload_runner.sh CLOUD_IMAGE_NAME=${CLOUD_IMAGE_NAME?} +bash src/dependencies/scripts/docker_upload_runner.sh CLOUD_IMAGE_NAME=${CLOUD_IMAGE_NAME?} ``` ## Submit your RL workload via Pathways diff --git a/docs/tutorials/posttraining/sft_on_multi_host.md b/docs/tutorials/posttraining/sft_on_multi_host.md index 85e7964c86..d2c42dd54d 100644 --- a/docs/tutorials/posttraining/sft_on_multi_host.md +++ b/docs/tutorials/posttraining/sft_on_multi_host.md @@ -52,7 +52,7 @@ docker run hello-world Then run the following command to create a local Docker image named `maxtext_base_image`. This build process takes approximately 10 to 15 minutes. ```bash -bash dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-training +bash src/dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-training ``` ### 1.3. Upload the Docker image to Artifact Registry @@ -61,7 +61,7 @@ bash dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-trainin ```bash export DOCKER_IMAGE_NAME= -bash dependencies/scripts/docker_upload_runner.sh CLOUD_IMAGE_NAME=${DOCKER_IMAGE_NAME?} +bash src/dependencies/scripts/docker_upload_runner.sh CLOUD_IMAGE_NAME=${DOCKER_IMAGE_NAME?} ``` The `docker_upload_runner.sh` script uploads your Docker image to Artifact Registry. diff --git a/pyproject.toml b/pyproject.toml index 09c93c7f41..68857b9502 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,10 +26,10 @@ classifiers = [ dependencies = [] [tool.hatch.metadata.hooks.requirements_txt.optional-dependencies] -tpu = ["dependencies/requirements/generated_requirements/tpu-requirements.txt"] -tpu-post-train = ["dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt"] -cuda12 = ["dependencies/requirements/generated_requirements/cuda12-requirements.txt"] -docs = ["dependencies/requirements/requirements_docs.txt"] +tpu = ["src/dependencies/requirements/generated_requirements/tpu-requirements.txt"] +tpu-post-train = ["src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt"] +cuda12 = ["src/dependencies/requirements/generated_requirements/cuda12-requirements.txt"] +docs = ["src/dependencies/requirements/requirements_docs.txt"] [project.urls] Repository = "https://github.com/AI-Hypercomputer/maxtext.git" @@ -39,7 +39,7 @@ Repository = "https://github.com/AI-Hypercomputer/maxtext.git" allow-direct-references = true [tool.hatch.build.targets.wheel] -packages = ["src/MaxText", "src/maxtext", "src/install_maxtext_extra_deps"] +packages = ["src/MaxText", "src/maxtext", "src/install_maxtext_extra_deps", "src/dependencies"] # TODO: Add this hook back when it handles device-type parsing # [tool.hatch.build.targets.wheel.hooks.custom] @@ -49,3 +49,5 @@ packages = ["src/MaxText", "src/maxtext", "src/install_maxtext_extra_deps"] install_maxtext_tpu_github_deps = "install_maxtext_extra_deps.install_github_deps:main" install_maxtext_cuda12_github_deps = "install_maxtext_extra_deps.install_github_deps:main" install_maxtext_tpu_post_train_extra_deps = "install_maxtext_extra_deps.install_post_train_extra_deps:main" +docker_build_dependency_image = "dependencies.scripts.docker_build_dependency_image:main" +docker_upload_runner = "dependencies.scripts.docker_upload_runner:main" diff --git a/dependencies/dockerfiles/clean_py_env.Dockerfile b/src/dependencies/dockerfiles/clean_py_env.Dockerfile similarity index 100% rename from dependencies/dockerfiles/clean_py_env.Dockerfile rename to src/dependencies/dockerfiles/clean_py_env.Dockerfile diff --git a/dependencies/dockerfiles/clean_py_env_gpu.Dockerfile b/src/dependencies/dockerfiles/clean_py_env_gpu.Dockerfile similarity index 100% rename from dependencies/dockerfiles/clean_py_env_gpu.Dockerfile rename to src/dependencies/dockerfiles/clean_py_env_gpu.Dockerfile diff --git a/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile b/src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile similarity index 97% rename from dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile rename to src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile index 968ed78407..c1a0ed2b71 100644 --- a/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile +++ b/src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile @@ -48,7 +48,7 @@ WORKDIR /deps # Copy setup files and dependency files separately for better caching COPY tools/setup tools/setup/ -COPY dependencies/requirements/ dependencies/requirements/ +COPY src/dependencies/requirements/ src/dependencies/requirements/ COPY src/install_maxtext_extra_deps/extra_deps_from_github.txt src/install_maxtext_extra_deps/ # Install dependencies - these steps are cached unless the copied files change diff --git a/dependencies/dockerfiles/maxtext_post_training_local_dependencies.Dockerfile b/src/dependencies/dockerfiles/maxtext_post_training_local_dependencies.Dockerfile similarity index 100% rename from dependencies/dockerfiles/maxtext_post_training_local_dependencies.Dockerfile rename to src/dependencies/dockerfiles/maxtext_post_training_local_dependencies.Dockerfile diff --git a/dependencies/dockerfiles/maxtext_runner.Dockerfile b/src/dependencies/dockerfiles/maxtext_runner.Dockerfile similarity index 100% rename from dependencies/dockerfiles/maxtext_runner.Dockerfile rename to src/dependencies/dockerfiles/maxtext_runner.Dockerfile diff --git a/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile b/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile similarity index 97% rename from dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile rename to src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile index 61ca472818..6de3313d75 100644 --- a/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile +++ b/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile @@ -45,7 +45,7 @@ WORKDIR /deps # Copy setup files and dependency files separately for better caching COPY tools/setup tools/setup/ -COPY dependencies/requirements/ dependencies/requirements/ +COPY src/dependencies/requirements/ src/dependencies/requirements/ COPY src/install_maxtext_extra_deps/ src/install_maxtext_extra_deps/ COPY src/maxtext/integration/vllm/ src/maxtext/integration/vllm/ diff --git a/dependencies/requirements/base_requirements/gpu-base-requirements.txt b/src/dependencies/requirements/base_requirements/gpu-base-requirements.txt similarity index 100% rename from dependencies/requirements/base_requirements/gpu-base-requirements.txt rename to src/dependencies/requirements/base_requirements/gpu-base-requirements.txt diff --git a/dependencies/requirements/base_requirements/requirements.txt b/src/dependencies/requirements/base_requirements/requirements.txt similarity index 100% rename from dependencies/requirements/base_requirements/requirements.txt rename to src/dependencies/requirements/base_requirements/requirements.txt diff --git a/dependencies/requirements/base_requirements/tpu-base-requirements.txt b/src/dependencies/requirements/base_requirements/tpu-base-requirements.txt similarity index 100% rename from dependencies/requirements/base_requirements/tpu-base-requirements.txt rename to src/dependencies/requirements/base_requirements/tpu-base-requirements.txt diff --git a/dependencies/requirements/generated_requirements/cuda12-requirements.txt b/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt similarity index 100% rename from dependencies/requirements/generated_requirements/cuda12-requirements.txt rename to src/dependencies/requirements/generated_requirements/cuda12-requirements.txt diff --git a/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt b/src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt similarity index 100% rename from dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt rename to src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt diff --git a/dependencies/requirements/generated_requirements/tpu-requirements.txt b/src/dependencies/requirements/generated_requirements/tpu-requirements.txt similarity index 100% rename from dependencies/requirements/generated_requirements/tpu-requirements.txt rename to src/dependencies/requirements/generated_requirements/tpu-requirements.txt diff --git a/dependencies/requirements/requirements.txt b/src/dependencies/requirements/requirements.txt similarity index 100% rename from dependencies/requirements/requirements.txt rename to src/dependencies/requirements/requirements.txt diff --git a/dependencies/requirements/requirements_decoupled_jax_0_7.1.txt b/src/dependencies/requirements/requirements_decoupled_jax_0_7.1.txt similarity index 100% rename from dependencies/requirements/requirements_decoupled_jax_0_7.1.txt rename to src/dependencies/requirements/requirements_decoupled_jax_0_7.1.txt diff --git a/dependencies/requirements/requirements_docs.txt b/src/dependencies/requirements/requirements_docs.txt similarity index 100% rename from dependencies/requirements/requirements_docs.txt rename to src/dependencies/requirements/requirements_docs.txt diff --git a/src/dependencies/scripts/docker_build_dependency_image.py b/src/dependencies/scripts/docker_build_dependency_image.py new file mode 100644 index 0000000000..0f1487dd3a --- /dev/null +++ b/src/dependencies/scripts/docker_build_dependency_image.py @@ -0,0 +1,27 @@ +# Copyright 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. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Wrapper to run docker_build_dependency_image.sh from pip install.""" + +import os +import sys + + +def main(): + script_path = os.path.join(os.path.dirname(__file__), "docker_build_dependency_image.sh") + if not os.path.exists(script_path): + raise FileNotFoundError(f"Script not found at {script_path}") + + cmd = ["bash", script_path] + sys.argv[1:] + os.execvp("bash", cmd) diff --git a/dependencies/scripts/docker_build_dependency_image.sh b/src/dependencies/scripts/docker_build_dependency_image.sh similarity index 83% rename from dependencies/scripts/docker_build_dependency_image.sh rename to src/dependencies/scripts/docker_build_dependency_image.sh index 697c8f3b05..1321e77d9f 100644 --- a/dependencies/scripts/docker_build_dependency_image.sh +++ b/src/dependencies/scripts/docker_build_dependency_image.sh @@ -23,34 +23,34 @@ # ================================== # Build docker image with stable dependencies -## bash dependencies/scripts/docker_build_dependency_image.sh DEVICE={{gpu|tpu}} MODE=stable +## bash src/dependencies/scripts/docker_build_dependency_image.sh DEVICE={{gpu|tpu}} MODE=stable # Build docker image with nightly dependencies -## bash dependencies/scripts/docker_build_dependency_image.sh DEVICE={{gpu|tpu}} MODE=nightly +## bash src/dependencies/scripts/docker_build_dependency_image.sh DEVICE={{gpu|tpu}} MODE=nightly # Build docker image with stable dependencies and, a pinned JAX_VERSION for TPUs -## bash dependencies/scripts/docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.4.13 +## bash src/dependencies/scripts/docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.4.13 # Build docker image with a pinned JAX_VERSION and, a pinned LIBTPU_VERSION for TPUs -## bash dependencies/scripts/docker_build_dependency_image.sh MODE={{stable|nightly}} JAX_VERSION=0.8.1 LIBTPU_VERSION=0.0.31.dev20251119+nightly +## bash src/dependencies/scripts/docker_build_dependency_image.sh MODE={{stable|nightly}} JAX_VERSION=0.8.1 LIBTPU_VERSION=0.0.31.dev20251119+nightly # Build docker image with a custom libtpu.so for TPUs # Note: libtpu.so file must be present in the root directory of the MaxText repository -## bash dependencies/scripts/docker_build_dependency_image.sh MODE={{stable|nightly}} +## bash src/dependencies/scripts/docker_build_dependency_image.sh MODE={{stable|nightly}} # Build docker image with nightly dependencies and, a pinned JAX_VERSION for GPUs # Available versions listed at https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/simple/jax -## bash dependencies/scripts/docker_build_dependency_image.sh DEVICE=gpu MODE=nightly JAX_VERSION=0.4.36.dev20241109 +## bash src/dependencies/scripts/docker_build_dependency_image.sh DEVICE=gpu MODE=nightly JAX_VERSION=0.4.36.dev20241109 # ================================== # POST-TRAINING BUILD EXAMPLES # ================================== # Build docker image with stable pre-training dependencies and stable post-training dependencies -## bash dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-training +## bash src/dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-training # Build docker image with stable pre-training dependencies and post-training dependencies from GitHub head -## bash dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-training POST_TRAINING_SOURCE=local +## bash src/dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-training POST_TRAINING_SOURCE=local if [ "${BASH_SOURCE-}" ]; then this_file="${BASH_SOURCE[0]}" @@ -61,7 +61,7 @@ else this_file="${0}" fi -MAXTEXT_REPO_ROOT="${MAXTEXT_REPO_ROOT:-$(CDPATH='' cd -- "$(dirname -- "${this_file}")"'/../..' && pwd)}" +MAXTEXT_REPO_ROOT="${MAXTEXT_REPO_ROOT:-$(CDPATH='' cd -- "$(dirname -- "${this_file}")"'/../../..' && pwd)}" # Enable "exit immediately if any command fails" option set -e @@ -132,7 +132,7 @@ build_post_training_deps_from_local_github() { DOCKERFILE_NAME='maxtext_post_training_local_dependencies.Dockerfile' echo "Building local post-training dependencies: $DOCKERFILE_NAME" - run_docker_build "$MAXTEXT_REPO_ROOT/dependencies/dockerfiles/$DOCKERFILE_NAME" \ + run_docker_build "$MAXTEXT_REPO_ROOT/src/dependencies/dockerfiles/$DOCKERFILE_NAME" \ "MODE=${WORKFLOW}" "BASEIMAGE=${LOCAL_IMAGE_NAME}" } @@ -144,7 +144,7 @@ build_gpu_image() { fi echo "Building docker image with arguments: ${docker_build_args[*]}" - run_docker_build "$MAXTEXT_REPO_ROOT/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile" "${docker_build_args[@]}" + run_docker_build "$MAXTEXT_REPO_ROOT/src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile" "${docker_build_args[@]}" } # Function to build image for TPUs @@ -161,7 +161,7 @@ build_tpu_image() { fi echo "Building docker image with arguments: ${docker_build_args[*]}" - run_docker_build "$MAXTEXT_REPO_ROOT/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile" "${docker_build_args[@]}" + run_docker_build "$MAXTEXT_REPO_ROOT/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile" "${docker_build_args[@]}" # Handle post-training workflow if specified if [[ ${WORKFLOW} == "post-training" || ${WORKFLOW} == "post-training-experimental" ]]; then diff --git a/src/dependencies/scripts/docker_upload_runner.py b/src/dependencies/scripts/docker_upload_runner.py new file mode 100644 index 0000000000..70d31894a4 --- /dev/null +++ b/src/dependencies/scripts/docker_upload_runner.py @@ -0,0 +1,27 @@ +# Copyright 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. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Wrapper to run docker_upload_runner.sh from pip install.""" + +import os +import sys + + +def main(): + script_path = os.path.join(os.path.dirname(__file__), "docker_upload_runner.sh") + if not os.path.exists(script_path): + raise FileNotFoundError(f"Script not found at {script_path}") + + cmd = ["bash", script_path] + sys.argv[1:] + os.execvp("bash", cmd) diff --git a/dependencies/scripts/docker_upload_runner.sh b/src/dependencies/scripts/docker_upload_runner.sh similarity index 97% rename from dependencies/scripts/docker_upload_runner.sh rename to src/dependencies/scripts/docker_upload_runner.sh index 9e07fac93a..126e25acc5 100644 --- a/dependencies/scripts/docker_upload_runner.sh +++ b/src/dependencies/scripts/docker_upload_runner.sh @@ -32,7 +32,7 @@ else this_file="${0}" fi -MAXTEXT_REPO_ROOT="${MAXTEXT_REPO_ROOT:-$(CDPATH='' cd -- "$(dirname -- "${this_file}")"'/../..' && pwd)}" +MAXTEXT_REPO_ROOT="${MAXTEXT_REPO_ROOT:-$(CDPATH='' cd -- "$(dirname -- "${this_file}")"'/../../..' && pwd)}" set -e @@ -111,7 +111,7 @@ if ! docker image inspect "${LOCAL_IMAGE_NAME}" &> /dev/null; then fi docker build --no-cache --build-arg BASEIMAGE=${LOCAL_IMAGE_NAME} \ - -f "$MAXTEXT_REPO_ROOT"'/dependencies/dockerfiles/maxtext_runner.Dockerfile' \ + -f "$MAXTEXT_REPO_ROOT"'/src/dependencies/dockerfiles/maxtext_runner.Dockerfile' \ -t ${LOCAL_IMAGE_NAME_RUNNER} . docker tag ${LOCAL_IMAGE_NAME_RUNNER} gcr.io/$PROJECT/${CLOUD_IMAGE_NAME}:latest diff --git a/tools/setup/setup.sh b/tools/setup/setup.sh index 9d5acb69a0..a4e24937c0 100644 --- a/tools/setup/setup.sh +++ b/tools/setup/setup.sh @@ -67,7 +67,7 @@ export NEEDRESTART_MODE=l # Directory Validation Check echo "Checking current directory..." -if [[ ! -d "dependencies" || ! -d "src" ]]; then +if [[ ! -d "src" || ! -d "src/dependencies" ]]; then echo -e "\n\e[31mERROR: Critical directories not found!\e[0m" echo "Please run this script from the root of the MaxText repository." echo "Expected to find './dependencies' and './src' folders." @@ -210,9 +210,9 @@ install_maxtext_with_deps() { fi echo "Setting up MaxText in $MODE mode for $DEVICE device" if [ "$DEVICE" = "gpu" ]; then - dep_name='dependencies/requirements/generated_requirements/cuda12-requirements.txt' + dep_name='src/dependencies/requirements/generated_requirements/cuda12-requirements.txt' else - dep_name='dependencies/requirements/generated_requirements/tpu-requirements.txt' + dep_name='src/dependencies/requirements/generated_requirements/tpu-requirements.txt' fi echo "Installing requirements from $dep_name" python3 -m uv pip install --resolution=lowest -r "$dep_name" \ @@ -227,7 +227,7 @@ install_post_training_deps() { exit 1 fi echo "Setting up MaxText post-training workflow for $DEVICE device" - dep_name='dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt' + dep_name='src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt' echo "Installing requirements from $dep_name" python3 -m uv pip install --resolution=lowest -r "$dep_name" python3 -m src.install_maxtext_extra_deps.install_post_train_extra_deps