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", "torch"]
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.12,<3.13"
license = { text = "Apache 2.0" }
dependencies = [
    "megatron-core",
    "megatron-bridge",
    "nvidia-modelopt[torch]; sys_platform != 'darwin'",
    "nvidia-resiliency-ext; sys_platform != 'darwin'",
    "transformer-engine[pytorch, core_cu13]; sys_platform != 'darwin'",
    "accelerate",
    "fastapi",
    "pydantic-settings",
    "ray",
    "ray[serve]",
    "uvicorn",
    "tensorstore",
    "zarr>=2.18.2,<3.0.0",
    # Lightning deps
    "cloudpickle",
    "fiddle",
    "hydra-core>1.3,<=1.3.2",
    "lightning<2.5.0",
    "omegaconf>=2.3.0",
    "peft",
    "torch>=2.6.0",
    "torchvision",
    "torchmetrics>=0.11.0",
    "wandb",
    "webdataset>=0.2.86",
    "nvidia-pytriton ; platform_system != 'Darwin' ",
    "flashinfer-python>=0.2.5 ; platform_system != 'Darwin'",
    "Pillow ; platform_system != 'Darwin' and platform_machine != 'aarch64'",
    "decord ; platform_system != 'Darwin' and platform_machine != 'aarch64'",
    "pyparsing>2.0.2",
    "sentencepiece",
    "tiktoken",
    "einops",
    "ijson",
    "pyarrow>=20.0.0,<21.0.0",
]

[project.optional-dependencies]
inframework = []
vllm = ["vllm~=0.18.0", "pandas", "timm"]
trtllm = ["tensorrt-llm~=1.1.0", "cuda-python~=13.0.0"]
trt-onnx = [
    "tensorrt==10.14.1.48.post1",
    "onnx==1.21.0",
    "onnxscript>=0.6.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", "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[common-only,nlp-only,multimodal-only]>=2.5.0a0,<2.7.0",
    "nv_one_logger_core>=2.3.1",
    "nv_one_logger_training_telemetry>=2.3.1",
    "nv_one_logger_pytorch_lightning_integration>=2.3.1",
]
nemo-run = ["nemo-run"]

[tool.uv.sources]
xformers = [{ index = "pytorch-cu130" }]
torch = [{ index = "pytorch-cu130" }]
vllm = [
    { index = "pytorch-cu130", 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 = "1b66a8c119f1b1337ef968d0053bffc7529ba9f3" }
# 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",
    "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 = [
    "torch; sys_platform == 'never'",
    "torchvision; sys_platform == 'never'",
    "triton; sys_platform == 'never'",
    "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",
    "transformer-engine[pytorch]>=2.12.0a0,<2.14.0",
    "open-clip-torch>=3.2.0",
    "megatron-energon[av-decode]>=6.0,<7.dev0",
    "flashinfer-python>=0.3.0,<0.4.0",
    "datasets>=3.3.0",
    "flash-linear-attention>=0.3.0,<0.4.dev0",
    "patchelf; sys_platform=='never'",
    "nvidia-resiliency-ext>=0.3.0,<0.6.0",
    "transformer-engine[pytorch,core_cu13]>=2.12.0a0,<2.15.0; sys_platform != 'darwin'",
    "transformer-engine-cu13>=2.12.0a0,<2.15.0; sys_platform != 'darwin'",
    "transformer-engine-cu12; sys_platform == 'never'",
    "mamba-ssm>=2.3.0,<2.4.0",
    "transformers>=5.0.0",
    "protobuf~=6.33.5",
    "opencv-python-headless; sys_platform == 'never'",
    "cryptography>=43.0.0,<47",
    "onnx<1.20.0",
]
prerelease = "allow"

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

[[tool.uv.index]]
name = "pytorch-cu130"
url = "https://download.pytorch.org/whl/cu130"
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
    "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.lint.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"]