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

Update MaxText's installation scripts and setup process to detect
'uv.lock' in the current working directory. When detected, the scripts
now use 'uv add --frozen' to correctly track extra GitHub-based
dependencies in the user's pyproject.toml.

Key changes:
- Updated install_pre_train_deps.py and install_post_train_deps.py
  to support 'uv add'.
- Updated setup.sh to use 'uv add' for requirements while maintaining
  'uv pip install -e .' for the package itself to avoid conflicts.
- Updated docs/install_maxtext.md with instructions for the modern
  uv workflow and fixed various outdated command names and paths.

Author: igorts-git

docs/install_maxtext.md

@@ -45,7 +45,7 @@ install_maxtext_tpu_github_deps
 
 # Option 2: Installing maxtext[cuda12]
 uv pip install maxtext[cuda12] --resolution=lowest
-install_maxtext_cuda12_github_dep
+install_maxtext_cuda12_github_deps
 
 # Option 3: Installing maxtext[tpu-post-train]
 uv pip install maxtext[tpu-post-train] --resolution=lowest
@@ -55,10 +55,27 @@ install_maxtext_tpu_post_train_extra_deps
 uv pip install maxtext[runner] --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.
 
+## Modern UV Project (Recommended for New Projects)
+
+If you are starting a new project and want to use `uv`'s project management features (with a `pyproject.toml` and `uv.lock` in your own project), you can use `uv add`. MaxText's helper scripts will detect your `uv.lock` and correctly add their extra dependencies to your `pyproject.toml`.
+
+```bash
+# 1. Initialize your project
+mkdir my-maxtext-project && cd my-maxtext-project
+uv init
+
+# 2. Add MaxText as a dependency
+uv add maxtext[tpu] --resolution=lowest
+
+# 3. Install MaxText's extra GitHub dependencies
+# These will be automatically added to your pyproject.toml
+install_maxtext_tpu_github_deps
+```
+
 > **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.
 
 ## macOS Installation
@@ -88,7 +105,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
@@ -114,7 +131,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.
 
@@ -170,8 +187,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

@@ -49,26 +49,48 @@ def main():
     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",
-  ]
+  # Check for uv.lock in the current directory to decide between 'uv add' and 'uv pip install'
+  if os.path.exists("uv.lock"):
+    print("Found uv.lock. Using 'uv add' to track dependencies in pyproject.toml.")
+    command = [
+        sys.executable,
+        "-m",
+        "uv",
+        "add",
+        "--frozen",
+        "-r",
+        str(extra_deps_path),
+    ]
+    local_vllm_install_command = [
+        sys.executable,
+        "-m",
+        "uv",
+        "add",
+        "--frozen",
+        "--editable",
+        f"{repo_root}/maxtext/integration/vllm",
+    ]
+  else:
+    print("uv.lock not found. Using 'uv pip install' for direct environment installation.")
+    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

src/dependencies/github_deps/install_pre_train_deps.py

@@ -46,16 +46,30 @@ def main():
     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",
-  ]
+  # Check for uv.lock in the current directory to decide between 'uv add' and 'uv pip install'
+  if os.path.exists("uv.lock"):
+    print("Found uv.lock. Using 'uv add' to track dependencies in pyproject.toml.")
+    command = [
+        sys.executable,
+        "-m",
+        "uv",
+        "add",
+        "--frozen",
+        "-r",
+        str(extra_deps_path),
+    ]
+  else:
+    print("uv.lock not found. Using 'uv pip install' for direct environment installation.")
+    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

src/dependencies/scripts/setup.sh

@@ -199,6 +199,9 @@ install_maxtext_package_without_deps() {
     # to significantly faster image builds.
     if [ -f 'pyproject.toml' ]; then
         echo "Installing MaxText package without installing the dependencies (already installed)"
+        # We use 'uv pip install' for the editable install to avoid self-dependency
+        # conflicts that can occur with 'uv add' when the package name matches
+        # the project name in pyproject.toml.
         python3 -m uv pip install --no-deps -e .
     fi
 }
@@ -215,8 +218,12 @@ 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'
+    if [ -f "uv.lock" ]; then
+        python3 -m uv add --frozen -r "$dep_name" -r 'src/dependencies/github_deps/pre_train_deps.txt'
+    else
+        python3 -m uv pip install --resolution=lowest -r "$dep_name" \
+            -r 'src/dependencies/github_deps/pre_train_deps.txt'
+    fi
 
     install_maxtext_package_without_deps
 }
@@ -229,7 +236,11 @@ install_post_training_deps() {
     echo "Setting up MaxText post-training workflow for $DEVICE device"
     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"
+    if [ -f "uv.lock" ]; then
+        python3 -m uv add --frozen -r "$dep_name"
+    else
+        python3 -m uv pip install --resolution=lowest -r "$dep_name"
+    fi
     python3 -m src.dependencies.github_deps.install_post_train_deps
 }