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

@@ -61,6 +61,8 @@ 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.
 
+> **Note:** MaxText is only tested on Linux during releases.
+
 ## 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"