pyproject.toml

[build-system]
requires = ["flit_core >=3.2,<4"]
build-backend = "flit_core.buildapi"

[project]
name = "axlearn"
version = "0.1.1"
description = "AXLearn"
readme = "README.md"
requires-python = ">=3.9"

# Production dependencies.
# Minimal requirments for axlearn/common/config.py.
dependencies = [
    "attrs>=23.1.0", # We use `type` in `attrs.field`
    "numpy==1.26.4",  # 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.1.0",
    "chex==0.1.86",  # chex 0.1.86 is required for jax 0.4.25.
    "importlab==0.7",  # breaks pytype on 0.8
    "jax==0.4.30",
    "jaxlib==0.4.30",
    "nltk==3.7",  # for text preprocessing
    "optax==0.1.7",  # optimizers (0.1.0 has known bugs).
    "portpicker",
    "protobuf>=3.20.3",
    "tensorboard-plugin-profile==2.15.1",
    # This has both x86 and arm64 wheels. Underneath the hood it uses tensorflow-macos since 2.13.
    "tensorflow==2.16.1",
    "tensorflow-datasets>=4.9.2",
    "tensorflow-io>=0.37.1",  # for tensorflow-2.16. Note that 0.37.0 results in "pure virtual method called".
    "tensorflow_text==2.16.1; platform_machine == 'x86_64'", # 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.11.0",
    "scipy==1.12.0",  # to avoid "module 'scipy.linalg' has no attribute 'tril'"
    "seqio==0.0.18",  # used for inputs
    "aqtp==0.4.0", # aqtp>=0.5.0 requires python>=3.10 due to the use of Union shorthand "|" in type hints.
    "flax==0.8.5", # for AQT, param converter and adapter.
    "prefixed==0.7.1", # For formatting file sizes, param counts, etc.
]

# Apple Silicon dependencies.
# Need to run `conda install -c apple tensorflow-deps` separately.
apple-silicon = [
    "tensorflow-metal==1.1.0; platform_machine == 'arm64'",
    # As of Nov 15, 2023, tensorflow_text has no arm64 build.
    # We need to manually build the wheel and install it, following instructions at https://github.com/tensorflow/text/tree/2.14#build-from-source-steps.
]

# Requirements for testing and development.
dev = [
    "axlearn[orbax]",  # checkpointer tests
    "black==23.1a1",  # formatting
    "einops",
    "evaluate",
    "isort",  # formatting
    "pre-commit",  # local pre commit hooks
    "pycocotools",  # COCO evaluation tools
    "pylint==2.17.7",
    "pytest",  # test runner
    "pytest-xdist",  # pytest plugin for test parallelism
    "pytype==2022.4.22",  # type checking
    "scikit-learn>=1.1.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==0.6.12",  # DiT Dependency test-only
    "torch>=1.12.1",  # test-only
    "torchvision==0.16.1",  # test-only
    "transformers==4.44.1",  # 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",
    "google-api-core==2.17.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.16.0",
    "google-cloud-core==2.3.3",
    "google-cloud-build==3.24.1",
    "ml_goodput_measurement==0.0.2",
    "pyOpenSSL>=22.1.0",  # compat with cryptography version.
]
# For TPU training.
# Note: Specify -f https://storage.googleapis.com/jax-releases/libtpu_releases.html during install.
tpu = [
    "axlearn[gcp]",
    "jax[tpu]==0.4.30",  # must be >=0.4.19 for compat with v5p.
]
# Vertex AI tensorboard.
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 = [
    "apache-beam==2.55.1",
    "apache-beam[gcp]",
    "google-apitools",  # for beam pipeline
    "orjson==3.9.10",
]
# GPU custom kernel dependency.
gpu = [
    "triton==2.1.0",
    "jax[cuda12_pip]==0.4.30",
]
# 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",
    "transformers==4.44.1",
]
# Orbax checkpointing.
orbax = [
    "humanize==4.10.0",
    "orbax-checkpoint==0.5.23",
]

[tool.flit.module]
# This defines the import name. https://flit.pypa.io/en/stable/pyproject_toml.html#module-section
name = "axlearn"

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

[tool.black]
line-length = 100
target-version = ['py38', 'py39']

[tool.ruff]
line-length = 100
target-version = 'py39'

[tool.pytest.ini_options]
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.",
    "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"