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: Created a helper module uv_utils which provides a run_install helper function to use the correct uv command to run based on the presence of a uv.lock file in the current working directory. Changed install_pre_train_deps.py and install_post_train_deps.py to use the new uv_utils.run_install().
• 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 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_tpu_pre_train_extra_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 WORKFLOW=pre-training

# 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

# Test that the other installation modes work
install_tpu_post_train_extra_deps
build_maxtext_docker_image DEVICE=tpu MODE=stable WORKFLOW=post-training

install_cuda12_pre_train_extra_deps 
build_maxtext_docker_image DEVICE=gpu MODE=stable WORKFLOW=pre-training

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

# Test that the other installation modes work
uv run install_tpu_post_train_extra_deps
uv run build_maxtext_docker_image DEVICE=tpu MODE=stable WORKFLOW=post-training

uv run install_cuda12_pre_train_extra_deps 
uv run build_maxtext_docker_image DEVICE=gpu MODE=stable WORKFLOW=pre-training

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