[WIP Do not review yet] Support modern uv project management (uv init + uv add)

[igorts-git]

Description

Objective

Enable users to use MaxText seamlessly within their own uv-managed projects. Previously, MaxText's installation scripts relied exclusively on uv pip install, which does not update a project's pyproject.toml or uv.lock. This PR introduces automatic detection of uv projects and uses uv add to correctly track "extra" dependencies (like GitHub-based requirements) when applicable.

Background

There is a modern uv usage mode that is different from the traditional uv pip install flow.
The new mode requires uv init and uv add commands instead of uv pip install.
This PR makes it compatible with the new flow.

Key changes:

• uv.lock Detection: Added logic to install_pre_train_deps.py, install_post_train_deps.py, and setup.sh to check for the presence of a uv.lock file in the current working directory.
• uv add Integration:
  • When a uv project is detected, the scripts now use uv add --frozen -r <requirements> to add extra dependencies to the user's pyproject.toml.
  • The --frozen flag is used to ensure that only the requested dependencies are added/updated without triggering a full environment resolution/sync, which is critical given MaxText's complex optional dependency tree.
• Backward Compatibility: Maintained the existing uv pip install workflow as a fallback for standard virtual environments, Docker builds, and MaxText development.
• Documentation: Updated docs/install_maxtext.md with instructions for the modern uv workflow and corrected several outdated paths and command names.

Pair-programmed with Gemini-CLI.

FIXES: #3441

Tests

1. Build a new Wheel

To verify that the package metadata and entry points are correctly defined:

# Build the wheel
python3 -m build
# The wheel will be in the dist/ directory
ls -l dist/*.whl

2. Verify Legacy Installation (uv pip install)

Ensures no regressions for existing users:

# Create a fresh standard venv
uv venv --python 3.12 test_legacy_venv
source test_legacy_venv/bin/activate

# Install from source
uv pip install -e .[tpu] --resolution=lowest

# Run extra deps script (should detect no uv.lock and use uv pip)
install_maxtext_tpu_github_deps

# Verify dependencies are installed but no uv.lock is created
uv pip show google-jetstream
ls uv.lock # Should not exist
deactivate

3. Verify Modern Installation (uv add)

Tests the new project management workflow:

# Create a new uv project
mkdir my_maxtext_app && cd my_maxtext_app
uv init

# Add MaxText as a dependency
# (Points to your local clone)
uv add "maxtext[tpu] @ file://$(pwd -P)/../path/to/maxtext" --resolution=lowest

# Run extra deps script (should detect uv.lock and use uv add)
# Ensure the script is in your path or use the entry point
install_maxtext_tpu_github_deps

# Verify pyproject.toml was updated
grep "google-jetstream" pyproject.toml
grep -A 5 "tool.uv.sources" pyproject.toml

# Verify uv.lock exists
ls -l uv.lock

4. Local Smoke Training

Verify that the installed environment is functional:

# Run a short smoke test training
python3 -m MaxText.train MaxText/configs/base.yml \
    run_name=test_smoke \
    base_output_directory=/tmp/maxtext_test \
    dataset_type=synthetic \
    steps=5 \
    enable_checkpointing=False \
    enable_goodput_recording=False

5. Docker Image Validation

Verify that Docker builds still work correctly and produce consistent results:

# Build the TPU dependency image
# This will trigger setup.sh inside the container
build_maxtext_docker_image DEVICE=tpu MODE=stable

# Capture the list of installed packages in the new image
docker run --rm maxtext_base_image uv pip list > image_pkgs.txt

# Compare with a known good list or verify critical packages
grep "google-jetstream" image_pkgs.txt
grep "jax" image_pkgs.txt

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!