Make Docker build and upload scripts runnable through pip

Author: bvandermoon

.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 }}

.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}

.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: |

.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

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?}

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():

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:

docs/install_maxtext.md

@@ -127,7 +127,7 @@ Run the following command, replacing `<jax-build-commit-hash>` 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=<jax-build-commit-hash> \
   --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=<jax-build-commit-hash> \
   --python-version=3.12 \

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

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
      ```
 
 ______________________________________________________________________