Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 6be43fd

Browse files
committed
Create a CI workflow that build maxtext package and run tests against it
Save the built maxtext distributions as action artifacts. Artifacts are retained within the workflow run for 90 days. Create a docker image for the package testing Test with local configs/assets/test_assets files for now. Move github deps Install github deps using the cli command Update README.md and setup.sh
1 parent 94539bb commit 6be43fd

7 files changed

Lines changed: 145 additions & 6 deletions

File tree

.github/workflows/run_tests_against_package.yml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ jobs:
6868
- name: Install the maxtext wheel
6969
run: |
7070
python3 -m uv pip install maxtext-*-py3-none-any.whl --resolution=lowest --system
71+
install_maxtext_github_deps
7172
python3 --version
7273
python3 -m pip freeze
7374
- name: Copy test assets files

README.md

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,41 @@ MaxText is the launching point for ambitious LLM projects both in research and p
2626

2727
Check out our [Read The Docs site](https://maxtext.readthedocs.io/en/latest/) or directly [Get Started](https://maxtext.readthedocs.io/en/latest/tutorials/first_run.html) with your first MaxText run. If you’re interested in Diffusion models (Wan 2.1, Flux, etc), see the [MaxDiffusion](https://github.com/AI-Hypercomputer/maxdiffusion) repository in our AI Hypercomputer GitHub organization.
2828

29+
## Installation
30+
31+
We recommend installing MaxText inside a Python virtual environment.
32+
33+
### From PyPI (Recommended)
34+
This is the easiest way to get started with the latest stable version.
35+
36+
```bash
37+
# 1. Install uv, a fast Python package installer
38+
pip install uv
39+
40+
# 2. Install MaxText and its dependencies
41+
uv pip install maxtext --resolution=lowest
42+
install_maxtext_github_deps
43+
```
44+
> **Note:** The `install_maxtext_github_deps` command is temporarily required to install dependencies directly from GitHub that are not yet available on PyPI.
45+
46+
> **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.
47+
48+
### From Source
49+
If you plan to contribute to MaxText or need the latest unreleased features, install from source.
50+
51+
```bash
52+
# 1. Clone the repository
53+
git clone https://github.com/AI-Hypercomputer/maxtext.git
54+
cd maxtext
55+
56+
# 2. Install dependencies in editable mode
57+
pip install uv
58+
uv pip install -e . --resolution=lowest
59+
install_maxtext_github_deps
60+
```
61+
62+
After installation, you can verify the package is available with `python3 -c "import MaxText"` and run training jobs with `python3 -m MaxText.train ...`.
63+
2964
## 🔥 Latest news 🔥
3065

3166
* \[September 5, 2025\] MaxText has moved to an `src` layout as part of [RESTRUCTURE.md](RESTRUCTURE.md). For existing environments, please run `pip install -e .` from MaxText root.

pyproject.toml

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,6 @@ dependencies = [
5353
"contourpy>=1.3.2",
5454
"coverage>=7.10.6",
5555
"cycler>=0.12.1",
56-
"datasets @ https://github.com/huggingface/datasets/archive/6790e138c00b87a1ddc72184f89e7814cf784360.zip",
5756
"decorator>=5.2.1",
5857
"dill>=0.3.8",
5958
"distlib>=0.4.0",
@@ -89,7 +88,6 @@ dependencies = [
8988
"google-cloud-storage>=2.19.0",
9089
"google-crc32c>=1.7.1",
9190
"google-genai>=1.33.0",
92-
"google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/daedc21c393f23449fb54ddc4f75fca34348ea9c.zip",
9391
"google-pasta>=0.2.0",
9492
"google-resumable-media>=2.7.2",
9593
"googleapis-common-protos>=1.70.0",
@@ -136,7 +134,6 @@ dependencies = [
136134
"ml-collections>=1.1.0",
137135
"ml-dtypes>=0.5.3",
138136
"ml-goodput-measurement>=0.0.14",
139-
"mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip",
140137
"more-itertools>=10.8.0",
141138
"mpmath>=1.3.0",
142139
"msgpack>=1.1.1",
@@ -196,7 +193,6 @@ dependencies = [
196193
"pytype>=2024.10.11",
197194
"pytz>=2025.2",
198195
"pyyaml>=6.0.2",
199-
"qwix @ https://github.com/google/qwix/archive/f2fd7b9114ff8d09e5b0131a453351578502da8a.zip",
200196
"regex>=2025.9.1",
201197
"requests-oauthlib>=2.0.0",
202198
"requests>=2.32.5",
@@ -236,7 +232,6 @@ dependencies = [
236232
"tqdm>=4.67.1",
237233
"transformers>=4.56.1",
238234
"treescope>=0.1.10",
239-
"tunix @ https://github.com/google/tunix/archive/d770659621eb16ef6588268e26fa687fa068df20.zip",
240235
"typing-extensions>=4.14.1",
241236
"typing-inspection>=0.4.1",
242237
"tzdata>=2025.2",
@@ -263,4 +258,7 @@ Repository = "https://github.com/AI-Hypercomputer/maxtext.git"
263258
allow-direct-references = true
264259

265260
[tool.hatch.build.targets.wheel]
266-
packages = ["src/MaxText"]
261+
packages = ["src/MaxText", "src/maxtext_extra_deps"]
262+
263+
[project.scripts]
264+
install_maxtext_github_deps = "maxtext_extra_deps.install_github_deps:main"

setup.sh

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,7 @@ python3 -m uv pip install --no-cache-dir -U -r requirements.txt
144144
# Install maxtext package
145145
if [ -f 'pyproject.toml' ]; then
146146
python3 -m uv pip install -e . --no-deps --resolution=lowest
147+
install_maxtext_github_deps
147148
fi
148149

149150
# Uninstall existing jax, jaxlib and libtpu-nightly

src/maxtext_extra_deps/__init__.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
# Copyright 2025 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+
Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
datasets @ https://github.com/huggingface/datasets/archive/6790e138c00b87a1ddc72184f89e7814cf784360.zip
2+
google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/daedc21c393f23449fb54ddc4f75fca34348ea9c.zip
3+
mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip
4+
qwix @ https://github.com/google/qwix/archive/f2fd7b9114ff8d09e5b0131a453351578502da8a.zip
5+
tunix @ https://github.com/google/tunix/archive/d770659621eb16ef6588268e26fa687fa068df20.zip
Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
# Copyright 2025 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+
import subprocess
16+
import sys
17+
from pathlib import Path
18+
19+
def main():
20+
"""
21+
Installs extra dependencies specified in extra_deps.txt using uv.
22+
23+
This script looks for 'extra_deps.txt' relative to its own location.
24+
It executes 'uv pip install -r <path_to_extra_deps.txt> --resolution=lowest'.
25+
"""
26+
script_dir = Path(__file__).resolve().parent
27+
28+
# Adjust this path if your extra_deps.txt is in a different location,
29+
# e.g., script_dir / "data" / "extra_deps_from_github.txt"
30+
extra_deps_file = script_dir / "extra_deps_from_github.txt"
31+
32+
if not extra_deps_file.exists():
33+
print(f"Error: '{extra_deps_file}' not found.")
34+
print("Please ensure 'extra_deps.txt' is in the correct location relative to the script.")
35+
sys.exit(1)
36+
# Check if 'uv' is available in the system's PATH
37+
try:
38+
subprocess.run(["uv", "--version"], check=True, capture_output=True)
39+
except FileNotFoundError:
40+
print("Error: 'uv' command not found.")
41+
subprocess.run(["pip", "install", "uv"], check=True)
42+
subprocess.run(["uv", "--version"], check=True, capture_output=True)
43+
sys.exit(1)
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_file),
57+
"--no-deps",
58+
]
59+
60+
print(f"Installing extra dependencies from '{extra_deps_file}' using uv...")
61+
print(f"Running command: {' '.join(command)}")
62+
63+
try:
64+
# Run the command
65+
process = subprocess.run(command, check=True, capture_output=True, text=True)
66+
print("Extra dependencies installed successfully!")
67+
print("--- Output from uv ---")
68+
print(process.stdout)
69+
if process.stderr:
70+
print("--- Errors/Warnings from uv (if any) ---")
71+
print(process.stderr)
72+
except subprocess.CalledProcessError as e:
73+
print(f"Failed to install extra dependencies.")
74+
print(f"Command '{' '.join(e.cmd)}' returned non-zero exit status {e.returncode}.")
75+
print("--- Stderr ---")
76+
print(e.stderr)
77+
print("--- Stdout ---")
78+
print(e.stdout)
79+
sys.exit(e.returncode)
80+
except Exception as e:
81+
print(f"An unexpected error occurred: {e}")
82+
sys.exit(1)
83+
84+
if __name__ == "__main__":
85+
main()

0 commit comments

Comments
 (0)