PR #3449: Move install_maxtext_extra_deps to dependencies directory

bvandermoon · Google-ML-Automation · commit 4a4956379c58 · 2026-03-19T17:34:44.000-07:00
Imported from GitHub PR #3449 # Description Move the `install_maxtext_extra_deps` deps directory to `dependencies` after `dependencies` was added to the PyPI package. # Tests This command still works/runs the expected installation: ``` uv pip install -e .[tpu] --resolution=lowest install_maxtext_tpu_github_deps # For testing locally: install_maxtext_tpu_post_train_extra_deps ``` CI also builds this command # Checklist Before submitting this PR, please make sure (put X in square brackets): - [x] I have performed a self-review of my code. For an optional AI review, add the `gemini-review` label. - [x] I have necessary comments in my code, particularly in hard-to-understand areas. - [x] I have run end-to-end tests tests and provided workload links above if applicable. - [x] I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation. Copybara import of the project: -- 277b08e by Branden Vandermoon <bvandermoon@google.com>: Move install_maxtext_extra_deps to dependencies directory Merging this change closes #3449 COPYBARA_INTEGRATE_REVIEW=#3449 from AI-Hypercomputer:bvandermoon-github-deps 277b08e PiperOrigin-RevId: 886463160
diff --git a/.github/workflows/build_and_test_maxtext.yml b/.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:
diff --git a/.github/workflows/run_pathways_tests.yml b/.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
diff --git a/.github/workflows/run_tests_against_package.yml b/.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
diff --git a/docs/conf.py b/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"),
 ]
 
 
diff --git a/docs/install_maxtext.md b/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
diff --git a/pyproject.toml b/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"
diff --git a/src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile b/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
diff --git a/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile b/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
diff --git a/src/dependencies/github_deps/__init__.py b/src/dependencies/github_deps/__init__.py
diff --git a/src/dependencies/github_deps/install_post_train_deps.py b/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:
diff --git a/src/dependencies/github_deps/install_pre_train_deps.py b/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:
diff --git a/src/dependencies/github_deps/post_train_base_deps.txt b/src/dependencies/github_deps/post_train_base_deps.txt
diff --git a/src/dependencies/github_deps/post_train_deps.txt b/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
diff --git a/src/dependencies/github_deps/pre_train_deps.txt b/src/dependencies/github_deps/pre_train_deps.txt
diff --git a/src/dependencies/scripts/setup.sh b/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 ----------
diff --git a/src/maxtext/examples/demo_decoding.ipynb b/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",

Original file line number	Diff line number	Diff line change
`@@ -120,8 +120,8 @@`
`120`	`120`	`os.path.join("run_maxtext", "run_maxtext_via_multihost_runner.md"),`
`121`	`121`	`os.path.join("reference", "core_concepts", "llm_calculator.ipynb"),`
`122`	`122`	`os.path.join("reference", "api_generated", "modules.rst"),`
`123`		`- os.path.join("reference", "api_generated", "install_maxtext_extra_deps.rst"),`
`124`		`- os.path.join("reference", "api_generated", "install_maxtext_extra_deps.install_github_deps.rst"),`
	`123`	`+ os.path.join("reference", "api_generated", "dependencies.github_deps.rst"),`
	`124`	`+ os.path.join("reference", "api_generated", "dependencies.github_deps.install_pre_train_deps.rst"),`
`125`	`125`	`]`
`126`	`126`
`127`	`127`