Merge pull request #2541 from SamuelMarks:setup-deps-refactor

PiperOrigin-RevId: 827675322

Author: Google-ML-Automation

.github/workflows/UploadDockerImages.yml

@@ -52,26 +52,26 @@ jobs:
         include:
           # TPU Image Builds
           - image_name: maxtext_jax_stable
-            dockerfile: ./maxtext_dependencies.Dockerfile
+            dockerfile: ./dependencies/dockerfiles/maxtext_dependencies.Dockerfile
             build_args: |
               MODE=stable
               JAX_VERSION=NONE
               LIBTPU_GCS_PATH=NONE
           - image_name: maxtext_jax_nightly
-            dockerfile: ./maxtext_dependencies.Dockerfile
+            dockerfile: ./dependencies/dockerfiles/maxtext_dependencies.Dockerfile
             build_args: |
               MODE=nightly
               JAX_VERSION=NONE
               LIBTPU_GCS_PATH=NONE
           # TPU Image builds using JAX AI Image
           - image_name: maxtext_jax_stable_stack
-            dockerfile: ./maxtext_jax_ai_image.Dockerfile
+            dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile
             base_image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
           - image_name: maxtext_stable_stack_nightly_jax
-            dockerfile: ./maxtext_jax_ai_image.Dockerfile
+            dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile
             base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/tpu/jax_nightly:latest
           - image_name: maxtext_stable_stack_candidate
-            dockerfile: ./maxtext_jax_ai_image.Dockerfile
+            dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile
             base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/tpu:latest
 
     # Setup for GKE runners per b/412986220#comment82 and b/412986220#comment90
@@ -130,13 +130,13 @@ jobs:
         # GPU Image Builds using JAX AI Image
         include:
           - image_name: maxtext_gpu_jax_stable_stack
-            dockerfile: ./maxtext_jax_ai_image.Dockerfile
+            dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile
             base_image: us-central1-docker.pkg.dev/deeplearning-images/jax-ai-image/gpu:latest
           - image_name: maxtext_gpu_stable_stack_nightly_jax
-            dockerfile: ./maxtext_jax_ai_image.Dockerfile
+            dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile
             base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/gpu/jax_nightly:latest
           - image_name: maxtext_stable_stack_candidate_gpu
-            dockerfile: ./maxtext_jax_ai_image.Dockerfile
+            dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile
             base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/gpu:latest
 
     steps:

.github/workflows/build_and_upload_images.sh

@@ -24,7 +24,16 @@
 # Example command:
 # bash build_and_upload_images.sh PROJECT=<project> MODE=stable DEVICE=tpu CLOUD_IMAGE_NAME=${USER}_runner
 
+if [ "${BASH_SOURCE-}" ]; then
+  this_file="${BASH_SOURCE[0]}"
+elif [ "${ZSH_VERSION-}" ]; then
+  # shellcheck disable=SC2296
+  this_file="${(%):-%x}"
+else
+  this_file="${0}"
+fi
 
+MAXTEXT_REPO_ROOT="${MAXTEXT_REPO_ROOT:-$(CDPATH='' cd -- "$(dirname -- "${this_file}")"'/../..' && pwd)}"
 export LOCAL_IMAGE_NAME=maxtext_base_image
 
 # Set environment variables
@@ -40,7 +49,7 @@ if [[ ! -v CLOUD_IMAGE_NAME ]] || [[ ! -v PROJECT ]] || [[ ! -v MODE ]] || [[ !
 fi
 
 gcloud auth configure-docker us-docker.pkg.dev --quiet
-bash docker_build_dependency_image.sh LOCAL_IMAGE_NAME=$LOCAL_IMAGE_NAME MODE=$MODE DEVICE=$DEVICE
+bash "$MAXTEXT_REPO_ROOT"'/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
@@ -56,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_runner.Dockerfile -t ${LOCAL_IMAGE_NAME}_runner .
+docker build --build-arg BASEIMAGE=${LOCAL_IMAGE_NAME} -f "$MAXTEXT_REPO_ROOT"'/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/build_upload_internal.yml

@@ -59,7 +59,7 @@ jobs:
         with:
           push: true
           context: .
-          file: ./maxtext_jax_ai_image.Dockerfile
+          file: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile
           tags: gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:${{ inputs.device_type }}
           provenance: false
           build-args: |

.github/workflows/check_docs_build.yml

@@ -24,7 +24,7 @@ jobs:
           cache: 'pip' # caching pip dependencies
 
       - name: Install dependencies
-        run: pip install -r requirements_docs.txt
+        run: pip install -r dependencies/requirements/requirements_docs.txt
 
       - name: Build documentation
         run: |

.github/workflows/run_tests_against_package.yml

@@ -85,7 +85,7 @@ jobs:
           source .venv/bin/activate
           maxtext_wheel=$(ls maxtext-*-py3-none-any.whl 2>/dev/null)
           uv pip install ${maxtext_wheel}[${MAXTEXT_PACKAGE_EXTRA}] --resolution=lowest
-          install_maxtext_github_deps
+          uv pip install -r src/install_maxtext_extra_deps/extra_deps_from_github.txt
           python3 --version
           python3 -m pip freeze
       - name: Copy test assets files

clean_py_env.Dockerfile …cies/dockerfiles/clean_py_env.Dockerfileclean_py_env.Dockerfile renamed to dependencies/dockerfiles/clean_py_env.Dockerfile clean_py_env.Dockerfile renamed to dependencies/dockerfiles/clean_py_env.Dockerfile

@@ -38,7 +38,7 @@ git clone https://github.com/AI-Hypercomputer/maxtext.git
 
 RUN cd maxtext/ && \
 git checkout ${MAXTEXT_VERSION} && \
-bash setup.sh
+bash ./tools/setup/setup.sh
 
 RUN cd /JetStream && \
 git checkout ${JETSTREAM_VERSION} && \

clean_py_env_gpu.Dockerfile …/dockerfiles/clean_py_env_gpu.Dockerfileclean_py_env_gpu.Dockerfile renamed to dependencies/dockerfiles/clean_py_env_gpu.Dockerfile clean_py_env_gpu.Dockerfile renamed to dependencies/dockerfiles/clean_py_env_gpu.Dockerfile

@@ -35,7 +35,7 @@ git clone https://github.com/AI-Hypercomputer/JetStream.git
 
 RUN cd maxtext/ && \
 git checkout ${MAXTEXT_VERSION} && \
-bash setup.sh
+bash ./tools/setup/setup.sh
 
 RUN cd /JetStream && \
 git checkout ${JETSTREAM_VERSION} && \
@@ -45,4 +45,4 @@ COPY maxengine_server_entrypoint.sh /usr/bin/
 
 RUN chmod +x /usr/bin/maxengine_server_entrypoint.sh
 
-ENTRYPOINT ["/usr/bin/maxengine_server_entrypoint.sh"]
+ENTRYPOINT ["/usr/bin/maxengine_server_entrypoint.sh"]