Skip to content

Commit 842f3ff

Browse files
melissawmmelissawm
committed
Review updating dependencies instructions
1 parent 89105c6 commit 842f3ff

2 files changed

Lines changed: 66 additions & 41 deletions

File tree

‎docs/development/update_dependencies.md‎

Lines changed: 65 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -18,40 +18,73 @@
1818

1919
## Introduction
2020

21-
This document provides a guide to updating dependencies in MaxText using the `seed-env` tool. `seed-env` helps generate deterministic and reproducible Python environments by creating fully-pinned `requirements.txt` files from a base set of requirements.
21+
This document provides a guide to updating dependencies in MaxText using the
22+
[seed-env](https://github.com/google-ml-infra/actions/tree/main/python_seed_env)
23+
tool. `seed-env` helps generate deterministic and reproducible Python
24+
environments by creating fully-pinned `requirements.txt` files from a base set
25+
of requirements.
2226

23-
Please keep dependencies updated throughout development. This will allow each commit to work properly from both a feature and dependency perspective. We will periodically upload commits to PyPI for stable releases. But it is also critical to keep dependencies in sync for users installing MaxText from source.
27+
Please keep dependencies updated throughout development. This will allow each
28+
commit to work properly from both a feature and dependency perspective. We will
29+
periodically upload commits to PyPI for stable releases. But it is also critical
30+
to keep dependencies in sync for users installing MaxText from source.
2431

25-
## Overview of the Process
32+
## Overview of the process
2633

2734
To update dependencies, you will follow these general steps:
2835

29-
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`).
30-
2. **Generate New Files**: Run the `seed-env` CLI tool to generate new, fully-pinned requirements files based on your changes.
31-
3. **Update Project Files**: Copy the newly generated files into the `generated_requirements/` directory.
32-
4. **Handle GitHub Dependencies**: Move any dependencies that are installed directly from GitHub from the generated files to `src/install_maxtext_extra_deps/extra_deps_from_github.txt`.
33-
5. **Verify**: Test the new dependencies to ensure the project installs and runs correctly.
36+
1. **Modify base requirements**: Update the desired dependencies in
37+
`src/dependencies/requirements/base_requirements/requirements.txt` or the hardware-specific files
38+
(`src/dependencies/requirements/base_requirements/tpu-base-requirements.txt`,
39+
`src/dependencies/requirements/base_requirements/gpu-base-requirements.txt`).
40+
2. **Find the JAX build commit hash**: The dependency generation process is
41+
pinned to a specific nightly build of JAX. You need to find the commit hash
42+
for the desired JAX build.
43+
3. **Generate the requirement files**: Run the `seed-env` CLI tool to generate
44+
new, fully-pinned requirements files based on your changes.
45+
4. **Update project files**: Copy the newly generated files into the
46+
`src/dependencies/requirements/generated_requirements/` directory. If
47+
necessary, also update any dependencies that are installed directly from
48+
GitHub from the generated files to `src/dependencies/extra_deps`.
49+
5. **Verify the new dependencies**: Test the new dependencies to ensure the
50+
project installs and runs correctly.
3451

3552
The following sections provide detailed instructions for each step.
3653

37-
## Step 1: Install seed-env
54+
## Step 0: Install `seed-env`
3855

39-
First, you need to install the `seed-env` command-line tool by running `pip install seed-env uv`. Or follow the instructions in the
40-
[seed-env repository](https://github.com/google-ml-infra/actions/tree/main/python_seed_env#install-the-seed-env-tool) if you want to build `seed-env` from source.
56+
First, you need to install the `seed-env` command-line tool by running
57+
`pip install seed-env uv`. Or follow the instructions in the
58+
[seed-env repository](https://github.com/google-ml-infra/actions/tree/main/python_seed_env#install-the-seed-env-tool)
59+
if you want to build `seed-env` from source.
4160

42-
## Step 2: Find the JAX Build Commit Hash
61+
## Step 1: Modify base requirements
4362

44-
The dependency generation process is pinned to a specific nightly build of JAX. You need to find the commit hash for the desired JAX build.
63+
Update the desired dependencies in
64+
`src/dependencies/requirements/base_requirements/requirements.txt` or the
65+
hardware-specific files
66+
(`src/dependencies/requirements/base_requirements/tpu-base-requirements.txt`,
67+
`src/dependencies/requirements/base_requirements/gpu-base-requirements.txt`).
4568

46-
You can find the latest commit hashes in the [JAX `build/` folder](https://github.com/jax-ml/jax/commits/main/build). Choose a recent, successful build and copy its full commit hash.
69+
## Step 2: Find the JAX build commit hash
4770

48-
## Step 3: Generate the Requirements Files
71+
The dependency generation process is pinned to a specific nightly build of JAX.
72+
You need to find the commit hash for the desired JAX build.
4973

50-
Next, run the `seed-env` CLI to generate the new requirements files. You will need to do this separately for the TPU and GPU environments. The generated files will be placed in a directory specified by `--output-dir`.
74+
You can find the latest commit hashes in the
75+
[JAX `build/` folder](https://github.com/jax-ml/jax/commits/main/build). Choose
76+
a recent, successful build and copy its full commit hash.
77+
78+
## Step 3: Generate the requirements files
79+
80+
Next, run the `seed-env` CLI to generate the new requirements files. You will
81+
need to do this separately for the TPU and GPU environments. The generated files
82+
will be placed in a directory specified by `--output-dir`.
5183

5284
### For TPU
5385

54-
Run the following command, replacing `<jax-build-commit-hash>` with the hash you copied in the previous step.
86+
Run the following command, replacing `<jax-build-commit-hash>` with the hash you
87+
copied in the previous step.
5588

5689
```bash
5790
seed-env \
@@ -69,7 +102,7 @@ Similarly, run the command for the GPU requirements.
69102

70103
```bash
71104
seed-env \
72-
--local-requirements=src/dependencies/requirements/base_requirements/cuda12-base-requirements.txt \
105+
--local-requirements=src/dependencies/requirements/base_requirements/gpu-base-requirements.txt \
73106
--host-name=MaxText \
74107
--seed-commit=<jax-build-commit-hash> \
75108
--python-version=3.12 \
@@ -78,37 +111,28 @@ seed-env \
78111
--output-dir=generated_gpu_artifacts
79112
```
80113

81-
## Step 4: Update Project Files
114+
## Step 4: Update project files
82115

83-
After generating the new requirements, you need to update the files in the MaxText repository.
116+
After generating the new requirements, you need to update the files in the
117+
MaxText repository.
84118

85119
1. **Copy the generated files:**
86120

87121
- Move `generated_tpu_artifacts/tpu-requirements.txt` to `generated_requirements/tpu-requirements.txt`.
88122
- Move `generated_gpu_artifacts/cuda12-requirements.txt` to `generated_requirements/cuda12-requirements.txt`.
89123

90-
2. **Update `extra_deps_from_github.txt` (if necessary):**
91-
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.
92-
93-
## Step 5: Verify the New Dependencies
124+
2. **Update `src/dependencies/extra_deps` (if necessary):**
125+
Currently, MaxText uses a few dependencies, such as `mlperf-logging` and
126+
`google-jetstream`, that are installed directly from GitHub source. These are
127+
defined in `base_requirements/requirements.txt`, and the `seed-env` tool will
128+
carry them over to the generated requirements files.
94129

95-
Finally, test that the new dependencies install correctly and that MaxText runs as expected.
130+
## Step 5: Verify the new dependencies
96131

97-
1. **Create a clean environment:** It's best to start with a fresh Python virtual environment.
132+
Finally, test that the new dependencies install correctly and that MaxText runs
133+
as expected.
98134

99-
```bash
100-
uv venv --python 3.12 --seed maxtext_venv
101-
source maxtext_venv/bin/activate
102-
```
103-
104-
2. **Run the setup script:** Execute `bash setup.sh` to install the new dependencies.
105-
106-
```bash
107-
# install the tpu package
108-
uv pip install -e .[tpu] --resolution=lowest
109-
# or install the gpu package by running the following line:
110-
# uv pip install -e .[cuda12] --resolution=lowest
111-
install_maxtext_github_deps
112-
```
135+
1. **Install MaxText:** Follow the instructions to
136+
[install MaxText from source](install-from-source).
113137

114-
3. **Run tests:** Run MaxText tests to ensure there are no regressions.
138+
2. **Run tests:** Run MaxText tests to ensure there are no regressions.

‎docs/install_maxtext.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,7 @@ benchmarks.
101101
```
102102
103103
(install-from-source)=
104+
104105
## From source
105106
106107
If you plan to contribute to MaxText or need the latest unreleased features,

0 commit comments

Comments
 (0)