diff --git a/.github/workflows/build_and_test_maxtext.yml b/.github/workflows/build_and_test_maxtext.yml index f758f428d5..7a1fe03114 100644 --- a/.github/workflows/build_and_test_maxtext.yml +++ b/.github/workflows/build_and_test_maxtext.yml @@ -262,7 +262,7 @@ jobs: tf_force_gpu_allow_growth: false container_resource_option: "--privileged" is_scheduled_run: ${{ github.event_name == 'schedule' }} - extra_pip_deps_file: 'src/install_maxtext_extra_deps/extra_post_train_base_deps_from_github.txt' + extra_pip_deps_file: 'src/dependencies/github_deps/post_train_base_deps.txt' maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }} maxtext_post_training_tpu_unit_tests: @@ -284,7 +284,7 @@ jobs: tf_force_gpu_allow_growth: false container_resource_option: "--privileged" is_scheduled_run: ${{ github.event_name == 'schedule' }} - extra_pip_deps_file: 'src/install_maxtext_extra_deps/extra_post_train_base_deps_from_github.txt' + extra_pip_deps_file: 'src/dependencies/github_deps/post_train_base_deps.txt' maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }} maxtext_gpu_integration_tests: diff --git a/.github/workflows/run_pathways_tests.yml b/.github/workflows/run_pathways_tests.yml index ecf4182f22..92a38bccf1 100644 --- a/.github/workflows/run_pathways_tests.yml +++ b/.github/workflows/run_pathways_tests.yml @@ -85,7 +85,7 @@ jobs: source .venv/bin/activate maxtext_wheel=$(ls maxtext-*-py3-none-any.whl 2>/dev/null) uv pip install ${maxtext_wheel}[tpu] --resolution=lowest - uv pip install -r src/install_maxtext_extra_deps/extra_deps_from_github.txt + uv pip install -r src/dependencies/github_deps/pre_train_deps.txt python3 --version python3 -m pip freeze - name: Copy test assets files diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index 0b85ccc94f..b19b8f41dd 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.github/workflows/run_tests_against_package.yml @@ -96,7 +96,7 @@ jobs: source .venv/bin/activate maxtext_wheel=$(ls maxtext-*-py3-none-any.whl 2>/dev/null) uv pip install ${maxtext_wheel}[${MAXTEXT_PACKAGE_EXTRA}] --resolution=lowest - uv pip install -r src/install_maxtext_extra_deps/extra_deps_from_github.txt + uv pip install -r src/dependencies/github_deps/pre_train_deps.txt python3 --version python3 -m pip freeze uv pip install pytest-cov diff --git a/docs/conf.py b/docs/conf.py index b6479fcd21..cc71573ee1 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -120,8 +120,8 @@ os.path.join("run_maxtext", "run_maxtext_via_multihost_runner.md"), os.path.join("reference", "core_concepts", "llm_calculator.ipynb"), os.path.join("reference", "api_generated", "modules.rst"), - os.path.join("reference", "api_generated", "install_maxtext_extra_deps.rst"), - os.path.join("reference", "api_generated", "install_maxtext_extra_deps.install_github_deps.rst"), + os.path.join("reference", "api_generated", "dependencies.github_deps.rst"), + os.path.join("reference", "api_generated", "dependencies.github_deps.install_pre_train_deps.rst"), ] diff --git a/docs/install_maxtext.md b/docs/install_maxtext.md index 7f779db7a2..05f720b104 100644 --- a/docs/install_maxtext.md +++ b/docs/install_maxtext.md @@ -111,7 +111,7 @@ To update dependencies, you will follow these general steps: 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`). 2. **Generate New Files**: Run the `seed-env` CLI tool to generate new, fully-pinned requirements files based on your changes. 3. **Update Project Files**: Copy the newly generated files into the `generated_requirements/` directory. -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`. +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`. 5. **Verify**: Test the new dependencies to ensure the project installs and runs correctly. The following sections provide detailed instructions for each step. @@ -169,7 +169,7 @@ After generating the new requirements, you need to update the files in the MaxTe - Move `generated_tpu_artifacts/tpu-requirements.txt` to `generated_requirements/tpu-requirements.txt`. - Move `generated_gpu_artifacts/cuda12-requirements.txt` to `generated_requirements/cuda12-requirements.txt`. -2. **Update `extra_deps_from_github.txt` (if necessary):** +2. **Update `pre_train_deps.txt` (if necessary):** 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. ## Step 5: Verify the New Dependencies diff --git a/pyproject.toml b/pyproject.toml index 547a85d3da..840cbca67f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,15 +40,15 @@ Repository = "https://github.com/AI-Hypercomputer/maxtext.git" allow-direct-references = true [tool.hatch.build.targets.wheel] -packages = ["src/MaxText", "src/maxtext", "src/install_maxtext_extra_deps", "src/dependencies"] +packages = ["src/MaxText", "src/maxtext", "src/dependencies/github_deps", "src/dependencies"] # TODO: Add this hook back when it handles device-type parsing # [tool.hatch.build.targets.wheel.hooks.custom] # path = "build_hooks.py" [project.scripts] -install_maxtext_tpu_github_deps = "install_maxtext_extra_deps.install_github_deps:main" -install_maxtext_cuda12_github_deps = "install_maxtext_extra_deps.install_github_deps:main" -install_maxtext_tpu_post_train_extra_deps = "install_maxtext_extra_deps.install_post_train_extra_deps:main" +install_maxtext_tpu_github_deps = "dependencies.github_deps.install_pre_train_deps:main" +install_maxtext_cuda12_github_deps = "dependencies.github_deps.install_pre_train_deps:main" +install_maxtext_tpu_post_train_extra_deps = "dependencies.github_deps.install_post_train_deps:main" build_maxtext_docker_image = "dependencies.scripts.build_maxtext_docker_image:main" upload_maxtext_docker_image = "dependencies.scripts.upload_maxtext_docker_image:main" diff --git a/src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile b/src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile index 7a08541a86..b9d7cdf80a 100644 --- a/src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile +++ b/src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile @@ -53,9 +53,9 @@ ENV MAXTEXT_REPO_ROOT=/deps WORKDIR /deps # Copy setup files and dependency files separately for better caching +COPY ${PACKAGE_DIR}/dependencies/github_deps/ src/dependencies/github_deps/ COPY ${PACKAGE_DIR}/dependencies/requirements/ src/dependencies/requirements/ COPY ${PACKAGE_DIR}/dependencies/scripts/ src/dependencies/scripts/ -COPY ${PACKAGE_DIR}/install_maxtext_extra_deps/ src/install_maxtext_extra_deps/ COPY ${PACKAGE_DIR}/maxtext/integration/vllm/ src/MaxText/integration/vllm/ # Install dependencies - these steps are cached unless the copied files change diff --git a/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile b/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile index 81c3f2d8b2..427640fefe 100644 --- a/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile +++ b/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile @@ -50,9 +50,9 @@ ENV MAXTEXT_REPO_ROOT=/deps WORKDIR /deps # Copy setup files and dependency files separately for better caching +COPY ${PACKAGE_DIR}/dependencies/github_deps/ src/dependencies/github_deps/ COPY ${PACKAGE_DIR}/dependencies/requirements/ src/dependencies/requirements/ COPY ${PACKAGE_DIR}/dependencies/scripts/ src/dependencies/scripts/ -COPY ${PACKAGE_DIR}/install_maxtext_extra_deps/ src/install_maxtext_extra_deps/ COPY ${PACKAGE_DIR}/maxtext/integration/vllm/ src/maxtext/integration/vllm/ # Copy the custom libtpu.so file if it exists diff --git a/src/install_maxtext_extra_deps/__init__.py b/src/dependencies/github_deps/__init__.py similarity index 100% rename from src/install_maxtext_extra_deps/__init__.py rename to src/dependencies/github_deps/__init__.py diff --git a/src/install_maxtext_extra_deps/install_post_train_extra_deps.py b/src/dependencies/github_deps/install_post_train_deps.py similarity index 84% rename from src/install_maxtext_extra_deps/install_post_train_extra_deps.py rename to src/dependencies/github_deps/install_post_train_deps.py index 0e1749cd0a..199cb7bed2 100644 --- a/src/install_maxtext_extra_deps/install_post_train_extra_deps.py +++ b/src/dependencies/github_deps/install_post_train_deps.py @@ -15,7 +15,7 @@ """Installs extra dependencies from a requirements file using uv. This script is designed to be run to install dependencies specified in -'extra_post_train_deps_from_github.txt', which is expected to be in the same directory. +'post_train_deps.txt', which is expected to be in the same directory. It first ensures 'uv' is installed and then uses it to install the packages listed in the requirements file. """ @@ -28,22 +28,22 @@ def main(): """ - Installs extra dependencies specified in extra_post_train_deps_from_github.txt using uv. + Installs extra dependencies specified in post_train_deps.txt using uv. - This script looks for 'extra_post_train_deps_from_github.txt' relative to its own location. + This script looks for 'post_train_deps.txt' relative to its own location. It executes 'uv pip install -r --resolution=lowest'. """ script_dir = Path(__file__).resolve().parent os.environ["VLLM_TARGET_DEVICE"] = "tpu" - # Adjust this path if your extra_post_train_deps_from_github.txt is in a different location, - # e.g., script_dir / "data" / "extra_post_train_deps_from_github.txt" - extra_deps_file = script_dir / "extra_post_train_deps_from_github.txt" + # Adjust this path if your post_train_deps.txt is in a different location, + # e.g., script_dir / "data" / "post_train_deps.txt" + extra_deps_file = script_dir / "post_train_deps.txt" if not extra_deps_file.exists(): print(f"Error: '{extra_deps_file}' not found.") - print("Please ensure 'extra_post_train_deps_from_github.txt' is in the correct location relative to the script.") + print("Please ensure 'post_train_deps.txt' is in the correct location relative to the script.") sys.exit(1) # Check if 'uv' is available in the environment try: diff --git a/src/install_maxtext_extra_deps/install_github_deps.py b/src/dependencies/github_deps/install_pre_train_deps.py similarity index 83% rename from src/install_maxtext_extra_deps/install_github_deps.py rename to src/dependencies/github_deps/install_pre_train_deps.py index 10b0b6d632..22a893bedd 100644 --- a/src/install_maxtext_extra_deps/install_github_deps.py +++ b/src/dependencies/github_deps/install_pre_train_deps.py @@ -15,7 +15,7 @@ """Installs extra dependencies from a requirements file using uv. This script is designed to be run to install dependencies specified in -'extra_deps_from_github.txt', which is expected to be in the same directory. +'pre_train_deps.txt', which is expected to be in the same directory. It first ensures 'uv' is installed and then uses it to install the packages listed in the requirements file. """ @@ -27,20 +27,20 @@ def main(): """ - Installs extra dependencies specified in extra_deps_from_github.txt using uv. + Installs extra dependencies specified in pre_train_deps.txt using uv. - This script looks for 'extra_deps_from_github.txt' relative to its own location. + This script looks for 'pre_train_deps.txt' relative to its own location. It executes 'uv pip install -r --resolution=lowest'. """ script_dir = Path(__file__).resolve().parent - # Adjust this path if your extra_deps_from_github.txt is in a different location, - # e.g., script_dir / "data" / "extra_deps_from_github.txt" - extra_deps_file = script_dir / "extra_deps_from_github.txt" + # Adjust this path if your pre_train_deps.txt is in a different location, + # e.g., script_dir / "data" / "pre_train_deps.txt" + extra_deps_file = script_dir / "pre_train_deps.txt" if not extra_deps_file.exists(): print(f"Error: '{extra_deps_file}' not found.") - print("Please ensure 'extra_deps_from_github.txt' is in the correct location relative to the script.") + print("Please ensure 'pre_train_deps.txt' is in the correct location relative to the script.") sys.exit(1) # Check if 'uv' is available in the environment try: diff --git a/src/install_maxtext_extra_deps/extra_post_train_base_deps_from_github.txt b/src/dependencies/github_deps/post_train_base_deps.txt similarity index 100% rename from src/install_maxtext_extra_deps/extra_post_train_base_deps_from_github.txt rename to src/dependencies/github_deps/post_train_base_deps.txt diff --git a/src/install_maxtext_extra_deps/extra_post_train_deps_from_github.txt b/src/dependencies/github_deps/post_train_deps.txt similarity index 90% rename from src/install_maxtext_extra_deps/extra_post_train_deps_from_github.txt rename to src/dependencies/github_deps/post_train_deps.txt index b990ed71d7..7bd07a345e 100644 --- a/src/install_maxtext_extra_deps/extra_post_train_deps_from_github.txt +++ b/src/dependencies/github_deps/post_train_deps.txt @@ -1,4 +1,4 @@ --r extra_post_train_base_deps_from_github.txt +-r post_train_base_deps.txt google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip tpu-inference @ https://github.com/vllm-project/tpu-inference/archive/0cae84fc9a883ba1bde02d4f07930e6af9e92958.zip diff --git a/src/install_maxtext_extra_deps/extra_deps_from_github.txt b/src/dependencies/github_deps/pre_train_deps.txt similarity index 100% rename from src/install_maxtext_extra_deps/extra_deps_from_github.txt rename to src/dependencies/github_deps/pre_train_deps.txt diff --git a/src/dependencies/scripts/setup.sh b/src/dependencies/scripts/setup.sh index a4e24937c0..13dd91a5e0 100644 --- a/src/dependencies/scripts/setup.sh +++ b/src/dependencies/scripts/setup.sh @@ -216,7 +216,7 @@ install_maxtext_with_deps() { fi echo "Installing requirements from $dep_name" python3 -m uv pip install --resolution=lowest -r "$dep_name" \ - -r 'src/install_maxtext_extra_deps/extra_deps_from_github.txt' + -r 'src/dependencies/github_deps/pre_train_deps.txt' install_maxtext_package_without_deps } @@ -230,7 +230,7 @@ install_post_training_deps() { dep_name='src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt' echo "Installing requirements from $dep_name" python3 -m uv pip install --resolution=lowest -r "$dep_name" - python3 -m src.install_maxtext_extra_deps.install_post_train_extra_deps + python3 -m src.dependencies.github_deps.install_post_train_deps } # ---------- Post-Training workflow installation ---------- diff --git a/src/maxtext/examples/demo_decoding.ipynb b/src/maxtext/examples/demo_decoding.ipynb index 4bd4dbad93..14c08899bc 100644 --- a/src/maxtext/examples/demo_decoding.ipynb +++ b/src/maxtext/examples/demo_decoding.ipynb @@ -86,7 +86,7 @@ "\n", "# Install MaxText and dependencies\n", "!uv pip install maxtext --resolution=lowest\n", - "!python3 -m MaxText.install_maxtext_extra_deps\n", + "!python3 -m dependencies.github_deps\n", "\n", "# Use nest_asyncio to allow nested event loops in notebooks\n", "!uv pip install nest_asyncio\n",