Move install_maxtext_extra_deps to dependencies directory

.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/github_deps/post_train_base_deps.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/github_deps/post_train_base_deps.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/github_deps/pre_train_deps.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/github_deps/pre_train_deps.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.github_deps.rst"),
+    os.path.join("reference", "api_generated", "dependencies.github_deps.install_pre_train_deps.rst"),
 ]
 
 

docs/install_maxtext.md

@@ -111,7 +111,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/github_deps/pre_train_deps.txt`.
 5. **Verify**: Test the new dependencies to ensure the project installs and runs correctly.
 
 The following sections provide detailed instructions for each step.
@@ -169,7 +169,7 @@ After generating the new requirements, you need to update the files in the MaxTe
    - Move `generated_tpu_artifacts/tpu-requirements.txt` to `generated_requirements/tpu-requirements.txt`.
    - Move `generated_gpu_artifacts/cuda12-requirements.txt` to `generated_requirements/cuda12-requirements.txt`.
 
-2. **Update `extra_deps_from_github.txt` (if necessary):**
+2. **Update `pre_train_deps.txt` (if necessary):**
    Currently, MaxText uses a few dependencies, such as `mlperf-logging` and `google-jetstream`, that are installed directly from GitHub source. These are defined in `base_requirements/requirements.txt`, and the `seed-env` tool will carry them over to the generated requirements files.
 
 ## Step 5: Verify the New Dependencies

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/github_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.github_deps.install_pre_train_deps:main"
+install_maxtext_cuda12_github_deps = "dependencies.github_deps.install_pre_train_deps:main"
+install_maxtext_tpu_post_train_extra_deps = "dependencies.github_deps.install_post_train_deps:main"
 build_maxtext_docker_image = "dependencies.scripts.build_maxtext_docker_image:main"
 upload_maxtext_docker_image = "dependencies.scripts.upload_maxtext_docker_image:main"

src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile

@@ -53,9 +53,9 @@ ENV MAXTEXT_REPO_ROOT=/deps
 WORKDIR /deps
 
 # Copy setup files and dependency files separately for better caching
+COPY ${PACKAGE_DIR}/dependencies/github_deps/ src/dependencies/github_deps/
 COPY ${PACKAGE_DIR}/dependencies/requirements/ src/dependencies/requirements/
 COPY ${PACKAGE_DIR}/dependencies/scripts/ src/dependencies/scripts/
-COPY ${PACKAGE_DIR}/install_maxtext_extra_deps/ src/install_maxtext_extra_deps/
 COPY ${PACKAGE_DIR}/maxtext/integration/vllm/ src/MaxText/integration/vllm/
 
 # Install dependencies - these steps are cached unless the copied files change

src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile

@@ -50,9 +50,9 @@ ENV MAXTEXT_REPO_ROOT=/deps
 WORKDIR /deps
 
 # Copy setup files and dependency files separately for better caching
+COPY ${PACKAGE_DIR}/dependencies/github_deps/ src/dependencies/github_deps/
 COPY ${PACKAGE_DIR}/dependencies/requirements/ src/dependencies/requirements/
 COPY ${PACKAGE_DIR}/dependencies/scripts/ src/dependencies/scripts/
-COPY ${PACKAGE_DIR}/install_maxtext_extra_deps/ src/install_maxtext_extra_deps/
 COPY ${PACKAGE_DIR}/maxtext/integration/vllm/ src/maxtext/integration/vllm/
 
 # Copy the custom libtpu.so file if it exists

src/install_maxtext_extra_deps/install_post_train_extra_deps.py → src/dependencies/github_deps/install_post_train_deps.py

@@ -15,7 +15,7 @@
 """Installs extra dependencies from a requirements file using uv.
 
 This script is designed to be run to install dependencies specified in
-'extra_post_train_deps_from_github.txt', which is expected to be in the same directory.
+'post_train_deps.txt', which is expected to be in the same directory.
 It first ensures 'uv' is installed and then uses it to install the packages
 listed in the requirements file.
 """
@@ -28,22 +28,22 @@
 
 def main():
   """
-  Installs extra dependencies specified in extra_post_train_deps_from_github.txt using uv.
+  Installs extra dependencies specified in post_train_deps.txt using uv.
 
-  This script looks for 'extra_post_train_deps_from_github.txt' relative to its own location.
+  This script looks for 'post_train_deps.txt' relative to its own location.
   It executes 'uv pip install -r <path_to_extra_deps.txt> --resolution=lowest'.
   """
   script_dir = Path(__file__).resolve().parent
 
   os.environ["VLLM_TARGET_DEVICE"] = "tpu"
 
-  # Adjust this path if your extra_post_train_deps_from_github.txt is in a different location,
-  # e.g., script_dir / "data" / "extra_post_train_deps_from_github.txt"
-  extra_deps_file = script_dir / "extra_post_train_deps_from_github.txt"
+  # Adjust this path if your post_train_deps.txt is in a different location,
+  # e.g., script_dir / "data" / "post_train_deps.txt"
+  extra_deps_file = script_dir / "post_train_deps.txt"
 
   if not extra_deps_file.exists():
     print(f"Error: '{extra_deps_file}' not found.")
-    print("Please ensure 'extra_post_train_deps_from_github.txt' is in the correct location relative to the script.")
+    print("Please ensure 'post_train_deps.txt' is in the correct location relative to the script.")
     sys.exit(1)
   # Check if 'uv' is available in the environment
   try:

src/install_maxtext_extra_deps/install_github_deps.py → src/dependencies/github_deps/install_pre_train_deps.py

@@ -15,7 +15,7 @@
 """Installs extra dependencies from a requirements file using uv.
 
 This script is designed to be run to install dependencies specified in
-'extra_deps_from_github.txt', which is expected to be in the same directory.
+'pre_train_deps.txt', which is expected to be in the same directory.
 It first ensures 'uv' is installed and then uses it to install the packages
 listed in the requirements file.
 """
@@ -27,20 +27,20 @@
 
 def main():
   """
-  Installs extra dependencies specified in extra_deps_from_github.txt using uv.
+  Installs extra dependencies specified in pre_train_deps.txt using uv.
 
-  This script looks for 'extra_deps_from_github.txt' relative to its own location.
+  This script looks for 'pre_train_deps.txt' relative to its own location.
   It executes 'uv pip install -r <path_to_extra_deps.txt> --resolution=lowest'.
   """
   script_dir = Path(__file__).resolve().parent
 
-  # Adjust this path if your extra_deps_from_github.txt is in a different location,
-  # e.g., script_dir / "data" / "extra_deps_from_github.txt"
-  extra_deps_file = script_dir / "extra_deps_from_github.txt"
+  # Adjust this path if your pre_train_deps.txt is in a different location,
+  # e.g., script_dir / "data" / "pre_train_deps.txt"
+  extra_deps_file = script_dir / "pre_train_deps.txt"
 
   if not extra_deps_file.exists():
     print(f"Error: '{extra_deps_file}' not found.")
-    print("Please ensure 'extra_deps_from_github.txt' is in the correct location relative to the script.")
+    print("Please ensure 'pre_train_deps.txt' is in the correct location relative to the script.")
     sys.exit(1)
   # Check if 'uv' is available in the environment
   try:

src/install_maxtext_extra_deps/extra_post_train_deps_from_github.txt → src/dependencies/github_deps/post_train_deps.txt

@@ -1,4 +1,4 @@
--r extra_post_train_base_deps_from_github.txt
+-r post_train_base_deps.txt
 google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip
 mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip
 tpu-inference @ https://github.com/vllm-project/tpu-inference/archive/0cae84fc9a883ba1bde02d4f07930e6af9e92958.zip

src/dependencies/scripts/setup.sh

@@ -216,7 +216,7 @@ install_maxtext_with_deps() {
     fi
     echo "Installing requirements from $dep_name"
     python3 -m uv pip install --resolution=lowest -r "$dep_name" \
-        -r 'src/install_maxtext_extra_deps/extra_deps_from_github.txt'
+        -r 'src/dependencies/github_deps/pre_train_deps.txt'
 
     install_maxtext_package_without_deps
 }
@@ -230,7 +230,7 @@ install_post_training_deps() {
     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
+    python3 -m src.dependencies.github_deps.install_post_train_deps
 }
 
 # ---------- Post-Training workflow installation ----------

src/maxtext/examples/demo_decoding.ipynb

@@ -86,7 +86,7 @@
         "\n",
         "# Install MaxText and dependencies\n",
         "!uv pip install maxtext --resolution=lowest\n",
-        "!python3 -m MaxText.install_maxtext_extra_deps\n",
+        "!python3 -m dependencies.github_deps\n",
         "\n",
         "# Use nest_asyncio to allow nested event loops in notebooks\n",
         "!uv pip install nest_asyncio\n",