pyproject.toml

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

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

# Production dependencies.
# Every time we upgrade JAX, we should try to bring the rest to the newest versions.
dependencies = [
    "attrs>=21.3.0",
    "absl-py",
    "chex<0.1.81",  # chex 0.1.81 depends on numpy>=1.25.0.
    "flax==0.7.4",  # only for checkpoints.
    "importlab==0.7",  # breaks pytype on 0.8
    "jax==0.4.21",
    "jaxlib==0.4.21",
    "nltk==3.7",  # for text preprocessing
    "numpy<1.24",  # needed to pin to < 1.24; tf ragged_tensor depends on deprecated np.object.
    "optax==0.1.7",  # optimizers (0.1.0 has known bugs).
    "portpicker",
    "protobuf>=3.20.3",
    "tensorboard-plugin-profile==2.14.0",
    # This has both x86 and arm64 wheels. Underneath the hood it uses tensorflow-macos since 2.13.
    "tensorflow==2.14.0",
    "tensorflow-datasets>=4.9.2",
    # tensorflow-io 0.34.0 enabled macOS arm64 builds.
    "tensorflow-io>=0.34.0",
    "tensorflow_text==2.14.0; platform_machine == 'x86_64'", # implied by seqio, but also used directly for text processing
    "tensorstore>=0.1.21",  # used for supporting GDA checkpoints
    "toml",  # for config management
    "typing-extensions>=4.2.0",  # needed for typing.Protocol
    "seqio==0.0.18",  # used for inputs
]

[project.optional-dependencies]
# 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 = [
    "black==23.1a1",  # formatting
    "diffusers==0.16.1", # diffuser dependency for golden config updates
    "einops",
    "evaluate",
    "fairseq==0.12.2",  # test-only; note Apple Silicon start.md instructions parse this specifier '==', please be careful when upgrading.
    "isort",  # formatting
    "pre-commit",  # local pre commit hooks
    "pycocotools",  # COCO evaluation tools
    "pylint==2.13.9",  # formatting.  Upgrading to 2.14.0 shows (likely) pylint bug.
    "pytest",  # test runner
    "pytest-xdist",  # pytest plugin for test parallelism
    "pytype==2022.4.22",  # type checking
    "scikit-learn>=1.1.1",  # test-only
    "scipy",
    "sentencepiece != 0.1.92",
    "tqdm",  # test-only
    "timm==0.6.12",  # DiT Dependency test-only
    "torch>=1.12.1",  # test-only
    "torchvision==0.14.1",  # test-only
    "transformers==4.36.0",  # 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.
    "google-api-python-client==2.109.0",
    "google-auth==2.23.0",
    "google-auth[pyopenssl]",  # Ensures that we have compatible pyopenssl/cryptography pins.
    "pyOpenSSL>=22.1.0",  # compat with cryptography version.
    "google-cloud-storage==2.11.0",
    "google-cloud-core==2.3.3",
]
# For TPU training.
# Note: Specify -f https://storage.googleapis.com/jax-releases/libtpu_releases.html during install.
tpu = [
    "axlearn[gcp]",
    "jax[tpu]==0.4.21",  # 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",
    "google-cloud-aiplatform[tensorboard]",
    "tensorboard",
]
# Dataflow dependencies.
dataflow = [
    "apache-beam==2.52.0",
    "apache-beam[gcp]",
    "google-apitools",  # for beam pipeline
    "orjson==3.9.10",
]

[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.pylint.basic]
# A comma-separated list of package or module names from where C extensions may
# be loaded. Extensions are loading into the active Python interpreter and may
# run arbitrary code.
extension-pkg-whitelist = "numpy,jax"

[tool.pylint.messages_control]
# Disable the message, report, category or checker with the given id(s). You
# can either give multiple identifiers separated by comma (,) or put this
# option multiple times (only on the command line, not in the configuration
# file where it should appear only once).You can also use "--disable=all" to
# disable everything first and then reenable specific checks. For example, if
# you want to run only the similarities checker, you can use "--disable=all
# --enable=similarities". If you want to run only the classes checker, but have
# no Warning level messages displayed, use"--disable=all --enable=classes
# --disable=W"
disable = [
    "blacklisted-name",
    "duplicate-code",
    "fixme",
    "invalid-name",
    "missing-docstring",
    "no-else-return",
    "no-member",  # Config objects use dynamic members generated by attrs.
    "protected-access",
    "protected-name",
    "redefined-builtin",
    "redefined-outer-name",
    "too-few-public-methods",
    "too-many-arguments",
    "too-many-locals",
    "unnecessary-lambda",
    "unsubscriptable-object",  # pylint considers jax.Array unsubscriptable.
]

[tool.pylint.typecheck]
# List of members which are set dynamically and missed by Pylint inference
# system, and so shouldn't trigger E1101 when accessed.
generated-members = "googleapiclient.*,numpy.*,torch.*,sentencepiece.*,jaxlib.xla_extension.*"

[tool.pylint.format]
indent-string = "    "
max-line-length = 100
fail-under = 10
# Disable ignore-long-lines for URLs and gs paths.
ignore-long-lines = "<?(http|https|gs)://\\S+>?"

[tool.pylint.similarities]
# Don't include imports or docstrings when looking for duplicate-code.
ignore-imports = "yes"
ignore-docstrings = "yes"
# By default, this is 4, which causes many false-positives.
min-similarity-lines = 20

[tool.pylint.MAIN]
load-plugins = [
    "pylint.extensions.docparams"
]

[tool.black]
line-length = 100
target-version = ['py38', '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"