Refactor install_extra_deps scripts

.github/workflows/run_jupyter_notebooks.yml

@@ -72,7 +72,7 @@ jobs:
 
           # 2. Install MaxText package and all the post training dependencies
           uv pip install ${maxtext_wheel}[tpu-post-train] --resolution=lowest
-          install_maxtext_tpu_post_train_extra_deps
+          install_tpu_post_train_extra_deps
           
           python3 -m pip freeze
       - name: Run Post-Training Notebooks

.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/dependencies/github_deps/pre_train_deps.txt
+          install_tpu_pre_train_extra_deps
           python3 --version
           python3 -m pip freeze
       - name: Copy test assets files

.github/workflows/run_tests_against_package.yml

@@ -110,9 +110,9 @@ jobs:
           echo "Installing ${maxtext_wheel} for ${MAXTEXT_PACKAGE_EXTRA}..."
           uv pip install ${maxtext_wheel}[${MAXTEXT_PACKAGE_EXTRA}] --resolution=lowest
           if [ "${MAXTEXT_PACKAGE_EXTRA}" == "tpu-post-train" ]; then
-            uv pip install -r src/dependencies/github_deps/post_train_base_deps.txt
+            install_tpu_post_train_extra_deps
           else
-            uv pip install -r src/dependencies/github_deps/pre_train_deps.txt
+            install_tpu_pre_train_extra_deps
           fi
           python3 --version
           python3 -m pip freeze

docs/install_maxtext.md

@@ -41,22 +41,22 @@ source maxtext_venv/bin/activate
 
 # Option 1: Installing maxtext[tpu]
 uv pip install maxtext[tpu]==0.2.1 --resolution=lowest
-install_maxtext_tpu_github_deps
+install_tpu_pre_train_extra_deps
 
 # Option 2: Installing maxtext[cuda12]
 uv pip install maxtext[cuda12]==0.2.1 --resolution=lowest
-install_maxtext_cuda12_github_dep
+install_cuda12_pre_train_extra_dep
 
 # Option 3: Installing maxtext[tpu-post-train]
 uv pip install maxtext[tpu-post-train]==0.2.1 --resolution=lowest
-install_maxtext_tpu_post_train_extra_deps
+install_tpu_post_train_extra_deps
 
 # Option 4: Installing maxtext[runner]
 uv pip install maxtext[runner]==0.2.1 --resolution=lowest
 ```
 
-> **Note:** The `install_maxtext_tpu_github_deps`, `install_maxtext_cuda12_github_dep`, and
-> `install_maxtext_tpu_post_train_extra_deps` commands are temporarily required to install dependencies directly from GitHub
+> **Note:** The `install_tpu_pre_train_extra_deps`, `install_cuda12_pre_train_extra_deps`, and
+> `install_tpu_post_train_extra_deps` commands are temporarily required to install dependencies directly from GitHub
 > that are not yet available on PyPI. As shown above, choose the one that corresponds to your use case.
 
 > **Note:** The maxtext package contains a comprehensive list of all direct and transitive dependencies, with lower bounds, generated by [seed-env](https://github.com/google-ml-infra/actions/tree/main/python_seed_env). We highly recommend the `--resolution=lowest` flag. It instructs `uv` to install the specific, tested versions of dependencies defined by MaxText, rather than the latest available ones. This ensures a consistent and reproducible environment, which is critical for stable performance and for running benchmarks.
@@ -82,15 +82,15 @@ source maxtext_venv/bin/activate
 
 # Option 1: Installing .[tpu]
 uv pip install -e .[tpu] --resolution=lowest
-install_maxtext_tpu_github_deps
+install_tpu_pre_train_extra_deps
 
 # Option 2: Installing .[cuda12]
 uv pip install -e .[cuda12] --resolution=lowest
-install_maxtext_cuda12_github_dep
+install_cuda12_pre_train_extra_deps
 
 # Option 3: Installing .[tpu-post-train]
 uv pip install -e .[tpu-post-train] --resolution=lowest
-install_maxtext_tpu_post_train_extra_deps
+install_tpu_post_train_extra_deps
 
 # Option 4: Installing maxtext[runner]
 uv pip install -e .[runner] --resolution=lowest

docs/tutorials/inference.md

@@ -47,7 +47,7 @@ Required-by:
 If the plugin is not installed, please run the install post training extra dependencies script again with the following command:
 
 ```bash
-install_maxtext_tpu_post_train_extra_deps
+install_tpu_post_train_extra_deps
 ```
 
 # Offline Inference

pyproject.toml

@@ -46,8 +46,8 @@ packages = ["src/maxtext", "src/dependencies"]
 path = "build_hooks.py"
 
 [project.scripts]
-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"
+install_tpu_pre_train_extra_deps = "dependencies.scripts.install_pre_train_extra_deps:main"
+install_cuda12_pre_train_extra_deps = "dependencies.scripts.install_pre_train_extra_deps:main"
+install_tpu_post_train_extra_deps = "dependencies.scripts.install_post_train_extra_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

@@ -50,7 +50,7 @@ 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/extra_deps/ src/dependencies/extra_deps/
 COPY ${PACKAGE_DIR}/dependencies/requirements/ src/dependencies/requirements/
 COPY ${PACKAGE_DIR}/dependencies/scripts/ src/dependencies/scripts/
 COPY ${PACKAGE_DIR}/maxtext/integration/vllm/ src/maxtext/integration/vllm/

src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile

@@ -47,7 +47,7 @@ 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/extra_deps/ src/dependencies/extra_deps/
 COPY ${PACKAGE_DIR}/dependencies/requirements/ src/dependencies/requirements/
 COPY ${PACKAGE_DIR}/dependencies/scripts/ src/dependencies/scripts/
 COPY ${PACKAGE_DIR}/maxtext/integration/vllm/ src/maxtext/integration/vllm/

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

@@ -14,10 +14,8 @@
 
 """Installs extra dependencies from a requirements file using uv.
 
-This script is designed to be run to install dependencies specified in
-'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.
+This script is designed to install dependencies specified in 'dependencies/extra_deps/post_train_*.txt'.
+It first ensures 'uv' is installed and then uses it to install the packages listed in the requirements file.
 """
 
 import os
@@ -27,18 +25,16 @@
 
 def main():
   """
-  Installs extra dependencies specified in post_train_deps.txt using uv.
-
-  This script looks for 'post_train_deps.txt' relative to its own location.
+  Installs extra dependencies specified in 'dependencies/extra_deps/post_train_*.txt' using uv.
   It executes 'uv pip install -r <path_to_extra_deps.txt> --resolution=lowest'.
   """
   os.environ["VLLM_TARGET_DEVICE"] = "tpu"
 
   current_dir = os.path.dirname(os.path.abspath(__file__))
   repo_root = os.path.abspath(os.path.join(current_dir, "..", ".."))
-  extra_deps_path = os.path.join(current_dir, "post_train_deps.txt")
-  if not os.path.exists(extra_deps_path):
-    raise FileNotFoundError(f"Dependencies file not found at {extra_deps_path}")
+  github_deps_path = os.path.join(repo_root, "dependencies", "extra_deps", "post_train_github_deps.txt")
+  if not os.path.exists(github_deps_path):
+    raise FileNotFoundError(f"Github dependencies file not found at {github_deps_path}")
 
   # Check if 'uv' is available in the environment
   try:
@@ -49,14 +45,14 @@ def main():
     print(f"Stderr: {e.stderr.decode()}")
     sys.exit(1)
 
-  command = [
+  github_deps_command = [
       sys.executable,  # Use the current Python executable's pip to ensure the correct environment
       "-m",
       "uv",
       "pip",
       "install",
       "-r",
-      str(extra_deps_path),
+      str(github_deps_path),
       "--no-deps",
   ]
 
@@ -72,9 +68,9 @@ def main():
 
   try:
     # Run the command to install Github dependencies
-    print(f"Installing extra dependencies: {' '.join(command)}")
-    _ = subprocess.run(command, check=True, capture_output=True, text=True)
-    print("Extra dependencies installed successfully!")
+    print(f"Installing Github dependencies: {' '.join(github_deps_command)}")
+    _ = subprocess.run(github_deps_command, check=True, capture_output=True, text=True)
+    print("Github dependencies installed successfully!")
 
     # Run the command to install the MaxText vLLM directory
     print(f"Installing MaxText vLLM dependency: {' '.join(local_vllm_install_command)}")

src/dependencies/github_deps/install_pre_train_deps.py → src/dependencies/scripts/install_pre_train_extra_deps.py

@@ -14,10 +14,8 @@
 
 """Installs extra dependencies from a requirements file using uv.
 
-This script is designed to be run to install dependencies specified in
-'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.
+This script is designed to install dependencies specified in 'dependencies/extra_deps/pre_train_*.txt'.
+It first ensures 'uv' is installed and then uses it to install the packages listed in the requirements file.
 """
 
 import os
@@ -27,15 +25,14 @@
 
 def main():
   """
-  Installs extra dependencies specified in pre_train_deps.txt using uv.
-
-  This script looks for 'pre_train_deps.txt' relative to its own location.
+  Installs extra dependencies specified in 'dependencies/extra_deps/pre_train_*.txt' using uv.
   It executes 'uv pip install -r <path_to_extra_deps.txt> --resolution=lowest'.
   """
   current_dir = os.path.dirname(os.path.abspath(__file__))
-  extra_deps_path = os.path.join(current_dir, "pre_train_deps.txt")
-  if not os.path.exists(extra_deps_path):
-    raise FileNotFoundError(f"Dependencies file not found at {extra_deps_path}")
+  repo_root = os.path.abspath(os.path.join(current_dir, "..", ".."))
+  github_deps_path = os.path.join(repo_root, "dependencies", "extra_deps", "pre_train_github_deps.txt")
+  if not os.path.exists(github_deps_path):
+    raise FileNotFoundError(f"Github dependencies file not found at {github_deps_path}")
 
   # Check if 'uv' is available in the environment
   try:
@@ -46,22 +43,21 @@ def main():
     print(f"Stderr: {e.stderr.decode()}")
     sys.exit(1)
 
-  command = [
+  github_deps_command = [
       sys.executable,  # Use the current Python executable's pip to ensure the correct environment
       "-m",
       "uv",
       "pip",
       "install",
       "-r",
-      str(extra_deps_path),
+      str(github_deps_path),
       "--no-deps",
   ]
 
   try:
-    # Run the command
-    print(f"Installing extra dependencies: {' '.join(command)}")
-    _ = subprocess.run(command, check=True, capture_output=True, text=True)
-    print("Extra dependencies installed successfully!")
+    print(f"Installing Github dependencies: {' '.join(github_deps_command)}")
+    _ = subprocess.run(github_deps_command, check=True, capture_output=True, text=True)
+    print("Github dependencies installed successfully!")
   except subprocess.CalledProcessError as e:
     print("Failed to install extra dependencies.")
     print(f"Command '{' '.join(e.cmd)}' returned non-zero exit status {e.returncode}.")

src/dependencies/scripts/setup.sh

@@ -215,8 +215,8 @@ install_maxtext_with_deps() {
         dep_name='src/dependencies/requirements/generated_requirements/tpu-requirements.txt'
     fi
     echo "Installing requirements from $dep_name"
-    python3 -m uv pip install --resolution=lowest -r "$dep_name" \
-        -r 'src/dependencies/github_deps/pre_train_deps.txt'
+    python3 -m uv pip install --resolution=lowest -r "$dep_name"
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps
 
     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.dependencies.github_deps.install_post_train_deps
+    python3 -m src.dependencies.scripts.install_post_train_extra_deps
 }
 
 # ---------- Post-Training workflow installation ----------

src/maxtext/examples/rl_llama3_demo.ipynb

@@ -108,7 +108,7 @@
     "    \n",
     "    # Install MaxText and post-training dependencies\n",
     "    !uv pip install -e .[tpu-post-train] --resolution=lowest\n",
-    "    !install_maxtext_tpu_post_train_extra_deps"
+    "    !install_tpu_post_train_extra_deps"
    ]
   },
   {

src/maxtext/examples/sft_llama3_demo_tpu.ipynb

@@ -6,7 +6,7 @@
         "id": "iBmRfde4Kgv4"
       },
       "source": [
-        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/AI-Hypercomputer/maxtext/blob/main/src/maxtext/examples/sft_llama3_demo.ipynb)\n",
+        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/AI-Hypercomputer/maxtext/blob/main/src/maxtext/examples/sft_llama3_demo_tpu.ipynb)\n",
         "\n",
         "# Llama3.1-8B-Instruct Supervised Fine-Tuning (SFT) Demo\n"
       ]
@@ -120,7 +120,7 @@
         "    \n",
         "    # Install MaxText and post-training dependencies\n",
         "    !uv pip install -e .[tpu-post-train] --resolution=lowest\n",
-        "    !install_maxtext_tpu_post_train_extra_deps"
+        "    !install_tpu_post_train_extra_deps"
       ]
     },
     {

src/maxtext/examples/sft_qwen3_demo.ipynb

@@ -126,7 +126,7 @@
         "    \n",
         "    # Install MaxText and post-training dependencies\n",
         "    !uv pip install -e .[tpu-post-train] --resolution=lowest\n",
-        "    !install_maxtext_tpu_post_train_extra_deps"
+        "    !install_tpu_post_train_extra_deps"
       ]
     },
     {