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-managed project 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.
• .gitignore: Piggybacking on this PR I also added maxtext_venv to .gitignore. This was needed to be able to run python3 -m build to build a wheel for testing below.

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 installing the package artifact:

# Create a fresh standard venv
mkdir ~/clean_install
cd ~/clean_install

uv venv --python 3.12 test_legacy_venv
source test_legacy_venv/bin/activate

# Install the wheel built in Step 1
uv pip install ~/git/maxtext/dist/maxtext-0.2.1-py3-none-any.whl[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

# Train for a few steps.
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

# Try to build a docker image
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

# Cleanup
deactivate
cd -
rm -rf ~/clean_install

3. Verify Modern Installation (uv add)

Tests the new project management workflow using the package artifact:

# Create a new uv project
mkdir ~/clean_install && cd ~/clean_install
uv init

# Add MaxText as a dependency using the wheel
uv add ~/git/maxtext/dist/maxtext-*.whl --extra tpu --resolution=lowest

# Run extra deps script (should detect uv.lock and use uv add)
uv run 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

# Train 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

# Build the TPU dependency image
# This will trigger setup.sh inside the container
uv run 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

# Cleanup
cd -
rm -rf ~/clean_install

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!