pyproject.toml

[project]
name = "esm"
version = "3.4.1.post1"
description = "EvolutionaryScale open model repository"
readme = "README.md"
requires-python = ">=3.12"
license = {file = "LICENSE.md"}

authors = [
  {name = "EvolutionaryScale Team"}
]

maintainers = [
  {name = "Zeming Lin", email = "zeming+esm@evolutionaryscale.ai" }
]

classifiers = [
  "Development Status :: 3 - Alpha",
  "Topic :: Scientific/Engineering :: Bio-Informatics",
  "Programming Language :: Python :: 3.12",
  "Programming Language :: Python :: 3.13",
  "Programming Language :: Python :: 3.14",
]

dependencies = [
  "torch>=2.11.0,<2.12.0",
  "transformers>=4.57.6,<5.0.0",
  "ipython",
  "einops",
  "biotite>=1.0.0",
  "rdkit",
  "msgpack-numpy",
  "biopython",
  "scikit-learn",
  "brotli",
  "attrs",
  "pandas",
  "cloudpathlib",
  "httpx",
  "tenacity",
  "zstd",
  "ipywidgets",
  "py3dmol",
  "pydssp",
  "boto3",
  # Both are imported directly by esm.models.hub, esm.models.esmc and
  # esm.models.esmfold2; they resolved only because transformers happens to pull
  # them in.
  "huggingface_hub",
  "safetensors",
  "pygtrie",
  "dna_features_viewer",
  "accelerate",
  "cuequivariance-torch>=0.8.1; sys_platform == 'linux' and platform_machine == 'x86_64'",
  "cuequivariance-ops-torch-cu13>=0.8.1; sys_platform == 'linux' and platform_machine == 'x86_64'",
]

[project.optional-dependencies]
# ESMFold2 "fused" kernel backend (set_kernel_backend("fused"))
# inference kernels for tri-mul / LN+SwiGLU / dropout-residual.
fused = ["triton>=3,<4"]
# ESMFold2 context-parallel ESM-C tensor parallelism (wrap_model_with_cp(tp_esmc=True)).
fold-cp = ["transformer-engine[pytorch]>=2,<3"]
# "cuequivariance" kernel backend — pick the build matching your CUDA toolkit.
cueq12 = ["cuequivariance-torch", "cuequivariance-ops-torch-cu12"]
cueq13 = ["cuequivariance-torch", "cuequivariance-ops-torch-cu13"]

# Pytest
[tool.pytest.ini_options]
addopts = """
  --cov=esm
	--cov-report term-missing:skip-covered
	-n auto
	-m 'not gpu and not manual and not merge_only'
	--ignore=tests/oss_pytests/test_oss_client.py
"""
markers = [
  "gpu: requires a CUDA device; run in the monorepo's gpu-ci",
  "manual: too heavy for a hosted runner; run in the monorepo's research-ci",
  "merge_only: needs a real published checkpoint; run in the monorepo's research-ci",
]

[tool.setuptools]
package-dir = {"" = "."}
include-package-data = true

[tool.setuptools.packages.find]
where = ["."]
include = ["esm*"]

[tool.setuptools.package-data]
esm = ["data/*"]

[tool.pixi.workspace]
channels = ["conda-forge"]
platforms = ["linux-64", "osx-arm64"]

# These are build dependencies, to ensure pip support, keep run-time deps above in `dependencies`
[tool.pixi.target.linux-64.activation.env]
LD_PRELOAD = "$CONDA_PREFIX/lib/libstdc++.so.6${LD_PRELOAD:+:$LD_PRELOAD}"

[tool.pixi.dependencies]
matplotlib-base = "*"
filelock = "*"
pandas = "*"
colorama = "*"
setuptools = "<82"
scipy = "*"
pyyaml = "*"
pkg-config = "*"
cmake = "*"
pip = "*"
twine = "*"

# Python is per-feature rather than a workspace dependency so one manifest can
# solve every version CI exercises.
[tool.pixi.feature.py312.dependencies]
python = "3.12.*"

[tool.pixi.feature.py313.dependencies]
python = "3.13.*"

[tool.pixi.feature.py314.dependencies]
python = "3.14.*"

[tool.pixi.pypi-dependencies]
esm = { path = ".", editable = true }

[tool.pixi.tasks]
build-wheel = "python -m pip wheel --no-deps -w dist ."
upload-wheel = "python -m twine upload --repository pypi"

[tool.pixi.feature.dev.dependencies]
matplotlib = "*"
numpy = ">=2.0.0,<2.4"
pre-commit = "*"
pytest = "*"
pytest-cov = "*"
pytest-xdist = "*"
seaborn = "*"
ty = "==0.0.49"

[tool.pixi.feature.dev.pypi-dependencies]
# DockQ is optional at runtime. Keep the NumPy 2-compatible fork in the dev
# environment for tests without publishing its direct URL in package metadata.
dockq = { git = "https://github.com/nrontsis/DockQ.git", rev = "ba4df5adaad7c77fd60851d0b7b05f2b77061ba2" }

[tool.pixi.feature.dev.tasks]
lint-all = "pre-commit run --all-files --show-diff-on-failure"
cov-test = "pytest -v --junitxml=pytest.xml --cov=esm"

[tool.pixi.feature.fused]
platforms = ["linux-64"]

[tool.pixi.feature.fused.pypi-dependencies]
flash_attn = { url = "https://github.com/evolutionaryscale/wheels/releases/download/py312-pt211-cu13-sm80-90/flash_attn-2.7.4.post1-cp312-cp312-linux_x86_64.whl" }
xformers = { url = "https://download.pytorch.org/whl/cu130/xformers-0.0.35-py39-none-manylinux_2_28_x86_64.whl" }
cuequivariance-torch = ">=0.8.1"
cuequivariance-ops-torch-cu13 = ">=0.8.1"

# `gpu` stays on 3.12: the prebuilt flash_attn wheel in `fused` is cp312-only.
[tool.pixi.environments]
default = {features = ["py312"], solve-group = "default"}
dev = {features = ["dev", "py312"], solve-group = "default"}
gpu = {features = ["dev", "fused", "py312"], solve-group = "default"}
dev-py313 = {features = ["dev", "py313"], solve-group = "py313"}
dev-py314 = {features = ["dev", "py314"], solve-group = "py314"}

[tool.ruff]
extend-include = ["*.ipynb"]

[tool.ruff.lint]
# Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`), sort imports ('I')
select = ["E4", "E7", "E9", "F", "I"]
ignore = [
    # allow variable == False (tensors should do this)
    "E712",
    # allow assigning of lambdas
    "E731",
    # Allow ambiguous variables, e.g. we use O for oxygen
    "E741",
    # Ignore errors from jaxtyping hints
    # https://docs.kidger.site/jaxtyping/faq/#flake8-or-ruff-are-throwing-an-error
    "F722",
    # TODO: Fix the few offenders in a follow up PR
    "E721",
]

[tool.ruff.lint.isort]
split-on-trailing-comma = false
known-third-party = ["wandb"]

[tool.ruff.format]
skip-magic-trailing-comma = true
docstring-code-format = true
docstring-code-line-length = "dynamic"

[tool.isort]
known_third_party = ["wandb"]

[tool.ty.src]
# gitignore-style globs. Skip notebooks, tests, and the Modal-app tutorial whose deps
# (modal, abnumber) and dynamic decorators aren't resolvable in the lint env.
exclude = [
    "**/*.ipynb",
    "**/*_test.py",
    "**/test_*.py",
    "**/tests/",
    "**/conftest.py",
    "cookbook/tutorials/binder_design.py",
]

[tool.ty.rules]
# Strict LSP override-compatibility checking that pyright `basic` never enforced; low
# value for duck-typed ML code.
invalid-method-override = "ignore"
# flash_attn resolves only where it's installed (a GPU environment); the ignores it
# requires there read as unused/redundant in this CPU lint env where it's absent.
unused-ignore-comment = "ignore"
unused-type-ignore-comment = "ignore"

[tool.ty.analysis]
# ty can't introspect optional/compiled deps absent from this lint env (flash_attn is
# GPU-only; zstd is a C-extension), so it falsely reports them unresolved. Treat as Any
# rather than disabling unresolved-import, so genuinely-broken first-party imports still
# surface.
replace-imports-with-any = [
    "biotite.database",
    "biotite.database.**",
    "biotite.structure",
    "biotite.structure.**",
    "biotite.sequence.align",
    "biotite.sequence.align.**",
    "scipy.spatial",
    "transformer_engine",
    "transformer_engine.**",
    "transformer_engine_torch",
    "xformers",
    "xformers.**",
    "cuequivariance_torch",
    "cuequivariance_torch.**",
    "flash_attn",
    "flash_attn.**",
    "zstd",
]

[tool.importlinter]
root_package = "esm"