# ---- All project specifications ---- #
[project]
name = "torchtitan"
description = "A PyTorch native platform for training generative AI models"
readme = "README.md"
requires-python = ">=3.11"
license = {file = "LICENSE"}
authors = [
    { name = "PyTorch Team", email = "packages@pytorch.org" },
]
keywords = ["pytorch", "training", "llm"]
dependencies = [
    # Stateful Dataloader
    "grain==0.2.18",

    # Checkpointing
    "torch_checkpointing>=0.1.0",

    # Hugging Face integrations
    "datasets>=3.6.0,<4.8.0",
    "tokenizers>=0.15.0",
    "renderers==0.1.11",
    "safetensors",

    # Miscellaneous
    "tensorboard",
    "wandb",
    "einops",
    "pillow",
    "spmd_types @ git+https://github.com/meta-pytorch/spmd_types.git@8b8a56e79b117cfac3ed289fa209f8358558d482",
    "attn-gym[linear]==0.0.16",
    # Keep this commit pin in sync with requirements.txt.
    "torch_remat @ git+https://github.com/meta-pytorch/remat.git@4fd000acba6dc66762b7f3217ecadb8e07243dd4",
]
dynamic = ["version"]

[project.urls]
GitHub = "https://github.com/pytorch/torchtitan"
Documentation = "https://github.com/pytorch/torchtitan/tree/main/docs"
Issues = "https://github.com/pytorch/torchtitan/issues"

[project.optional-dependencies]
dev = [
    "pre-commit",
    "pytest",
    "pytest-cov",
    "expecttest", # test_tokenizer
    "pyrefly==0.45.1",
]

[tool.setuptools.dynamic]
version = {file = "assets/version.txt"}

[tool.codespell]
skip = ["*.json"]
ignore-words-list = ["assertin", "datas" ,"indx", "inpt", "nd", "socio-economic"]

# ---- Explicit project build information ---- #
[build-system]
requires = ["setuptools>=61.0"]
build-backend = "setuptools.build_meta"

[tool.setuptools.packages.find]
where = [""]
# torchtitan_recipes is a second top-level package, listed so that
# tightening the first pattern to "torchtitan.*" cannot silently drop it.
include = ["torchtitan*", "torchtitan_recipes*"]

[tool.pytest.ini_options]
addopts = ["--showlocals"]  # show local variables in tracebacks
testpaths = ["tests"]
markers = [
    "multi_gpu: tests that require multiple physical GPUs",
]

[tool.pyrefly]
python-version = "3.11"
project-excludes = ["torchtitan/experiments", "torchtitan/rl", "**/tests/**"]
replace-imports-with-any = [
    "torchao.*", "torchft", "torchvision.*", "deep_ep.*", "jinja2.*", "helion", "helion.*", "batch_invariant_ops", "verifiers.*", # optional dependencies
    "dist_moe", "dist_moe.*", # optional dependency without Pyrefly stubs
    # torch_remat is required, but does not publish Pyrefly stubs.
    "torch_remat",
]
search-path = ["../pytorch"]  # local built pytorch
