Support "modern" uv project management (uv init + uv add)

- Update MaxText installation scripts to detect 'uv.lock' and use 'uv add --frozen'.
- Consolidate uv detection and installation logic into a shared 'uv_utils.py' module.
- Improve uv detection robustness by prioritizing 'python -m uv' and path lookup.
- Update docs/install_maxtext.md with uv project management workflow instructions and fix outdated paths.

Author: igorts-git

docs/install_maxtext.md

@@ -106,6 +106,29 @@ uv pip install -e .[runner] --resolution=lowest
 
 After installation, you can verify the package is available with `python3 -c "import maxtext"` and run training jobs with `python3 -m maxtext.trainers.pre_train.train ...`.
 
+## UV Project management
+
+For simplicity this guide uses the traditional `uv pip install` syntax.
+If you are using the `uv` project management features (with a `pyproject.toml` and `uv.lock` in your own project), you would need to run your commands differently.
+You would need to use `uv add` instead of `uv pip install` and `uv run` as a prefix to all other commands.
+For example, to install and run MaxText in a `uv`-managed project:
+
+```bash
+# 1. Initialize your uv project
+mkdir my-maxtext-project && cd my-maxtext-project
+uv init
+
+# 2. Add MaxText as a dependency
+uv add maxtext[tpu]==0.2.1 --resolution=lowest
+
+# 3. Install MaxText's extra GitHub dependencies. These will be automatically added to your pyproject.toml
+uv run install_tpu_pre_train_extra_deps # This will be added to your pyproject.toml
+
+# 4. Run MaxText training for a few steps
+uv run python3 -m maxtext.trainers.pre_train.train run_name=${RUN_NAME?} base_output_directory=${BASE_OUTPUT_DIRECTORY?} \
+    dataset_type=synthetic steps=5 per_device_batch_size=1 model_name=llama2-7b
+```
+
 # Update MaxText dependencies
 
 ## Introduction
@@ -120,7 +143,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.
+3. **Update Project Files**: Copy the newly generated files into the `src/dependencies/requirements/generated_requirements/` directory.
 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.
 
@@ -176,8 +199,8 @@ After generating the new requirements, you need to update the files in the MaxTe
 
 1. **Copy the generated files:**
 
-   - 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`.
+   - Move `generated_tpu_artifacts/tpu-requirements.txt` to `src/dependencies/requirements/generated_requirements/tpu-requirements.txt`.
+   - Move `generated_gpu_artifacts/cuda12-requirements.txt` to `src/dependencies/requirements/generated_requirements/cuda12-requirements.txt`.
 
 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.

src/dependencies/scripts/install_post_train_extra_deps.py

@@ -19,14 +19,21 @@
 """
 
 import os
-import subprocess
-import sys
+
+# This block makes the script a bit more flexible. It allows `uv_utils` to be imported whether this module is run as a
+# standalone script or as part of a larger Python package. It also allows us to not worry whether the full package name
+# starts with "src." (this happens when running inside a docker image as part of setup.sh).
+try:
+  from . import uv_utils
+except ImportError:
+  import uv_utils
 
 
 def main():
   """
   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'.
+
+  It executes 'uv add' (if uv.lock is present) or 'uv pip install'.
   """
   os.environ["VLLM_TARGET_DEVICE"] = "tpu"
 
@@ -36,57 +43,17 @@ def main():
   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:
-    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)
-
-  github_deps_command = [
-      sys.executable,  # Use the current Python executable's pip to ensure the correct environment
-      "-m",
-      "uv",
-      "pip",
-      "install",
-      "-r",
-      str(github_deps_path),
-      "--no-deps",
-  ]
-
-  local_vllm_install_command = [
-      sys.executable,  # Use the current Python executable's pip to ensure the correct environment
-      "-m",
-      "uv",
-      "pip",
-      "install",
-      f"{repo_root}/maxtext/integration/vllm",  # MaxText on vllm installations
-      "--no-deps",
-  ]
-
-  try:
-    # Run the command to install Github dependencies
-    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)}")
-    _ = subprocess.run(local_vllm_install_command, check=True, capture_output=True, text=True)
-    print("MaxText vLLM dependency 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}.")
-    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)
+  # Install both requirements file and the local vLLM integration
+  # Note: The original logic in HEAD didn't use --editable for vLLM,
+  # but my commit did. I'll stick to my commit's logic if it was intended.
+  # Actually, let's see what the original logic in my commit was.
+  # In my commit: paths=[f"{repo_root}/maxtext/integration/vllm"], is_editable=True
+  # But the `run_install` signature is `(requirements_files=None, paths=None, editable_paths=None)`.
+  # Wait, I should check `uv_utils.py` again.
+  uv_utils.run_install(
+      requirements_files=[github_deps_path],
+      editable_paths=[os.path.join(repo_root, "maxtext", "integration", "vllm")],
+  )
 
 
 if __name__ == "__main__":

src/dependencies/scripts/install_pre_train_extra_deps.py

@@ -19,56 +19,29 @@
 """
 
 import os
-import subprocess
-import sys
+
+# This block makes the script a bit more flexible. It allows `uv_utils` to be imported whether this module is run as a
+# standalone script or as part of a larger Python package. It also allows us to not worry whether the full package name
+# starts with "src." (this happens when running inside a docker image as part of setup.sh).
+try:
+  from . import uv_utils
+except ImportError:
+  import uv_utils
 
 
 def main():
   """
   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'.
+
+  It executes 'uv add' (if uv.lock is present) or 'uv pip install'.
   """
   current_dir = os.path.dirname(os.path.abspath(__file__))
   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:
-    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)
-
-  github_deps_command = [
-      sys.executable,  # Use the current Python executable's pip to ensure the correct environment
-      "-m",
-      "uv",
-      "pip",
-      "install",
-      "-r",
-      str(github_deps_path),
-      "--no-deps",
-  ]
-
-  try:
-    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}.")
-    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)
+  uv_utils.run_install(requirements_files=[github_deps_path])
 
 
 if __name__ == "__main__":

src/dependencies/scripts/uv_utils.py

@@ -0,0 +1,126 @@
+# Copyright 2026 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.
+
+"""Helper utilities for working with uv in installation scripts."""
+
+import os
+import shutil
+import subprocess
+import sys
+
+
+def get_uv_command():
+  """
+  Returns the command to run uv, either as a binary in PATH or as a module.
+  Attempts to install uv via pip if not found.
+  """
+  # 1. Try finding 'uv' in PATH
+  uv_binary = shutil.which("uv")
+  if uv_binary:
+    return [uv_binary]
+
+  # 2. Try running it as a module
+  try:
+    subprocess.run([sys.executable, "-m", "uv", "--version"], check=True, capture_output=True)
+    return [sys.executable, "-m", "uv"]
+  except (subprocess.CalledProcessError, FileNotFoundError):
+    pass
+
+  # 3. Fall back to installing via pip
+  try:
+    print("uv not found in PATH or as a module. Attempting to install it via pip...")
+    subprocess.run([sys.executable, "-m", "pip", "install", "uv"], check=True, capture_output=True)
+    # Check PATH again after installation
+    uv_binary = shutil.which("uv")
+    if uv_binary:
+      return [uv_binary]
+    return [sys.executable, "-m", "uv"]
+  except subprocess.CalledProcessError as e:
+    print(f"Error installing uv via pip: {e}")
+    print(f"Stderr: {e.stderr.decode()}")
+    sys.exit(1)
+
+
+def run_install(
+    requirements_files=None, paths=None, editable_paths=None, resolution=None, index_url=None, find_links=None, pre=False
+):
+  """
+  Executes the appropriate uv install command (uv add or uv pip install).
+
+  Args:
+    requirements_files: List of paths to requirements.txt files.
+    paths: List of paths to local packages or directories (non-editable).
+    editable_paths: List of paths to local packages or directories (editable).
+    resolution: Strategy for selecting between compatible versions (e.g., 'lowest').
+    index_url: URL to use for package index.
+    find_links: Locations to search for candidate distributions.
+    pre: Whether to allow pre-release versions.
+  """
+  uv_command = get_uv_command()
+  is_uv_project = os.path.exists("uv.lock")
+
+  # Step 1: Standard installations
+  if requirements_files or paths:
+    if is_uv_project:
+      cmd = uv_command + ["add", "--frozen"]
+    else:
+      cmd = uv_command + ["pip", "install", "--no-deps"]
+
+    _apply_common_flags(cmd, requirements_files, paths, resolution, index_url, find_links, pre)
+    _execute_command(cmd)
+
+  # Step 2: Editable installations
+  if editable_paths:
+    if is_uv_project:
+      cmd = uv_command + ["add", "--frozen", "--editable"]
+    else:
+      cmd = uv_command + ["pip", "install", "--no-deps", "-e"]
+
+    _apply_common_flags(cmd, None, editable_paths, resolution, index_url, find_links, pre)
+    _execute_command(cmd)
+
+
+def _apply_common_flags(cmd, requirements_files, paths, resolution, index_url, find_links, pre):
+  """Applies common uv flags to the command."""
+  if requirements_files:
+    for req in requirements_files:
+      cmd.extend(["-r", str(req)])
+  if paths:
+    cmd.extend(paths)
+  if resolution:
+    cmd.extend(["--resolution", resolution])
+  if index_url:
+    cmd.extend(["--index-url", index_url])
+  if find_links:
+    cmd.extend(["--find-links", find_links])
+  if pre:
+    cmd.append("--pre")
+
+
+def _execute_command(cmd):
+  """Helper to execute a command with logging and error handling."""
+  try:
+    print(f"Executing: {' '.join(cmd)}")
+    subprocess.run(cmd, check=True, capture_output=True, text=True)
+    print("Success!")
+  except subprocess.CalledProcessError as e:
+    print(f"Command failed with 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: {e}")
+    sys.exit(1)