Add post-training dependencies to MaxText

Author: bvandermoon

dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt

@@ -41,8 +41,11 @@ allow-direct-references = true
 [tool.hatch.build.targets.wheel]
 packages = ["src/MaxText", "src/maxtext", "src/install_maxtext_extra_deps"]
 
-[tool.hatch.build.targets.wheel.hooks.custom]
-path = "build_hooks.py"
+# 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_github_deps = "install_maxtext_extra_deps.install_github_deps:main"
+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_github_deps = "install_maxtext_extra_deps.install_post_train_github_deps:main"

pyproject.toml

@@ -0,0 +1,5 @@
+google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip
+google-tunix @ https://github.com/google/tunix/archive/b12123542511e45814ec3ff798ec4fb30b0b48e8.zip
+mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip
+tpu-inference @ https://github.com/vllm-project/tpu-inference/archive/eb2bd54c70f443775d8870caf56b78ab0b671269.zip
+vllm @ git+https://github.com/vllm-project/vllm@3bbb2046ff320395c80c139e55e7c1947c3fb5e1

src/install_maxtext_extra_deps/extra_post_train_deps_from_github.txt

@@ -0,0 +1,94 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""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.
+It first ensures 'uv' is installed and then uses it to install the packages
+listed in the requirements file.
+"""
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+
+def main():
+  """
+  Installs extra dependencies specified in extra_post_train_deps_from_github.txt using uv.
+
+  This script looks for 'extra_post_train_deps_from_github.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"
+
+  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.")
+    sys.exit(1)
+  # Check if 'uv' is available in the environment
+  try:
+    subprocess.run([sys.executable, "-m", "pip", "install", "uv"], check=True, capture_output=True)
+    subprocess.run([sys.executable, "-m", "uv", "--version"], check=True, capture_output=True)
+  except subprocess.CalledProcessError as e:
+    print(f"Error checking uv version: {e}")
+    print(f"Stderr: {e.stderr.decode()}")
+    sys.exit(1)
+
+  command = [
+      sys.executable,  # Use the current Python executable's pip to ensure the correct environment
+      "-m",
+      "uv",
+      "pip",
+      "install",
+      "-r",
+      str(extra_deps_file),
+      "--no-deps",
+  ]
+
+  print(f"Installing extra dependencies from '{extra_deps_file}' using uv...")
+  print(f"Running command: {' '.join(command)}")
+
+  try:
+    # Run the command
+    process = subprocess.run(command, check=True, capture_output=True, text=True)
+    print("Extra dependencies installed successfully!")
+    print("--- Output from uv ---")
+    print(process.stdout)
+    if process.stderr:
+      print("--- Errors/Warnings from uv (if any) ---")
+      print(process.stderr)
+  except subprocess.CalledProcessError as e:
+    print("Failed to install extra dependencies.")
+    print(f"Command '{' '.join(e.cmd)}' returned non-zero exit status {e.returncode}.")
+    print("--- Stderr ---")
+    print(e.stderr)
+    print("--- Stdout ---")
+    print(e.stdout)
+    sys.exit(e.returncode)
+  except (OSError, FileNotFoundError) as e:
+    print(f"An OS-level error occurred while trying to run uv: {e}")
+    sys.exit(1)
+
+
+if __name__ == "__main__":
+  main()

src/install_maxtext_extra_deps/install_post_train_github_deps.py

@@ -14,6 +14,11 @@
 
 """Unit tests for Distillation Checkpointing logic."""
 
+import pytest
+
+pytest.importorskip("tunix")
+pytestmark = [pytest.mark.tpu_only]
+
 import json
 import os
 import shutil

tests/unit/distillation_checkpointing_test.py

@@ -15,6 +15,7 @@
 """Tests for training and data loading hooks for SFT"""
 import pytest
 
+pytest.importorskip("tunix")
 pytestmark = [pytest.mark.tpu_only, pytest.mark.external_training]
 
 import jax

tests/unit/sft_hooks_test.py

@@ -111,6 +111,8 @@ def compare_sharding_jsons(json1: dict, model1_name: str, json2: dict, model2_na
   return has_diff
 
 
+# Requires JAX TPU support to generate the simulated TPU topology.
+@pytest.mark.tpu_only
 @pytest.mark.parametrize("model_name, topology, num_slice", TEST_CASES)
 def test_sharding_dump_for_model(model_name: str, topology: str, num_slice: str) -> None:
   """
@@ -217,6 +219,8 @@ def abstract_state_and_shardings(request):
 class TestGetAbstractState:
   """Test class for get_abstract_state function and sharding comparison."""
 
+  # Requires JAX TPU support to generate the simulated TPU topology.
+  @pytest.mark.tpu_only
   def test_get_abstract_state_sharding(self, abstract_state_and_shardings):  # pylint: disable=redefined-outer-name
     """Tests that get_abstract_state returns a state with the correct abstract structure and compares sharding."""
 

tests/unit/sharding_compare_test.py

@@ -15,6 +15,11 @@
 
 """Unit tests for the Distillation Trainer."""
 
+import pytest
+
+pytest.importorskip("tunix")
+pytestmark = [pytest.mark.tpu_only]
+
 import shutil
 import tempfile
 import unittest