Fix macOS installation by dynamically skipping legacy MaxText directory

Author: igorts-git

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"

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,10 @@ 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
+
+Due to macOS's case-insensitive filesystem we only support installing MaxText 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.

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"

src/dependencies/requirements/generated_requirements/cuda12-requirements.txt

@@ -233,9 +233,9 @@ toml>=0.10.2
 tomlkit>=0.13.3
 toolz>=1.1.0
 tqdm>=4.67.1
-transformer-engine-cu12>=2.9.0
-transformer-engine-jax>=2.9.0
-transformer-engine>=2.9.0
+transformer-engine-cu12>=2.9.0 ; sys_platform == 'linux'
+transformer-engine-jax>=2.9.0 ; sys_platform == 'linux'
+transformer-engine>=2.9.0 ; sys_platform == 'linux'
 transformers>=4.57.3
 treescope>=0.1.10
 typeguard>=2.13.3

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

src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt

@@ -164,7 +164,7 @@ lark>=1.2.2
 latex2sympy2_extended>=1.11.0
 libclang>=18.1.1
 libcst>=1.8.6
-libtpu>=0.0.32
+libtpu>=0.0.32 ; platform_machine == 'x86_64' and sys_platform == 'linux'
 llguidance>=1.3.0
 llvmlite>=0.45.1
 lm-format-enforcer>=0.11.3
@@ -204,21 +204,21 @@ nodeenv>=1.9.1
 numba>=0.62.1
 numpy>=2.2.6
 numpy-typing-compat>=20250818.2.2
-nvidia-cublas-cu12>=12.8.4.1
-nvidia-cuda-cupti-cu12>=12.8.90
-nvidia-cuda-nvrtc-cu12>=12.8.93
-nvidia-cuda-runtime-cu12>=12.8.90
-nvidia-cudnn-cu12>=9.10.2.21
-nvidia-cufft-cu12>=11.3.3.83
-nvidia-cufile-cu12>=1.13.1.3
-nvidia-curand-cu12>=10.3.9.90
-nvidia-cusolver-cu12>=11.7.3.90
-nvidia-cusparse-cu12>=12.5.8.93
-nvidia-cusparselt-cu12>=0.7.1
-nvidia-nccl-cu12>=2.27.5
-nvidia-nvjitlink-cu12>=12.8.93
-nvidia-nvshmem-cu12>=3.3.20
-nvidia-nvtx-cu12>=12.8.90
+nvidia-cublas-cu12>=12.8.4.1 ; sys_platform == 'linux'
+nvidia-cuda-cupti-cu12>=12.8.90 ; sys_platform == 'linux'
+nvidia-cuda-nvrtc-cu12>=12.8.93 ; sys_platform == 'linux'
+nvidia-cuda-runtime-cu12>=12.8.90 ; sys_platform == 'linux'
+nvidia-cudnn-cu12>=9.10.2.21 ; sys_platform == 'linux'
+nvidia-cufft-cu12>=11.3.3.83 ; sys_platform == 'linux'
+nvidia-cufile-cu12>=1.13.1.3 ; sys_platform == 'linux'
+nvidia-curand-cu12>=10.3.9.90 ; sys_platform == 'linux'
+nvidia-cusolver-cu12>=11.7.3.90 ; sys_platform == 'linux'
+nvidia-cusparse-cu12>=12.5.8.93 ; sys_platform == 'linux'
+nvidia-cusparselt-cu12>=0.7.1 ; sys_platform == 'linux'
+nvidia-nccl-cu12>=2.27.5 ; sys_platform == 'linux'
+nvidia-nvjitlink-cu12>=12.8.93 ; sys_platform == 'linux'
+nvidia-nvshmem-cu12>=3.3.20 ; sys_platform == 'linux'
+nvidia-nvtx-cu12>=12.8.90 ; sys_platform == 'linux'
 oauthlib>=3.3.1
 omegaconf>=2.3.0
 openai>=2.24.0
@@ -369,7 +369,7 @@ tqdm>=4.67.3
 traitlets>=5.14.3
 transformers>=4.57.1
 treescope>=0.1.10
-triton>=3.5.0
+triton>=3.5.0 ; sys_platform == 'linux'
 typeguard>=2.13.3
 typer>=0.24.1
 typing-inspect>=0.9.0