pyproject.toml

# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

[build-system]
requires = ["setuptools>=42", "wheel", "cython>=3.0.0"]
build-backend = "setuptools.build_meta"

[tool.setuptools]
include-package-data = true

[tool.setuptools.packages.find]
include = [
    "nemo_export",
    "nemo_export.*",
    "nemo_deploy",
    "nemo_deploy.*",
    "nemo_export_deploy_common",
    "nemo_export_deploy_common.*",
]

[tool.setuptools.dynamic]
version = { attr = "nemo_export_deploy_common.package_info.__version__" } # any module attribute compatible with ast.literal_eval
readme = { file = "README.md", content-type = "text/markdown" }


[project]
name = "NeMo-Export-Deploy"
dynamic = ["version", "readme"]
description = "NeMo Export and Deploy - a library to export and deploy LLMs and MMs"
requires-python = ">=3.10,<3.13"
license = { text = "Apache 2.0" }
dependencies = [
    "megatron-bridge>=0.2.0a0,<0.3.0",
    "megatron-core[mlm,dev]>=0.15.0a0,<0.16.0",
    "fastapi",
    "pydantic-settings",
    "ray",
    "ray[serve]",
    "uvicorn",
    "torch==2.7.1",
    "torchvision",
    "torchmetrics>=0.11.0",
    "wandb",
    "webdataset>=0.2.86",
    "nvidia-pytriton ; platform_system != 'Darwin' ",
    "Pillow ; platform_system != 'Darwin' and platform_machine != 'aarch64'",
    "decord ; platform_system != 'Darwin' and platform_machine != 'aarch64'",
    "pyparsing>2.0.2",
    "ijson",
    "pyarrow<21.0.0",
    "peft",

]

[project.optional-dependencies]
inframework = []
vllm = ["vllm~=0.10.0", "pandas", "timm"]
trtllm = ["tensorrt-llm>=1.0.0a0,<1.1.0,>=1.0.0rc6", "cuda-python~=12.8.0"]
trt-onnx = ["tensorrt==10.11.0.33", "onnx==1.18.0", "transformers==4.51.3"]

[dependency-groups]
# This is a default group so that we install these even with bare `uv sync`
build = ["setuptools", "torch==2.7.1", "pybind11", "Cython>=3.0.0", "ninja"]
docs = [
    "sphinx",
    "sphinx-autobuild",    # For live doc serving while editing docs
    "sphinx-autodoc2",     # For documenting Python API
    "sphinx-copybutton",   # Adds a copy button for code blocks
    "myst_parser",         # For our markdown docs
    "nvidia-sphinx-theme", # Our NVIDIA theme
]
linting = ["pre-commit>=3.6.0", "ruff~=0.9.0"]
test = ["pytest", "pytest-mock", "coverage", "click"]
nemo-toolkit = [
    "nemo-toolkit[automodel,common-only,nlp-only,eval,multimodal-only]>=2.5.0a0,<2.6.0",
    # Lightning deps
    "cloudpickle",
    "fiddle",
    "hydra-core>1.3,<=1.3.2",
    "lightning",
    "omegaconf>=2.3.0",
]
nemo-run = ["nemo-run"]

[tool.uv.sources]
xformers = [{ index = "pytorch-cu128" }]
torch = [{ index = "pytorch-cu128" }]
vllm = [
    { index = "pytorch-cu128", marker = "python_version < '3.9' and platform_machine == 'x86_64'" },
    { index = "pypi", marker = "platform_machine == 'aarch64'" },
]
# megatron-bridge = { git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git", rev = "ecf05926b4765aada82c8eabab4a374e8e83a9c5" }
transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "0289e76380088358a584d809faf69effab1a7cda" } # on release_v2.7.0
# nemo-toolkit = { git = "https://github.com/NVIDIA/NeMo.git", rev = "main" }

[tool.uv]
# Currently, TE must be built with no build-isolation b/c it requires torch
no-build-isolation-package = [
    "transformer-engine",
    "transformer-engine-torch",
    "flash-attn",
    "mamba-ssm",
    "causal-conv1d",
    "nv-grouped-gemm",
]
# Always apply the build group since dependencies like TE/mcore/nemo-run require build dependencies
# and this lets us assume they are implicitly installed with a simply `uv sync`. Ideally, we'd
# avoid including these in the default dependency set, but for now it's required.
default-groups = ["linting", "build", "test"]
# Users may use different link-modes depending on their scenario:
#  --link-mode=hardlink (default on linux; may get warnings about switching to --link-mode copy if uv cache and venv on different file-systems)
#  --link-mode=copy (slower but more reliable; supresses warning)
#  --link-mode=symlink (fastest option when uv cache and venv on different file-system; caveat: venv is brittle since it depends on the environment/container)
link-mode = "copy"
conflicts = [[{ extra = "trtllm" }, { extra = "vllm" }, { extra = "trt-onnx" }]]
override-dependencies = [
    "urllib3>1.27.0",
    "tiktoken>=0.9.0",                          # because nemo-toolkit and megatron-bridge disagree on tiktoken, we need to pin it here,
    "fsspec[http]>=2023.1.0,<=2024.9.0",
    "megatron-energon[av-decode]>=6.0,<7.dev0", # because nemo-toolkit and megatron-core disagree on megatron-energon, we need to pin it here,
]
prerelease = "allow"

[[tool.uv.index]]
name = "pypi"
url = "https://pypi.org/simple"
explicit = true

[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true

[tool.pytest.ini_options]
addopts = "--durations=15 -s -rA -x"
testpaths = ["tests"]
python_files = "test_*.py"

[tool.coverage.run]
concurrency = ["thread", "multiprocessing"]
omit = [
    "/tmp/*",
    "/workspace/tests/*",
    "/workspace/nemo_export_deploy_common/*",
    "/workspace/*.py",
    "*_ray.py",
    "*package_info.py",
]

[tool.coverage.paths]
source = [".", "/workspace", "/home/runner/work/Export-Deploy/Export-Deploy"]

[tool.ruff]
line-length = 120

[tool.ruff.format]
quote-style = "double"
skip-magic-trailing-comma = false
line-ending = "auto"

[tool.ruff.lint]
# Enable all `pydocstyle` rules, limiting to those that adhere to the
# Google convention via `convention = "google"`, below.
select = [
    "F541", # f-string without any placeholders
    "F841", # local variable assigned but never used
    "F401", # imported but unused
    "E741", # ambiguous variable name
    "F821", # undefined name
    "E266", # too many leading '#' for block comment
    "I",    # isort
    "D101", # docstring
    "D103",
]

ignore = [
    "E501", # Line too long - handled by formatter
    "D101",
    "D103",
]

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

# Section to exclude errors for different file types
[tool.ruff.per-file-ignores]
# Ignore all directories named `tests`.
"tests/**" = ["D"]
# Ignore all files that end in `_test.py`.
"*_test.py" = ["D"]
# Ignore F401 (import but unused) in __init__.py
"__init__.py" = ["F401"]