pyproject.toml

[build-system]
requires = ["uv_build>=0.7.19,<0.10.0"]
build-backend = "uv_build"

# Explicit module-name override because the dist name (marin-levanter) no
# longer matches the import name (levanter); uv_build's default would look
# for src/marin_levanter/ which doesn't exist.
[tool.uv.build-backend]
module-name = "levanter"

[project]
name = "marin-levanter"
version = "0.2.0"
description = "Scalable Training for Foundation Models with Named Tensors and JAX"
readme = "README.md"
requires-python = ">=3.11"

authors = [
    { name = "David Hall", email = "dlwh@cs.stanford.edu" },
    { name = "Jason Wang", email = "jsywang@cs.stanford.edu" },
    { name = "Ahmed Ahmed" },
    { name = "Ivan Zhou", email = "ivanz@stanford.edu" },
    { name = "Will Held" },
    { name = "Virginia Adams" },
]

classifiers = [
    "Programming Language :: Python :: 3",
    "License :: OSI Approved :: Apache Software License",
    "Operating System :: POSIX :: Linux",
    "Operating System :: MacOS :: MacOS X",
    "Development Status :: 4 - Beta",
    "Intended Audience :: Science/Research",
]

dependencies = [
    "marin-haliax",
    "equinox>=0.11.7,!=0.12.0",
    "jax>=0.9.2,<0.11",
    "marin-fray",
    "marin-rigging",
    "marin-zephyr",
    "einops",
    "jaxtyping>=0.2.34",
    "tokenizers>=0.15.2",
    "transformers>=4.57.1,<5.0",
    "chex>=0.1.86",
    "optax>=0.1.9,<0.2.7",
    "wandb>0.24.0",
    "draccus>=0.11.5",
    "pyarrow>=11.0.0",
    "zstandard>=0.18.0",
    "datasets>=3.1.0,<5.0",
    "gcsfs>=2024.2,<2027",
    "braceexpand>=0.1.7",
    "chex>=0.1.86",
    "jmp>=0.0.3",
    "fsspec[http]>=2024.2,<2027",
    "tensorstore>=0.1.73,<0.1.82",
    "pytimeparse>=1.1.8",
    "humanfriendly==10.0",
    "safetensors[numpy]>=0.4.2,<0.7.0",
    "tblib>=1.7.0,<4.0.0",
    "dataclasses-json~=0.6.4",
    "pydantic<3",
    "filelock~=3.13",
    "async-lru~=2.0",
    "tqdm-loggable>=0.2",
    "deepdiff>=8.6.2",
    "lenses",
    "jinja2",  # needed for apply_chat_template
    "protobuf>=6",
    "immutabledict",
]

[project.urls]
Homepage = "https://github.com/stanford-crfm/levanter"
"Bug Tracker" = "https://github.com/stanford-crfm/levanter/issues"

[project.optional-dependencies]
gpu = [
  "jax[cuda13]==0.10.0",
  # B200 emits a cuBLAS warning with older CUDA 13 cuBLAS builds.
  "nvidia-cublas>=13.2.0.9; sys_platform == 'linux'",
  # Preserve the CoreWeave H100 all-to-all guard under CUDA 13.
  "nvidia-nccl-cu13>=2.28.3; sys_platform == 'linux'",
  # Optional raw Sonic MoE gather/combine backend. Keep jax-triton exact:
  # 0.3.0 is missing the CUDA backend API needed by our current stack.
  "jax-triton==0.3.1; sys_platform == 'linux'",
  # This module imports Triton directly; stay on the tested minor without
  # forcing an exact patch version.
  "triton>=3.6.0,<3.7; sys_platform == 'linux'",
]
deepep = [
  # DeepEP itself remains an external source checkout addressed by DEEPEP_SRC_ROOT.
  # The Levanter wheel ships the JAX FFI shim sources; CUDA toolchain dependencies
  # come from the host GPU environment.
]
tpu = ["jax==0.9.2", "jaxlib==0.9.2", "libtpu==0.0.39"]
torch_test = [
  "torch>=2.7.0",
  "peft>=0.12.0",
]
profiling = [
    "xprof",
  # TensorBoard UI for viewing traces and profiles
  "tensorboard>=2.16",
  # Writer used by the TensorBoard tracker implementation
  "tensorboardX>=2.6",
]
serve = [
  "fastapi>=0.100.0",
  "uvicorn[standard]>=0.23.0",
  "openai>=1.0.0",
]

lm_eval = [
  "lm-eval@git+https://github.com/stanford-crfm/lm-evaluation-harness.git@d5e3391f22cde186c827674d5c3ec7c5f4fe0cab",
]

kernels = [
    "tokamax",
]

gcp = [
    "google-api-python-client>=2.175.0",
    "google-cloud-storage",
    "google-cloud-storage-transfer",
    "google-auth",
]

[dependency-groups]
test = [
  "pytest",
  "pytest-forked",
  "pytest-asyncio",
  "soundfile",
  "librosa>=0.11.0",
  "torchcodec>=0.10,<0.11",
  "tensorboardX>=2.6",
  "chex>=0.1.86",
  "httpx>=0.28.1",
  "pytest-profiling>=1.8.1",
  "pytest-xdist",
  "trackio>=0.5.0",
  # need this or else the resolver gets a bad version
  "numba>=0.62.1"
]

docs = [
  "mkdocs>=1.4.3",
  "mkdocs-material>=7.3.3",
  "mkdocstrings>=0.22.0",
  "mkdocs-literate-nav>=0.6.0",
  "mkdocs-macros-plugin>=0.7.0",
  "mkdocstrings-python>=1.1.2",
  "mkdocs-include-markdown-plugin",
  "pymdown-extensions",
  "pygments",
]

dev = [
  "pygments",
  "click>=8.0.0",
  "prompt_toolkit>=3.0.0",
  "matplotlib>=3.7.0",
]




[tool.uv]
conflicts = [
  [
    { extra = "gpu" },
    { extra = "tpu" },
  ],
]  # Never solve for both GPU and TPU at once

[tool.ruff]
line-length = 119
target-version = "py311"
extend-exclude = ["scripts/"]

[tool.ruff.lint]
# Local-import hygiene (issue #6219), mirrored from the root config so these
# rules apply repo-wide: TID252 (absolute over parent-relative imports) and
# ICN001 (conventional import aliases, e.g. `import numpy as np`).
extend-select = ["TID252", "ICN001"]
ignore = [
  "E203", "E501", "W605", "F821", "E266", "F722", "E731", "E741"
]

[tool.black]
line-length = 119
target-version = ["py311"]

[tool.mypy]
python_version = "3.11"
mypy_path = ["src"]

[tool.mypy-haliax.core]
ignore_missing_imports = true

[tool.pytest.ini_options]
addopts = "-n auto"
pythonpath = ["src", "tests"]
filterwarnings = [
    "ignore:Transparent hugepages are not enabled\\..*:UserWarning:jax._src.cloud_tpu_init",
    "ignore:Read the `app_url` setting from the appropriate Settings object\\.:DeprecationWarning:wandb.analytics.sentry",
    "ignore:websockets\\.legacy is deprecated.*:DeprecationWarning:websockets.legacy",
    "ignore:websockets\\.server\\.WebSocketServerProtocol is deprecated:DeprecationWarning:uvicorn.protocols.websockets.websockets_impl",
    "ignore:'HTTP_422_UNPROCESSABLE_ENTITY' is deprecated\\. Use 'HTTP_422_UNPROCESSABLE_CONTENT' instead\\.:DeprecationWarning:gradio.routes",
    "ignore:Passing 'msg' argument to Task\\.cancel\\(\\) is deprecated since Python 3\\.11, and scheduled for removal in Python 3\\.14\\.:DeprecationWarning",
    "ignore:Passing 'msg' argument to Future\\.cancel\\(\\) is deprecated since Python 3\\.11, and scheduled for removal in Python 3\\.14\\.:DeprecationWarning",
]
markers = [
    "slow: marks tests as slow (deselect with '-m \"not slow\"')",
    "torch: mark tests that use Torch (deselect with '-m \"not torch\"')",
    "timeout: override the default per-test timeout",
]
asyncio_default_fixture_loop_scope = "function"