pyproject.toml

[build-system]
requires = ["hatchling", "hatch-vcs"]
build-backend = "hatchling.build"

[project]
name = "axlearn"
dynamic = ["version"]
description = "AXLearn"
readme = "README.md"
requires-python = "==3.12.*"

# Production dependencies.
# Minimal requirments for axlearn/common/config.py.
dependencies = [
    "attrs>=23.1.0", # We use `type` in `attrs.field`
    "numpy>=2.0,<3",  # verified with tensorflow 2.14 RaggedTensor
]

[project.optional-dependencies]
# Core dependencies needed for jax based training and evaluation.
# Every time we upgrade JAX, we should try to bring the rest to the newest versions.
core = [
    "absl-py==2.3.1",
    "chex==0.1.91",
    "etils[epy]==1.13.0",  # pin to ensure importlib-resources is installed (needed by tensorflow-datasets)
    "importlab==0.8.1",  # breaks pytype on 0.8
    "jax==0.8.3",
    "jaxlib==0.8.3",
    "ml-dtypes==0.5.3",
    "msgpack==1.1.0",  # for checkpointing.
    "nltk==3.7",  # for text preprocessing
    "optax==0.2.6",  # Fix for TPU 7x
    "portpicker",
    "pyarrow<21.0.0",  # Pin to v20.x to avoid PyExtensionType -> ExtensionType breaking change in v21
    "protobuf>4",
    "tensorboard-plugin-profile==2.20.4",
    # This has both x86 and arm64 wheels. Underneath the hood it uses tensorflow-macos since 2.13.
    "tensorflow==2.19.1.3",
    "tensorflow-datasets==4.9.9",
    "tensorflow-io==0.37.3",  # for tensorflow-2.16. Note that 0.37.0 results in "pure virtual method called".
    "tensorflow-metadata==1.17.3", # Required by seqio/tfds
    "tensorflow_text==2.19.0", # implied by seqio, but also used directly for text processing
    "tensorstore>=0.1.63",  # used for supporting GDA checkpoints
    "toml",  # for config management
    "typing-extensions==4.15.0",
    "scipy>=1.12.0",  # to avoid "module 'scipy.linalg' has no attribute 'tril'"
    "seqio==0.0.20",  # used for inputs
    "aqtp==0.9.0", # Updated from 0.4.0; compatible with Python 3.10
    "flax>=0.12.0", # for AQT, param converter and adapter.
    "prefixed==0.9.0", # For formatting file sizes, param counts, etc.
    "grain==0.2.11", # Grain input processing. 0.2.6 onwards supports macos.
    "tokamax",
]
# Apple Silicon dependencies.
# Need to run `conda install -c apple tensorflow-deps` separately.
apple-silicon = [
    "array-record<0.8; platform_machine == 'arm64'", # Prevent mutex.cc lock
]
# Requirements for testing and development.
dev = [
    "einops==0.8.0",  # for unittests only. Do not use it in core impl because it's not thread-safe.
    "evaluate",
    "pika==1.3.2",  # used by event queue
    "pre-commit",  # local pre commit hooks
    "pycocotools",  # COCO evaluation tools
    "pylint==4.0.3",
    "pytest",  # test runner
    "pytest-xdist",  # pytest plugin for test parallelism
    "pytest-timeout",  # pytest plugin for forcing timeout of tests
    "pytype==2024.10.11",  # type checking
    "scikit-learn==1.6.1",  # test-only
    # Fix AttributeError: module 'scipy.linalg' has no attribute 'tril' and related scipy import errors.
    "scipy>=1.12.0",
    "sentencepiece != 0.1.92",
    "tqdm",  # test-only
    "timm==1.0.17",  # DiT Dependency test-only
    "torch==2.7.1",  # test-only
    "torchvision==0.22.1",  # test-only
    "safetensors<=0.5.3",  # TODO: Remove once torch dependency is >=2.3.0
    "transformers==4.51.3",  # test-only
    "wandb",  # test-only
    "wrapt",  # implied by tensorflow-datasets, but also used in config tests.
]
# Requirements for GCP job management.
gcp = [
    "cloud-tpu-client",  # for creating and managing TPUs.
    "crcmod",  # for file uploads.
    # kubernetes==32.0.0 will cause "kube_config.py:520] Object of type ConfigNode is not JSON serializable" error
    # when loading kube config, hence failing kubernetes client authentication.
    # See also https://github.com/kubernetes-client/python/issues/2333
    "kubernetes==31.0.0",
    "google-api-core==2.24.1",
    "google-api-python-client==2.109.0",
    "google-auth==2.29.0",
    "google-auth[pyopenssl]",  # Ensures that we have compatible pyopenssl/cryptography pins.
    "google-cloud-storage==2.19.0",
    "google-cloud-compute==1.19.2", # Needed for region discovery for CloudBuild API access.
    "google-cloud-core==2.3.3",
    "google-cloud-build==3.24.1",
    "ml-goodput-measurement==0.0.15",
    "pika==1.3.2",  # used by event queue
    "pyOpenSSL>=22.1.0",  # compat with cryptography version.
    "tpu-info==0.2.0", # For TPU monitoring from libtpu. https://github.com/AI-Hypercomputer/cloud-accelerator-diagnostics/tree/main/tpu_info
    "prometheus-client==0.21.0", # For TPU monitoring from tpu-device-plugin.
]
# For TPU training.
# Note: Specify -f https://storage.googleapis.com/jax-releases/libtpu_releases.html during install.
tpu = [
    "axlearn[gcp]",
    "jax[tpu]==0.8.3",  # must be >=0.4.19 for compat with v5p.
    "pathwaysutils==0.1.2",  # For JAX+Pathways single-controller accelerator coordinator.
]
# Vertex AI tensorboard. TODO(markblee): Merge with `gcp`.
vertexai_tensorboard = [
    # Required to fix a `distro-info` bug we run into when using `tb_gcp_tensorboard` from
    # `google_cloud_aiplatform[tensorboard]`. See:
    # [1] https://bugs.launchpad.net/ubuntu/+source/distro-info/+bug/1991606
    # [2] https://github.com/pypa/setuptools/issues/3772
    "setuptools==65.7.0",
    # Pin version to fix Tensorboard uploader TypeError: can only concatenate str (not "NoneType") to str
    # https://github.com/googleapis/python-aiplatform/commit/4f982ab254b05fe44a9d2ed959fca2793961b56c
    "google-cloud-aiplatform[tensorboard]>=1.61.0",
    "tensorboard",
]
# Dataflow dependencies.
dataflow = [
    "pyarrow<21.0.0",  # Pin to v20.x to avoid PyExtensionType -> ExtensionType breaking change in v21
    "apache-beam==2.69.0",
    "apache-beam[gcp]",
    "google-apitools",  # for beam pipeline
    "orjson==3.9.10",
]
# GPU custom kernel dependency.
gpu = [
    "triton>=2.1.0",
    "jax[cuda12]==0.8.3",
    "nvidia-ml-py==12.560.30",
    # pin nccl version, otherwise jax[cuda12] will pull latest version
    "nvidia-nccl-cu12==2.27.5",
    "nvidia-cudnn-cu12>=9.8.0.87" # Pin CuDNN to at least 9.8 for Jax >= 0.6.2
]
# Open API inference.
open_api = [
    "openai==1.35.1",
]
# MMAU benchmark.
mmau = [
    "openai==1.35.14",
    "anthropic==0.31.2",
    "nbformat==5.10.4",
    "plotly==5.22.0",
    "nbconvert==7.16.4",
    "google-cloud-aiplatform==1.50.0",
    "google-genai==1.25.0",
    "transformers==4.51.3",
    "huggingface-hub==0.24.6",
]
# Orbax checkpointing.
orbax = [
    "humanize==4.10.0",
    "orbax-checkpoint==0.11.33",
]
# Audio dependencies.
audio = [
    "levenshtein==0.25.1",
]

[tool.hatch]
# Read the version from git tags
version.source = "vcs"
version.fallback-version = "0.0.0"
build.targets.wheel.exclude = ["**/BUILD", "axlearn/experiments/testdata"]
# This is where the code is for historical reasons.
build.targets.wheel.packages = ["axlearn"]

[tool.pytype]
inputs = ["axlearn"]

[project.scripts]
# Entrypoint for CLI.
axlearn = "axlearn.cli:main"

[tool.black]
line-length = 100
target-version = ['py310', 'py312']

[tool.ruff]
line-length = 100
target-version = "py312"

[tool.pytest.ini_options]
timeout = 300
addopts = "-rs -s -p no:warnings --junitxml=test-results/testing.xml"
markers = [
    "gs_login: tests needing GS login.",
    "tpu: tests needing access to a TPU device.",
    "high_cpu: tests that require a lot of CPU.",
    "fp64: tests that require 64-bit floating point precision.",
    "for_8_devices: tests that run on host platform device count of 8.",
    "golden_config: golden config tests.",
    "golden_init: golden init tests.",
    "golden_regularizer: golden regularizer scale tests.",
    "golden_param_update: golden param update scale tests.",
]
testpaths = "."
junit_family="xunit2"

[tool.isort]
line_length = 100
profile = "black"

[tool.uv.pip]
find-links = [
    "https://storage.googleapis.com/axlearn-wheels/wheels.html",
    "https://storage.googleapis.com/jax-releases/libtpu_releases.html",
]
no-build-isolation-package = [
    "apache-beam",
]