Skip to content

Commit b8113cc

Browse files
melissawmmelissawm
committed
Review updating dependencies instructions
1 parent 6c68242 commit b8113cc

1 file changed

Lines changed: 64 additions & 41 deletions

File tree

‎docs/development/update_dependencies.md‎

Lines changed: 64 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -18,40 +18,72 @@
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`, `src/dependencies/requirements/base_requirements/gpu-base-requirements.txt`).
39+
2. **Find the JAX build commit hash**: The dependency generation process is
40+
pinned to a specific nightly build of JAX. You need to find the commit hash
41+
for the desired JAX build.
42+
3. **Generate the requirement files**: Run the `seed-env` CLI tool to generate
43+
new, fully-pinned requirements files based on your changes.
44+
4. **Update project files**: Copy the newly generated files into the
45+
`src/dependencies/requirements/generated_requirements/` directory. If
46+
necessary, also update any dependencies that are installed directly from
47+
GitHub from the generated files to `src/dependencies/extra_deps`.
48+
5. **Verify the new dependencies**: Test the new dependencies to ensure the
49+
project installs and runs correctly.
3450

3551
The following sections provide detailed instructions for each step.
3652

37-
## Step 1: Install seed-env
53+
## Step 0: Install `seed-env`
3854

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.
55+
First, you need to install the `seed-env` command-line tool by running
56+
`pip install seed-env uv`. Or follow the instructions in the
57+
[seed-env repository](https://github.com/google-ml-infra/actions/tree/main/python_seed_env#install-the-seed-env-tool)
58+
if you want to build `seed-env` from source.
4159

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

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.
62+
Update the desired dependencies in
63+
`src/dependencies/requirements/base_requirements/requirements.txt` or the
64+
hardware-specific files
65+
(`src/dependencies/requirements/base_requirements/tpu-base-requirements.txt`,
66+
`src/dependencies/requirements/base_requirements/gpu-base-requirements.txt`).
4567

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.
68+
## Step 2: Find the JAX build commit hash
4769

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

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

5283
### For TPU
5384

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

5688
```bash
5789
seed-env \
@@ -69,7 +101,7 @@ Similarly, run the command for the GPU requirements.
69101

70102
```bash
71103
seed-env \
72-
--local-requirements=src/dependencies/requirements/base_requirements/cuda12-base-requirements.txt \
104+
--local-requirements=src/dependencies/requirements/base_requirements/gpu-base-requirements.txt \
73105
--host-name=MaxText \
74106
--seed-commit=<jax-build-commit-hash> \
75107
--python-version=3.12 \
@@ -78,37 +110,28 @@ seed-env \
78110
--output-dir=generated_gpu_artifacts
79111
```
80112

81-
## Step 4: Update Project Files
113+
## Step 4: Update project files
82114

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

85118
1. **Copy the generated files:**
86119

87120
- Move `generated_tpu_artifacts/tpu-requirements.txt` to `generated_requirements/tpu-requirements.txt`.
88121
- Move `generated_gpu_artifacts/cuda12-requirements.txt` to `generated_requirements/cuda12-requirements.txt`.
89122

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

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

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

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-
```
134+
1. **Install MaxText:** Follow the instructions to
135+
[install MaxText from source](install-from-source).
113136

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

0 commit comments

Comments
 (0)