From ea6c7852bd181606600e369af8919af251a724d4 Mon Sep 17 00:00:00 2001 From: Igor Tsvetkov Date: Fri, 20 Mar 2026 10:43:50 -0700 Subject: [PATCH] Fix macOS installation by dynamically skipping legacy MaxText directory --- build_hooks.py | 27 ++++++++++++++++--- docs/install_maxtext.md | 8 +++++- pyproject.toml | 7 +++-- .../runner-requirements.txt | 2 +- 4 files changed, 34 insertions(+), 10 deletions(-) diff --git a/build_hooks.py b/build_hooks.py index 616a27aff7..79a7c52d7c 100644 --- a/build_hooks.py +++ b/build_hooks.py @@ -15,6 +15,7 @@ """Custom build hooks for PyPI.""" import os +import sys from hatchling.builders.hooks.plugin.interface import BuildHookInterface TPU_REQUIREMENTS_PATH = "src/dependencies/requirements/generated_requirements/tpu-requirements.txt" @@ -33,9 +34,27 @@ def get_tpu_dependencies(): class CustomBuildHook(BuildHookInterface): - """A custom hook to inject TPU dependencies into the core wheel dependencies.""" + """A custom hook to handle platform-specific package configuration for MaxText.""" def initialize(self, version, build_data): # pylint: disable=unused-argument - tpu_deps = get_tpu_dependencies() - build_data["dependencies"] = tpu_deps - print(f"Successfully injected {len(tpu_deps)} TPU dependencies into the wheel's core requirements.") + """Adjusts the build_data dictionary to customize the wheel's package structure.""" + # The following TPU dependency injection is disabled because TPU-specific requirements + # are now managed via optional dependencies (extras) in pyproject.toml + # (e.g., pip install maxtext[tpu]). + # tpu_deps = get_tpu_dependencies() + # build_data["dependencies"] = tpu_deps + # print(f"Successfully injected {len(tpu_deps)} TPU dependencies into the wheel's core requirements.") + + # macOS specific logic to avoid case-sensitivity issues with MaxText and maxtext directories + build_data["force_include"] = build_data.get("force_include", {}) + if sys.platform == "darwin": + print("macOS detected. Skipping legacy MaxText shims to avoid case-sensitivity conflicts.") + # Always include the __init__.py in the lowercase 'maxtext' package on macOS. + # This ensures that 'import maxtext' (and thus 'import MaxText' on macOS) + # has the proper version and metadata. + build_data["force_include"]["src/MaxText/__init__.py"] = "maxtext/__init__.py" + else: + # On other platforms, include 'src/MaxText' as its own top-level package for legacy support. + # We do NOT add __init__.py to 'maxtext' here to maintain exact parity with previous builds. + print("Included src/MaxText as a top-level package for non-macOS platforms.") + build_data["force_include"]["src/MaxText"] = "MaxText" diff --git a/docs/install_maxtext.md b/docs/install_maxtext.md index 05f720b104..296d636083 100644 --- a/docs/install_maxtext.md +++ b/docs/install_maxtext.md @@ -24,7 +24,7 @@ MaxText offers following installation modes: 3. maxtext[tpu-post-train]. Used for post-training on TPUs. Currently, this option should also be used for running vllm_decode on TPUs. 4. maxtext[runner]. Used for building MaxText's Docker images and scheduling workloads through XPK. -## From PyPI (Recommended) +## From PyPI (Recommended on Linux) This is the easiest way to get started with the latest stable version. @@ -61,6 +61,12 @@ uv pip install maxtext[runner] --resolution=lowest > **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. +## macOS Installation + +**Note:** macOS is not officially supported and not QA tested during releases. + +Due to macOS's case-insensitive filesystem MaxText can only be installed from source using the `.[runner]` configuration option listed below. + ## From Source If you plan to contribute to MaxText or need the latest unreleased features, install from source. diff --git a/pyproject.toml b/pyproject.toml index a32239f9d0..92f650980c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,11 +40,10 @@ Repository = "https://github.com/AI-Hypercomputer/maxtext.git" allow-direct-references = true [tool.hatch.build.targets.wheel] -packages = ["src/MaxText", "src/maxtext", "src/dependencies"] +packages = ["src/maxtext", "src/dependencies"] -# TODO: Add this hook back when it handles device-type parsing -# [tool.hatch.build.targets.wheel.hooks.custom] -# path = "build_hooks.py" +[tool.hatch.build.targets.wheel.hooks.custom] +path = "build_hooks.py" [project.scripts] install_maxtext_tpu_github_deps = "dependencies.github_deps.install_pre_train_deps:main" diff --git a/src/dependencies/requirements/generated_requirements/runner-requirements.txt b/src/dependencies/requirements/generated_requirements/runner-requirements.txt index 557642a685..4be9455cd0 100644 --- a/src/dependencies/requirements/generated_requirements/runner-requirements.txt +++ b/src/dependencies/requirements/generated_requirements/runner-requirements.txt @@ -71,7 +71,7 @@ mpmath>=1.3.0 namex>=0.1.0 numpy-typing-compat>=20251206.2.0 numpy>=2.0.2 -nvidia-cuda-cccl>=13.1.115 +nvidia-cuda-cccl>=13.1.115 ; sys_platform == 'linux' oauthlib>=3.3.1 omegaconf>=2.3.0 opt-einsum>=3.4.0