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

@@ -55,7 +55,7 @@ install_maxtext_tpu_post_train_extra_deps
 uv pip install maxtext[runner]==0.2.1 --resolution=lowest
 ```
 
-> **Note:** The `install_maxtext_tpu_github_deps`, `install_maxtext_cuda12_github_dep`, and
+> **Note:** The `install_maxtext_tpu_github_deps`, `install_maxtext_cuda12_github_deps`, and
 > `install_maxtext_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.
 
@@ -86,7 +86,7 @@ install_maxtext_tpu_github_deps
 
 # Option 2: Installing .[cuda12]
 uv pip install -e .[cuda12] --resolution=lowest
-install_maxtext_cuda12_github_dep
+install_maxtext_cuda12_github_deps
 
 # Option 3: Installing .[tpu-post-train]
 uv pip install -e .[tpu-post-train] --resolution=lowest
@@ -98,6 +98,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:
+
+```bash
+# 1. Initialize your uv-managed 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_maxtext_tpu_github_deps
+
+# 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
@@ -112,7 +135,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.
 
@@ -168,8 +191,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/github_deps/install_post_train_deps.py

@@ -21,16 +21,22 @@
 """
 
 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 post_train_deps.txt using uv.
 
   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'.
+  It executes 'uv add' (if uv.lock is present) or 'uv pip install'.
   """
   os.environ["VLLM_TARGET_DEVICE"] = "tpu"
 
@@ -40,57 +46,10 @@ def main():
   if not os.path.exists(extra_deps_path):
     raise FileNotFoundError(f"Dependencies file not found at {extra_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)
-
-  command = [
-      sys.executable,  # Use the current Python executable's pip to ensure the correct environment
-      "-m",
-      "uv",
-      "pip",
-      "install",
-      "-r",
-      str(extra_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 extra dependencies: {' '.join(command)}")
-    _ = subprocess.run(command, check=True, capture_output=True, text=True)
-    print("Extra 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
+  uv_utils.run_install(
+      requirements_files=[extra_deps_path], paths=[f"{repo_root}/maxtext/integration/vllm"], is_editable=True
+  )
 
 
 if __name__ == "__main__":

src/dependencies/github_deps/install_pre_train_deps.py

@@ -21,58 +21,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 pre_train_deps.txt using uv.
 
   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'.
+  It executes 'uv add' (if uv.lock is present) or 'uv pip install'.
   """
   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}")
 
-  # 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_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!")
-  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=[extra_deps_path])
 
 
 if __name__ == "__main__":

src/dependencies/github_deps/uv_utils.py

@@ -0,0 +1,111 @@
+# 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):
+  """
+  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).
+  """
+  uv_command = get_uv_command()
+  is_uv_project = os.path.exists("uv.lock")
+
+  # We run installations in two steps if we have both standard and editable items,
+  # because 'uv add --editable' cannot be mixed with non-local requirements.
+
+  # 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"]
+
+    if requirements_files:
+      for req in requirements_files:
+        cmd.extend(["-r", str(req)])
+    if paths:
+      cmd.extend(paths)
+
+    _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"]
+
+    cmd.extend(editable_paths)
+    _execute_command(cmd)
+
+
+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)