[build-system]
requires = ["setuptools", "wheel"]
build-backend = "setuptools.build_meta"

[project]
name = "ai2-olmo-core"
dynamic = ["version"]
readme = "README.md"
description = "Core training module for the Open Language Model (OLMo)"
authors = [
    { name = "Allen Institute for Artificial Intelligence", email = "olmo@allenai.org" },
]
requires-python = ">=3.10"
license = { file = "LICENSE" }
dependencies = [
    "numpy",
    "torch>=2.10.0,<2.14",
    "cached-path>=1.7.2",
    "filelock>=3.4,<4.0",
    "requests",
    "packaging",
    "dataclass-extensions>=0.3.0",
    "pyyaml",
    "rich",
    "safetensors",
    "importlib_resources",
    "bettermap",
    "pandas",
]

[project.urls]
Homepage = "https://github.com/allenai/Olmo-core"
Documentation = "https://olmo-core.readthedocs.io/en/latest/"
Changelog = "https://github.com/allenai/Olmo-core/blob/main/CHANGELOG.md"

[project.optional-dependencies]
dev = [
    "black>=23.1,<24.0",
    "boto3",
    "build",
    "furo==2024.8.6",
    "google-cloud-storage",
    "isort>=5.12,<5.14",
    "matplotlib",
    "mypy>=1.0,<1.4",
    "myst-parser>=1.0,<3.0",
    "pytest",
    "pytest-memray",
    "pytest-sphinx",
    "pytest-xdist",
    "ruff==0.15.9",
    "setuptools",
    "Sphinx>=6.0,<=8.1.3",
    "sphinx-autobuild",
    "sphinx-autodoc-typehints==1.23.3",
    "sphinx-copybutton",
    "sphinx-inline-tabs",
    "twine>=1.11.0",
    "wheel",
]
beaker = ["beaker-py>=2.5.4,<3.0", "beaker-gantry>=3.4.3,<4.0", "GitPython>=3.0,<4.0", "google-cloud-compute"]
comet = ["comet_ml"]
# The git source lives in [tool.uv.sources] rather than here: PyPI rejects any uploaded
# metadata containing a direct URL requirement, with an opaque "400 Bad Request".
dion = ["dion ; sys_platform == 'linux'"]
eval = ["ai2-olmo-eval==0.9.0"]
fa4 = [
    "flash-attn-4>=4.0.0b4 ; sys_platform == 'linux'",
    # quack-kernels (pulled in by flash-attn-4's `cute` path) references
    # `cutlass.cute.core.ThrMma`, which nvidia-cutlass-dsl removed in 4.6.0. Pin below that until
    # quack-kernels supports 4.6+, otherwise importing flash_attn.cute raises AttributeError and
    # breaks `import olmo_core.nn`.
    "nvidia-cutlass-dsl<4.6.0 ; sys_platform == 'linux'",
]
# NOTE: fla 0.5.2 requires triton >= 3.7.1 (a hard guard on Hopper). torch 2.13 bundles triton
# 3.7.1 (and nccl 2.29.7 + a matching torchaudio), so it works on the CUDA-12.9 / torch-2.13 image
# (see `beaker-image-cu129` in the Makefile). torch 2.10/2.11 ship triton 3.6.0 and torch 2.12 ships
# 3.7.0, none of which satisfy the guard -- so the torch>=2.13 pin below enforces the floor. Since
# `all` pulls in `fla`, this also makes `.[all]` require torch>=2.13. The <2.14 ceiling matches the
# torch we build/test images with and stops `.[all]` from floating up to an unshipped torch (2.14).
fla = ["flash-linear-attention==0.5.2", "torch>=2.13.0,<2.14"]
# The training image owns the CUDA stack. KDA additionally needs the CUDA 13 CuTe
# compiler (tested with 4.5.3); see KimiDeltaAttentionConfig for installation. Keep
# CUDA-specific extras out of this requirement so Gantry's all-extras resolution
# cannot replace a CUDA 12 PyTorch build with CUDA 13.
kernel-fun = ["kernel-fun==0.2.0 ; sys_platform == 'linux'"]
profiling = ["nvtx"]
torchao = ["torchao==0.15.0"]
transformers = ["transformers>=5.4.0"]
wandb = ["wandb"]
all = ["ai2-olmo-core[dev,beaker,comet,dion,eval,fa4,fla,profiling,torchao,transformers,wandb]"]

[tool.uv.sources]
dion = { git = "https://github.com/microsoft/dion.git", rev = "e64832041d8e01989abf609c9550f6307efbff2a" }

[tool.uv]
environments = [
    "sys_platform == 'linux'",
    "sys_platform == 'darwin' and platform_machine == 'arm64'",
]

[tool.setuptools]
include-package-data = true
license-files = ["LICENSE", "THIRD_PARTY_NOTICES.md"]

[tool.setuptools.package-data]
olmo_core = [
    "py.typed",
    "*.txt",
    "nn/_csrc/*.cpp",
    "kernels/cuda/*.cpp",
    "kernels/cuda/*.cu",
    "kernels/cuda/*.cuh",
    "kernels/cuda/*.h",
    "kernels/cuda/CMakeLists.txt",
]

[tool.setuptools.dynamic]
version = { attr = "olmo_core.version.VERSION" }

[tool.setuptools.packages.find]
where = ["src"]
include = ["olmo_core*"]
exclude = []

[tool.black]
line-length = 100
include = '\.pyi?$'
exclude = '''
(
      __pycache__
    | \.git
    | \.mypy_cache
    | \.pytest_cache
    | \.vscode
    | \.venv
    | \bdist\b
    | \bdoc\b
    | scratch/
)
'''

[tool.isort]
profile = "black"
multi_line_output = 3

[tool.ruff]
line-length = 100

[tool.ruff.lint]
ignore = ["F403", "F405", "E501"]
exclude = [
    ".bzr",
    ".direnv",
    ".eggs",
    ".git",
    ".venv",
    "venv",
    ".mypy_cache",
    "__pycache__",
    ".nox",
    ".pants.d",
    ".pytype",
    ".ruff_cache",
    ".svn",
    ".tox",
    "__pypackages__",
    "_build",
    "buck-out",
    "build",
    "dist",
    "node_modules",
    "doc",
    "pretrain_data",
    "inference",
]

[tool.ruff.lint.per-file-ignores]
"**/__init__.py" = ["F401"]

[tool.pyright]
reportPrivateImportUsage = false

[tool.ty.rules]
possibly-missing-attribute = "ignore"
possibly-missing-import = "ignore"
unresolved-attribute = "warn"

[[tool.ty.overrides]]
include = ["src/olmo_core/nn/**"]

[tool.ty.overrides.rules]
unresolved-attribute = "ignore"

[tool.mypy]
ignore_missing_imports = true
no_site_packages = true
check_untyped_defs = true
namespace_packages = false
disable_error_code = "has-type"

[[tool.mypy.overrides]]
module = "tests.*"
strict_optional = false

[tool.pytest.ini_options]
testpaths = "src/test/"
python_classes = ["Test*", "*Test"]
log_format = "%(asctime)s - %(levelname)s - %(name)s - %(message)s"
log_level = "DEBUG"
log_cli = false
log_cli_level = "DEBUG"
markers = ["gpu"]
filterwarnings = [
    'ignore::FutureWarning:huggingface_hub\.file_download',
    'ignore::DeprecationWarning:pkg_resources',
    'ignore::DeprecationWarning:google\.rpc',
    'ignore::DeprecationWarning:torchao\.dtypes.*',
    'ignore::FutureWarning:torch\.distributed\.checkpoint\.default_planner',
    'ignore::UserWarning:torch\.distributed\.checkpoint\.state_dict_saver',
    'ignore::UserWarning:torch\.distributed\.checkpoint\.state_dict_loader',
]