Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ tmp/
logs/
.venvs
venv*
maxtext_venv/
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
Expand Down
23 changes: 19 additions & 4 deletions build_hooks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
2 changes: 2 additions & 0 deletions docs/install_maxtext.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
7 changes: 3 additions & 4 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
Loading