Skip to content

Commit 3ac6587

Browse files
committed
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.
1 parent b5e833c commit 3ac6587

7 files changed

Lines changed: 173 additions & 102 deletions

File tree

.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,7 @@ celerybeat.pid
114114
# Environments
115115
.env
116116
.venv
117+
maxtext_venv/
117118
env/
118119
venv/
119120
ENV/

build_hooks.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
"""Custom build hooks for PyPI."""
1616

1717
import os
18+
import sys
1819
from hatchling.builders.hooks.plugin.interface import BuildHookInterface
1920

2021
TPU_REQUIREMENTS_PATH = "src/dependencies/requirements/generated_requirements/tpu-requirements.txt"

docs/install_maxtext.md

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ MaxText offers following installation modes:
2424
3. maxtext[tpu-post-train]. Used for post-training on TPUs. Currently, this option should also be used for running vllm_decode on TPUs.
2525
4. maxtext[runner]. Used for building MaxText's Docker images and scheduling workloads through XPK.
2626

27-
## From PyPI (Recommended)
27+
## From PyPI (Recommended on Linux)
2828

2929
This is the easiest way to get started with the latest stable version.
3030

@@ -55,7 +55,7 @@ install_maxtext_tpu_post_train_extra_deps
5555
uv pip install maxtext[runner]==0.2.1 --resolution=lowest
5656
```
5757

58-
> **Note:** The `install_maxtext_tpu_github_deps`, `install_maxtext_cuda12_github_dep`, and
58+
> **Note:** The `install_maxtext_tpu_github_deps`, `install_maxtext_cuda12_github_deps`, and
5959
> `install_maxtext_tpu_post_train_extra_deps` commands are temporarily required to install dependencies directly from GitHub
6060
> that are not yet available on PyPI. As shown above, choose the one that corresponds to your use case.
6161
@@ -86,7 +86,7 @@ install_maxtext_tpu_github_deps
8686

8787
# Option 2: Installing .[cuda12]
8888
uv pip install -e .[cuda12] --resolution=lowest
89-
install_maxtext_cuda12_github_dep
89+
install_maxtext_cuda12_github_deps
9090

9191
# Option 3: Installing .[tpu-post-train]
9292
uv pip install -e .[tpu-post-train] --resolution=lowest
@@ -98,6 +98,29 @@ uv pip install -e .[runner] --resolution=lowest
9898

9999
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 ...`.
100100

101+
## UV Project management
102+
103+
For simplicity this guide uses the traditional `uv pip install` syntax.
104+
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.
105+
You would need to use `uv add` instead of `uv pip install` and `uv run` as a prefix to all other commands.
106+
For example:
107+
108+
```bash
109+
# 1. Initialize your uv-managed project
110+
mkdir my-maxtext-project && cd my-maxtext-project
111+
uv init
112+
113+
# 2. Add MaxText as a dependency
114+
uv add maxtext[tpu] --resolution=lowest
115+
116+
# 3. Install MaxText's extra GitHub dependencies. These will be automatically added to your pyproject.toml
117+
uv run install_maxtext_tpu_github_deps
118+
119+
# 4. Run MaxText training for a few steps
120+
uv run python3 -m maxtext.trainers.pre_train.train run_name=${RUN_NAME?} base_output_directory=${BASE_OUTPUT_DIRECTORY?} \
121+
dataset_type=synthetic steps=5 per_device_batch_size=1 model_name=llama2-7b
122+
```
123+
101124
# Update MaxText dependencies
102125

103126
## Introduction
@@ -112,7 +135,7 @@ To update dependencies, you will follow these general steps:
112135

113136
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`).
114137
2. **Generate New Files**: Run the `seed-env` CLI tool to generate new, fully-pinned requirements files based on your changes.
115-
3. **Update Project Files**: Copy the newly generated files into the `generated_requirements/` directory.
138+
3. **Update Project Files**: Copy the newly generated files into the `src/dependencies/requirements/generated_requirements/` directory.
116139
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`.
117140
5. **Verify**: Test the new dependencies to ensure the project installs and runs correctly.
118141

@@ -168,8 +191,8 @@ After generating the new requirements, you need to update the files in the MaxTe
168191

169192
1. **Copy the generated files:**
170193

171-
- Move `generated_tpu_artifacts/tpu-requirements.txt` to `generated_requirements/tpu-requirements.txt`.
172-
- Move `generated_gpu_artifacts/cuda12-requirements.txt` to `generated_requirements/cuda12-requirements.txt`.
194+
- Move `generated_tpu_artifacts/tpu-requirements.txt` to `src/dependencies/requirements/generated_requirements/tpu-requirements.txt`.
195+
- Move `generated_gpu_artifacts/cuda12-requirements.txt` to `src/dependencies/requirements/generated_requirements/cuda12-requirements.txt`.
173196

174197
2. **Update `pre_train_deps.txt` (if necessary):**
175198
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

Lines changed: 10 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -21,16 +21,19 @@
2121
"""
2222

2323
import os
24-
import subprocess
25-
import sys
24+
25+
try:
26+
from . import uv_utils
27+
except ImportError:
28+
import uv_utils
2629

2730

2831
def main():
2932
"""
3033
Installs extra dependencies specified in post_train_deps.txt using uv.
3134
3235
This script looks for 'post_train_deps.txt' relative to its own location.
33-
It executes 'uv pip install -r <path_to_extra_deps.txt> --resolution=lowest'.
36+
It executes 'uv add' (if uv.lock is present) or 'uv pip install'.
3437
"""
3538
os.environ["VLLM_TARGET_DEVICE"] = "tpu"
3639

@@ -40,57 +43,10 @@ def main():
4043
if not os.path.exists(extra_deps_path):
4144
raise FileNotFoundError(f"Dependencies file not found at {extra_deps_path}")
4245

43-
# Check if 'uv' is available in the environment
44-
try:
45-
subprocess.run([sys.executable, "-m", "pip", "install", "uv"], check=True, capture_output=True)
46-
subprocess.run([sys.executable, "-m", "uv", "--version"], check=True, capture_output=True)
47-
except subprocess.CalledProcessError as e:
48-
print(f"Error checking uv version: {e}")
49-
print(f"Stderr: {e.stderr.decode()}")
50-
sys.exit(1)
51-
52-
command = [
53-
sys.executable, # Use the current Python executable's pip to ensure the correct environment
54-
"-m",
55-
"uv",
56-
"pip",
57-
"install",
58-
"-r",
59-
str(extra_deps_path),
60-
"--no-deps",
61-
]
62-
63-
local_vllm_install_command = [
64-
sys.executable, # Use the current Python executable's pip to ensure the correct environment
65-
"-m",
66-
"uv",
67-
"pip",
68-
"install",
69-
f"{repo_root}/maxtext/integration/vllm", # MaxText on vllm installations
70-
"--no-deps",
71-
]
72-
73-
try:
74-
# Run the command to install Github dependencies
75-
print(f"Installing extra dependencies: {' '.join(command)}")
76-
_ = subprocess.run(command, check=True, capture_output=True, text=True)
77-
print("Extra dependencies installed successfully!")
78-
79-
# Run the command to install the MaxText vLLM directory
80-
print(f"Installing MaxText vLLM dependency: {' '.join(local_vllm_install_command)}")
81-
_ = subprocess.run(local_vllm_install_command, check=True, capture_output=True, text=True)
82-
print("MaxText vLLM dependency installed successfully!")
83-
except subprocess.CalledProcessError as e:
84-
print("Failed to install extra dependencies.")
85-
print(f"Command '{' '.join(e.cmd)}' returned non-zero exit status {e.returncode}.")
86-
print("--- Stderr ---")
87-
print(e.stderr)
88-
print("--- Stdout ---")
89-
print(e.stdout)
90-
sys.exit(e.returncode)
91-
except (OSError, FileNotFoundError) as e:
92-
print(f"An OS-level error occurred while trying to run uv: {e}")
93-
sys.exit(1)
46+
# Install both requirements file and the local vLLM integration
47+
uv_utils.run_install(
48+
requirements_files=[extra_deps_path], paths=[f"{repo_root}/maxtext/integration/vllm"], is_editable=True
49+
)
9450

9551

9652
if __name__ == "__main__":

src/dependencies/github_deps/install_pre_train_deps.py

Lines changed: 7 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -21,58 +21,26 @@
2121
"""
2222

2323
import os
24-
import subprocess
25-
import sys
24+
25+
try:
26+
from . import uv_utils
27+
except ImportError:
28+
import uv_utils
2629

2730

2831
def main():
2932
"""
3033
Installs extra dependencies specified in pre_train_deps.txt using uv.
3134
3235
This script looks for 'pre_train_deps.txt' relative to its own location.
33-
It executes 'uv pip install -r <path_to_extra_deps.txt> --resolution=lowest'.
36+
It executes 'uv add' (if uv.lock is present) or 'uv pip install'.
3437
"""
3538
current_dir = os.path.dirname(os.path.abspath(__file__))
3639
extra_deps_path = os.path.join(current_dir, "pre_train_deps.txt")
3740
if not os.path.exists(extra_deps_path):
3841
raise FileNotFoundError(f"Dependencies file not found at {extra_deps_path}")
3942

40-
# Check if 'uv' is available in the environment
41-
try:
42-
subprocess.run([sys.executable, "-m", "pip", "install", "uv"], check=True, capture_output=True)
43-
subprocess.run([sys.executable, "-m", "uv", "--version"], check=True, capture_output=True)
44-
except subprocess.CalledProcessError as e:
45-
print(f"Error checking uv version: {e}")
46-
print(f"Stderr: {e.stderr.decode()}")
47-
sys.exit(1)
48-
49-
command = [
50-
sys.executable, # Use the current Python executable's pip to ensure the correct environment
51-
"-m",
52-
"uv",
53-
"pip",
54-
"install",
55-
"-r",
56-
str(extra_deps_path),
57-
"--no-deps",
58-
]
59-
60-
try:
61-
# Run the command
62-
print(f"Installing extra dependencies: {' '.join(command)}")
63-
_ = subprocess.run(command, check=True, capture_output=True, text=True)
64-
print("Extra dependencies installed successfully!")
65-
except subprocess.CalledProcessError as e:
66-
print("Failed to install extra dependencies.")
67-
print(f"Command '{' '.join(e.cmd)}' returned non-zero exit status {e.returncode}.")
68-
print("--- Stderr ---")
69-
print(e.stderr)
70-
print("--- Stdout ---")
71-
print(e.stdout)
72-
sys.exit(e.returncode)
73-
except (OSError, FileNotFoundError) as e:
74-
print(f"An OS-level error occurred while trying to run uv: {e}")
75-
sys.exit(1)
43+
uv_utils.run_install(requirements_files=[extra_deps_path])
7644

7745

7846
if __name__ == "__main__":
Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Helper utilities for working with uv in installation scripts."""
16+
17+
import os
18+
import shutil
19+
import subprocess
20+
import sys
21+
22+
23+
def get_uv_command():
24+
"""
25+
Returns the command to run uv, either as a binary in PATH or as a module.
26+
Attempts to install uv via pip if not found.
27+
"""
28+
# 1. Try finding 'uv' in PATH
29+
uv_binary = shutil.which("uv")
30+
if uv_binary:
31+
return [uv_binary]
32+
33+
# 2. Try running it as a module
34+
try:
35+
subprocess.run([sys.executable, "-m", "uv", "--version"], check=True, capture_output=True)
36+
return [sys.executable, "-m", "uv"]
37+
except (subprocess.CalledProcessError, FileNotFoundError):
38+
pass
39+
40+
# 3. Fall back to installing via pip
41+
try:
42+
print("uv not found in PATH or as a module. Attempting to install it via pip...")
43+
subprocess.run([sys.executable, "-m", "pip", "install", "uv"], check=True, capture_output=True)
44+
# Check PATH again after installation
45+
uv_binary = shutil.which("uv")
46+
if uv_binary:
47+
return [uv_binary]
48+
return [sys.executable, "-m", "uv"]
49+
except subprocess.CalledProcessError as e:
50+
print(f"Error installing uv via pip: {e}")
51+
print(f"Stderr: {e.stderr.decode()}")
52+
sys.exit(1)
53+
54+
55+
def run_install(requirements_files=None, paths=None, editable_paths=None):
56+
"""
57+
Executes the appropriate uv install command (uv add or uv pip install).
58+
59+
Args:
60+
requirements_files: List of paths to requirements.txt files.
61+
paths: List of paths to local packages or directories (non-editable).
62+
editable_paths: List of paths to local packages or directories (editable).
63+
"""
64+
uv_command = get_uv_command()
65+
is_uv_project = os.path.exists("uv.lock")
66+
67+
# We run installations in two steps if we have both standard and editable items,
68+
# because 'uv add --editable' cannot be mixed with non-local requirements.
69+
70+
# Step 1: Standard installations
71+
if requirements_files or paths:
72+
if is_uv_project:
73+
cmd = uv_command + ["add", "--frozen"]
74+
else:
75+
cmd = uv_command + ["pip", "install", "--no-deps"]
76+
77+
if requirements_files:
78+
for req in requirements_files:
79+
cmd.extend(["-r", str(req)])
80+
if paths:
81+
cmd.extend(paths)
82+
83+
_execute_command(cmd)
84+
85+
# Step 2: Editable installations
86+
if editable_paths:
87+
if is_uv_project:
88+
cmd = uv_command + ["add", "--frozen", "--editable"]
89+
else:
90+
cmd = uv_command + ["pip", "install", "--no-deps", "-e"]
91+
92+
cmd.extend(editable_paths)
93+
_execute_command(cmd)
94+
95+
96+
def _execute_command(cmd):
97+
"""Helper to execute a command with logging and error handling."""
98+
try:
99+
print(f"Executing: {' '.join(cmd)}")
100+
subprocess.run(cmd, check=True, capture_output=True, text=True)
101+
print("Success!")
102+
except subprocess.CalledProcessError as e:
103+
print(f"Command failed with exit status {e.returncode}.")
104+
print("--- Stderr ---")
105+
print(e.stderr)
106+
print("--- Stdout ---")
107+
print(e.stdout)
108+
sys.exit(e.returncode)
109+
except (OSError, FileNotFoundError) as e:
110+
print(f"An OS-level error occurred: {e}")
111+
sys.exit(1)

src/dependencies/scripts/setup.sh

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,9 @@ install_maxtext_package_without_deps() {
199199
# to significantly faster image builds.
200200
if [ -f 'pyproject.toml' ]; then
201201
echo "Installing MaxText package without installing the dependencies (already installed)"
202+
# We use 'uv pip install' for the editable install to avoid self-dependency
203+
# conflicts that can occur with 'uv add' when the package name matches
204+
# the project name in pyproject.toml.
202205
python3 -m uv pip install --no-deps -e .
203206
fi
204207
}
@@ -215,8 +218,12 @@ install_maxtext_with_deps() {
215218
dep_name='src/dependencies/requirements/generated_requirements/tpu-requirements.txt'
216219
fi
217220
echo "Installing requirements from $dep_name"
218-
python3 -m uv pip install --resolution=lowest -r "$dep_name" \
219-
-r 'src/dependencies/github_deps/pre_train_deps.txt'
221+
if [ -f "uv.lock" ]; then
222+
python3 -m uv add --frozen -r "$dep_name" -r 'src/dependencies/github_deps/pre_train_deps.txt'
223+
else
224+
python3 -m uv pip install --resolution=lowest -r "$dep_name" \
225+
-r 'src/dependencies/github_deps/pre_train_deps.txt'
226+
fi
220227

221228
install_maxtext_package_without_deps
222229
}
@@ -229,7 +236,11 @@ install_post_training_deps() {
229236
echo "Setting up MaxText post-training workflow for $DEVICE device"
230237
dep_name='src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt'
231238
echo "Installing requirements from $dep_name"
232-
python3 -m uv pip install --resolution=lowest -r "$dep_name"
239+
if [ -f "uv.lock" ]; then
240+
python3 -m uv add --frozen -r "$dep_name"
241+
else
242+
python3 -m uv pip install --resolution=lowest -r "$dep_name"
243+
fi
233244
python3 -m src.dependencies.github_deps.install_post_train_deps
234245
}
235246

0 commit comments

Comments
 (0)