pyproject.toml

[build-system]
requires = ["setuptools>=64", "setuptools_scm>=8"]
build-backend = "setuptools.build_meta"

[tool.setuptools.package-data]
oumi = ["LICENSE", "README.md", "*.jinja"]

[tool.setuptools_scm]
# default scm config.

[tool.setuptools]
# default setuptools config.

[project]
name = "oumi"
dynamic = ["version"]
readme = "README.md"
requires-python = ">=3.9"
description = "Oumi - Modeling Platform"
authors = [{ name = "Oumi Team", email = "contact@oumi.ai" }]
license = { file = "LICENSE" }
classifiers = [
    "Programming Language :: Python :: 3",
    "Programming Language :: Python :: 3.9",
    "Programming Language :: Python :: 3.10",
    "Programming Language :: Python :: 3.11",
    "Programming Language :: Python :: 3.12",
    "Programming Language :: Python :: 3.13",
    "Operating System :: OS Independent",
    "Development Status :: 3 - Alpha",
    "Intended Audience :: Developers",
    "Topic :: Scientific/Engineering :: Artificial Intelligence",
]

dependencies = [
    "accelerate>=1.2.1,<1.3",
    "aiohttp>=3.10,<3.12",          # Used by inference engine
    "aiofiles>=22.1.0,<25",         # Allows to use async file operations
    "aioresponses>=0.7.6,<0.8",     # User by inference engine tests
    "datasets>=3.2.0,<3.3",
    "jsonlines",
    "lm_eval[wandb]>=0.4.5,<0.5.0",
    "numpy>=1.26.4,<2.0",
    # The latest stable version, 2.3.0, is two years old and is causing dependency
    # conflicts. We'll use the dev version until a new stable version is released.
    # See https://github.com/oumi-ai/oumi/issues/1377
    "omegaconf>=2.4.0dev3,<2.5",
    "packaging",
    "pandas>=2.0.3,<3",
    "peft>=0.14.0,<0.15",
    "pexpect>=4.8.0,<4.9",          # Used by Polaris client
    "pillow>=10.3.0,<10.4",         # Used by image datasets
    "pydantic>=2.9.2,<2.10",
    "responses>=0.25.0,<0.26",
    "skypilot>=0.7.0,<0.8",         # Used by launcher
    "tensorboard>=2.18.0,<2.19",    # Optional, for monitoring training
    "torch>=2.5.0,<2.6.0",
    # torchdata 0.10 drops support for datapipes which we currently use
    "torchdata>=0.9.0,<0.10.0",  # Used by data pipes loader
    "torchvision>=0.20.0,<0.21", # Used by some VLM-s (multimodal)
    "tqdm",
    # Llama Vision attention is broken as late as 4.48.2 if gradient checkpointing is
    # enabled. See OPE-875 and https://github.com/huggingface/transformers/issues/36040.
    "transformers>=4.48.0,<4.49",
    "trl>=0.13.0,<0.14",
    "typer",                      # Used by CLI
    "typing_extensions",          # Backports of typing updates to python 3.9
    "wandb>=0.19.3,<0.20",        # Logging to Weights and Biases.
]

[project.urls]
Homepage = "https://github.com/oumi-ai/oumi"
Issues = "https://github.com/oumi-ai/oumi/issues"

[project.optional-dependencies]
dev = [
    "ipykernel",
    "ipywidgets",     # Interactive widgets for Jupyter notebooks
    "jupyterlab",
    "nbconvert",      # For notebook execution in tests
    "nbformat",       # For notebook execution in tests
    "pre-commit",
    "pyright",
    "pytest-asyncio", # Asyncio plugin for pytest
    "pytest-cov",     # Test coverage plugin for pytest
    "pytest-timeout", # Timeout plugin for pytest
    "pytest",
    "responses",
    "ruff",
    "torchfix",       # Tool for automatically fixing common PyTorch issues
]
docs = [
    "myst_parser",             # Allows us to write docs in markdown
    "myst-nb",                 # Allows us to include ipython notebooks in the docs
    "sphinx-book-theme",       # Alternative to sphinx-rtd-theme
    "sphinx-copybutton",       # Add copy button to code blocks in Sphinx
    "sphinx-design",           # Allows us to include tabs in the docs
    "sphinx-rtd-theme>=3.0.1", # Readthedocs theme for Sphinx
    "sphinx",                  # Used to build the docs
    "sphinxcontrib-bibtex",    # Allows us to cite bibtex references in the docs
    "sphinxcontrib-mermaid",   # Enable mermaid diagrams in the docs
    "sphinxcontrib-typer",     # Allows us to include typer CLI in the docs
]

# Dependencies that require a GPU to install
gpu = [
    "liger-kernel>=0.5.0,<0.6",
    "nvidia-ml-py>=12.560.30,<12.561",
    "bitsandbytes>=0.45.0,<0.46",      # Used for QLora, and PagedAdam implementation
]

# Targets for supported cloud providers
aws = ["skypilot[aws]"]
azure = ["skypilot[azure]"]
gcp = [
    "google-api-core>=2.19.0",
    "google-api-python-client",
    "google-auth>=2.30.0",
    "google-cloud-core>=2.4.1",
    "google-cloud-storage>=2.17.0",
    "skypilot[gcp]",
]
lambda = ["skypilot[lambda]"]
runpod = ["skypilot[runpod]"]

evaluation = [
    "alpaca-eval>=0.6.6,<0.7",
    "langdetect",              # leaderboard_ifeval
    "immutabledict",           # leaderboard_ifeval
    "nltk>=3.9.1",
    # "antlr4-python3-runtime==4.11",  # leaderboard_math_hard (incompatible)
    "sympy>=1.12",           # leaderboard_math_hard
    "sentencepiece>=0.1.98",
]

cambrian = [
    "timm==0.9.16",
    "open_clip_torch",
    "diffusers[torch]",
    "einops==0.6.1",
]

# CI targets
ci_cpu = [
    "oumi[dev,docs,gcp]",
    "vllm>=0.6.3,<0.7.0", # For VLLMInferenceEngine
    # This may fail to install. As a temporary workaround, run:
    # CMAKE_ARGS="-DLLAVA_BUILD=OFF" pip install -U llama-cpp-python
    "llama-cpp-python>=0.3.5,<0.4", # For LlamaCppInferenceEngine
]
# llama-cpp-python is not compatible with the github
# gpu actions runner, so we skip it for now
ci_gpu = [
    "oumi[dev,docs,gcp,gpu]",
    "vllm>=0.6.3,<0.7.0",
    "alpaca-eval>=0.6.6,<0.7",
]

[project.scripts]
oumi = "oumi.cli.main:run"

[tool.ruff]
extend-include = [
    "*.ipynb", # Include ipython notebooks
]
line-length = 88


[tool.ruff.lint]
select = [
    "ASYNC", # flake8-async: check for async/await syntax
    "D",     # pydocstyle: check docstring style
    "E",     # pycodestyle: errors
    "F",     # pyflakes: detect various errors
    "I",     # isort: check import order
    "ICN",   # flake-8 import conventions,
    "INT",   # flake8-gettext
    "LOG",   # flake8-logging-format: check for logging format,
    "NPY",   # NumPy-specific rules
    "PD",    # pandas vet
    "PTH",   # pathlib
    "Q",     # flake8-quotes
    "RSE",   # flake8-raise, use raise statements instead of return statements
    "TID",   # flake8-tidy-imports: check import tidiness
    "UP",    # pyupgrade: check for Python syntax updates
    "W",     # pycodestyle: warnings
]
ignore = [
    "D100",   # Missing docstring in public module, temporary, OPE-326
    "D101",   # Missing docstring in public class, temporary, OPE-326
    "NPY002", # Replace legacy numpy aliases
    "PTH123", # pathlib: `open()` should be replaced by `Path.open()`
]

[tool.ruff.lint.flake8-tidy-imports]
ban-relative-imports = "all" # Disallow all relative imports.

[tool.ruff.lint.per-file-ignores]
"tests/*" = ["D", "PTH"] # Ignore docstring checks in tests
"src/oumi/models/experimental/cambrian/**" = [
    "F401",   # Ignore pyflakes: `torch_xla` imported but unused;
    "F601",   # Ignore pyflakes: Dictionary key literal `"mistral_v2"` repeated
    "F841",   # Ignore pyflakes: Local variable is assigned but never used
    "E501",   # Ignore pycodestyle: Line too long
    "E701",   # Ignore pycodestyle: Multiple statements on one line
    "E721",   # Ignore pycodestyle: Do not compare types, use `isinstance()`
    "E722",   # Ignore pycodestyle: Do not use bare `except`
    "E731",   # Ignore pycodestyle: Do not assign a `lambda` expression, use a `def`
    "E741",   # Ignore pycodestyle: Ambiguous variable name
    "D",      # Ignore docstring checks in Cambrian
    "TID252", # Ignore flake8-tidy-imports: Prefer absolute imports over relative imports
    "W291",   # Ignore pycodestyle: warnings: Trailing whitespace
    "UP035",  # UP035 `typing.List` is deprecated, use `list` instead
    "UP006",  # UP006 Use `list` instead of `List` for type annotation
]

[tool.ruff.lint.pydocstyle]
convention = "google"

[tool.ruff.lint.isort]
known-first-party = ["oumi"]

[tool.pyright]
include = ["src/oumi/**", "tests/**"]
exclude = ["src/oumi/models/experimental/cambrian/**"]
typeCheckingMode = "basic"
pythonVersion = "3.9"
pythonPlatform = "All"                                 # Include linux, mac, windows

reportPrivateImportUsage = "none"

[tool.pytest.ini_options]
asyncio_default_fixture_loop_scope = "function"
testpaths = ["tests"]
addopts = ["--import-mode=prepend"]
markers = [
    "e2e: Slow e2e integration tests",
    "e2e_eternal: Extremely slow e2e integration tests (for manual/selective runs)",
    "single_gpu: The test uses max 1 GPU (can be potentially skipped on multi-GPU machine to conserve GPU resources)",
    "multi_gpu: The test should run on a machine with multiple GPU-s",
]