Move install_maxtext_extra_deps to dependencies directory

Author: bvandermoon

.github/workflows/build_and_test_maxtext.yml

@@ -262,7 +262,7 @@ jobs:
       tf_force_gpu_allow_growth: false
       container_resource_option: "--privileged"
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
-      extra_pip_deps_file: 'src/install_maxtext_extra_deps/extra_post_train_base_deps_from_github.txt'
+      extra_pip_deps_file: 'src/dependencies/install_maxtext_extra_deps/extra_post_train_base_deps_from_github.txt'
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
 
   maxtext_post_training_tpu_unit_tests:
@@ -284,7 +284,7 @@ jobs:
       tf_force_gpu_allow_growth: false
       container_resource_option: "--privileged"
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
-      extra_pip_deps_file: 'src/install_maxtext_extra_deps/extra_post_train_base_deps_from_github.txt'
+      extra_pip_deps_file: 'src/dependencies/install_maxtext_extra_deps/extra_post_train_base_deps_from_github.txt'
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
 
   maxtext_gpu_integration_tests:

.github/workflows/run_pathways_tests.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}[tpu] --resolution=lowest
-          uv pip install -r src/install_maxtext_extra_deps/extra_deps_from_github.txt
+          uv pip install -r src/dependencies/install_maxtext_extra_deps/extra_deps_from_github.txt
           python3 --version
           python3 -m pip freeze
       - name: Copy test assets files

.github/workflows/run_tests_against_package.yml

@@ -96,7 +96,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
-          uv pip install -r src/install_maxtext_extra_deps/extra_deps_from_github.txt
+          uv pip install -r src/dependencies/install_maxtext_extra_deps/extra_deps_from_github.txt
           python3 --version
           python3 -m pip freeze
           uv pip install pytest-cov

docs/conf.py

@@ -120,8 +120,8 @@
     os.path.join("run_maxtext", "run_maxtext_via_multihost_runner.md"),
     os.path.join("reference", "core_concepts", "llm_calculator.ipynb"),
     os.path.join("reference", "api_generated", "modules.rst"),
-    os.path.join("reference", "api_generated", "install_maxtext_extra_deps.rst"),
-    os.path.join("reference", "api_generated", "install_maxtext_extra_deps.install_github_deps.rst"),
+    os.path.join("reference", "api_generated", "dependencies.install_maxtext_extra_deps.rst"),
+    os.path.join("reference", "api_generated", "dependencies.install_maxtext_extra_deps.install_github_deps.rst"),
 ]
 
 

docs/install_maxtext.md

@@ -104,7 +104,7 @@ To update dependencies, you will follow these general steps:
 1. **Modify Base Requirements**: Update the desired dependencies in `base_requirements/requirements.txt` or the hardware-specific files (`base_requirements/tpu-base-requirements.txt`, `base_requirements/gpu-base-requirements.txt`).
 2. **Generate New Files**: Run the `seed-env` CLI tool to generate new, fully-pinned requirements files based on your changes.
 3. **Update Project Files**: Copy the newly generated files into the `generated_requirements/` directory.
-4. **Handle GitHub Dependencies**: Move any dependencies that are installed directly from GitHub from the generated files to `src/install_maxtext_extra_deps/extra_deps_from_github.txt`.
+4. **Handle GitHub Dependencies**: Move any dependencies that are installed directly from GitHub from the generated files to `src/dependencies/install_maxtext_extra_deps/extra_deps_from_github.txt`.
 5. **Verify**: Test the new dependencies to ensure the project installs and runs correctly.
 
 The following sections provide detailed instructions for each step.

pyproject.toml

@@ -40,15 +40,15 @@ 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", "src/dependencies"]
+packages = ["src/MaxText", "src/maxtext", "src/dependencies/install_maxtext_extra_deps", "src/dependencies"]
 
 # TODO: Add this hook back when it handles device-type parsing
 # [tool.hatch.build.targets.wheel.hooks.custom]
 # path = "build_hooks.py"
 
 [project.scripts]
-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"
+install_maxtext_tpu_github_deps = "dependencies.install_maxtext_extra_deps.install_github_deps:main"
+install_maxtext_cuda12_github_deps = "dependencies.install_maxtext_extra_deps.install_github_deps:main"
+install_maxtext_tpu_post_train_extra_deps = "dependencies.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"

src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile

@@ -49,7 +49,7 @@ WORKDIR /deps
 # Copy setup files and dependency files separately for better caching
 COPY tools/setup tools/setup/
 COPY src/dependencies/requirements/ src/dependencies/requirements/
-COPY src/install_maxtext_extra_deps/extra_deps_from_github.txt src/install_maxtext_extra_deps/
+COPY src/dependencies/install_maxtext_extra_deps/extra_deps_from_github.txt src/dependencies/install_maxtext_extra_deps/
 
 # Install dependencies - these steps are cached unless the copied files change
 RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION DEVICE=${ENV_DEVICE}"

src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile

@@ -46,7 +46,7 @@ WORKDIR /deps
 # Copy setup files and dependency files separately for better caching
 COPY tools/setup tools/setup/
 COPY src/dependencies/requirements/ src/dependencies/requirements/
-COPY src/install_maxtext_extra_deps/ src/install_maxtext_extra_deps/
+COPY src/dependencies/install_maxtext_extra_deps/ src/dependencies/install_maxtext_extra_deps/
 COPY src/maxtext/integration/vllm/ src/maxtext/integration/vllm/
 
 # Copy the custom libtpu.so file if it exists inside maxtext repository