Merge pull request #3469 from AI-Hypercomputer:igorts/fix_macos

PiperOrigin-RevId: 890048921

Author: Google-ML-Automation

.gitignore

@@ -3,6 +3,7 @@ tmp/
 logs/
 .venvs
 venv*
+maxtext_venv/
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

build_hooks.py

@@ -33,9 +33,24 @@ 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."""
+
+    # Avoid case-sensitivity issues with `MaxText` and `maxtext` directories on case-insensitive platforms.
+    build_data["force_include"] = build_data.get("force_include", {})
+
+    # Detect case-insensitivity by checking if this file can be accessed via a different case.
+    # On case-insensitive filesystems flipping the case of the filename still points to the same file.
+    is_case_insensitive = os.path.exists(__file__.upper()) and os.path.exists(__file__.lower())
+
+    if is_case_insensitive:
+      print("Skipping legacy MaxText shims to avoid case-sensitivity conflicts.")
+      # Always include the __init__.py in the lowercase 'maxtext'.
+      # This ensures that 'import maxtext' (and thus 'import MaxText') 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.
+      build_data["force_include"]["src/MaxText"] = "MaxText"

docs/install_maxtext.md

@@ -61,6 +61,8 @@ uv pip install maxtext[runner]==0.2.1 --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"