pyproject.toml

[build-system]
requires = ["setuptools>=78.1.1", "wheel>=0.46.2"]
build-backend = "setuptools.build_meta"

[tool.setuptools]
packages = ["nemo_rl"]

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

[project]
name = "nemo-rl"
dynamic = ["version", "readme"]
description = "NeMo RL: A Scalable and Efficient Post-Training Library for Models Ranging from 1 GPU to 1000s, and from Tiny to >100B Parameters"
requires-python = ">=3.13.13,<3.14"
license = { text = "Apache 2.0" }
maintainers = [
  { name = "Yi-Fu Wu", email = "yifuw@nvidia.com" },
  { name = "Terry Kong", email = "terryk@nvidia.com" },
  { name = "Yuki Huang", email = "yukih@nvidia.com" },
  { name = "NVIDIA", email = "nemo-toolkit@nvidia.com" },
]
dependencies = [
  "setuptools",
  "pip",                                                                                                 # Required for frozen environments; uv venv --seed may not reliably install pip
  "ninja",                                                                                               # for flash-attn parallel build
  "torch==2.11.0",
  "triton; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')",
  "colored==2.2.3",
  "ray[default]>=2.55.1",
  "wandb>=0.28.0",
  "numpy",
  "datasets>=4.0.0",
  "rich",
  "math-verify",
  "accelerate>=0.26",
  "tensorboard",
  "omegaconf",
  "torchdata",
  "nvidia-ml-py",
  "hydra-core",
  "tiktoken",
  "sentencepiece",
  "blobfile",
  "debugpy",
  "nvtx",
  "matplotlib",
  "plotly",
  "sympy>=1.14.0",
  "pillow>=12.3.0",                                                                                      # Address CVE (Pillow 12.3.0)
  "torchvision==0.26.0",
  # torchaudio and torchcodec are NOT installed in the container image.
  # Run tools/install_audio_deps.sh before using audio features or audio tests.
  # Keep the base range broad enough while allowing backend extras to tighten as needed.
  # Floor raised to 5.5.0 for Gemma 4 support.
  "transformers>=5.5.0,<5.9.0",
  "num2words>=0.5.14",                                                                                                       # for SmolVLM
  "timm",
  "open-clip-torch>=3.2.0",
  "mlflow>=3.13.0",
  "nvidia-nvshmem-cu13>=3.6.5; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')", # for deep_ep build
  "swanlab",
  "pyzmq",
  "soundfile>=0.13.1",
  "nccl4py; sys_platform != 'darwin'",                                                                                       # for non-colocated refit
  "cuda-bindings; sys_platform != 'darwin'",                                                                                 # for non-colocated refit
  # Worker venvs run a base sync before backend extras, so NIXL must be available here.
  "nixl==1.3.0; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')",
  "pybase64",                                                                                                 # for sglang refit
  "awscrt>=0.35.0",                                                                                           # for parallel S3 refit transport
  "zstandard",                                                                                                # for sparse refit body compression
  "fastokens-b10>=0.1.1",                                                                                     # Rust-backed BPE tokenizer (~10x faster encode); enable at runtime with NRL_USE_FASTOKENS=1
  "nvidia-cudnn-cu13==9.20.0.48; sys_platform != 'darwin'",                                                   # for transformer-engine no build isolation
  # tilelang — replacement Triton kernel mamba-ssm requires when
  # Triton >= 3.4.0 on Hopper, see github.com/state-spaces/mamba#640.
  # Without this, qwen3.5 / nano-v3 / moonlight megatron recipes
  # crash at first gated-chunk backward with a RuntimeError pointing
  # at this exact pip install. Linux x86_64 only — mamba-ssm itself
  # is gated to that pair.
  "tilelang; sys_platform == 'linux' and platform_machine == 'x86_64'",
  # Data-plane stack — promoted to base so worker venvs (built by
  # nemo_rl.utils.venvs.create_local_venv via bare `uv sync`, no extras)
  # automatically include them. Removes the need for a `[data-plane]`
  # extra and the corresponding plumbing in the per-worker venv builder.
  "tensordict",
  # Pinned to b266d39 (post-0.1.6, pre-0.1.7) for PR #77's MooncakeStore
  # refactor: `clear` switched from unanchored `remove_by_regex` to
  # exact-key `batch_remove`, which fixes a collateral-key-deletion bug
  # that breaks DAPO + mooncake_cpu. Bump to the 0.1.7 tag when released.
  "TransferQueue @ git+https://github.com/Ascend/TransferQueue.git@b266d39",
  # Backs data_plane.backend="mooncake_cpu". Default backend is "simple"
  # (in-process), but the mooncake_cpu path needs the `mooncake_master`
  # binary that ships in this wheel at <site-packages>/mooncake/. Bundled
  # with TQ rather than gated behind an extra so worker venvs (built
  # without extras) can be flipped to mooncake_cpu via config alone.
  # PyPI's base `mooncake-transfer-engine` is cu12-only (links
  # libcudart.so.12), which breaks on cu13 containers. Upstream now also
  # publishes a cu13 variant as a separate distribution name
  # `mooncake-transfer-engine-cuda13` (same `mooncake/` import namespace,
  # store.so linked against libcudart.so.13). Resolve from PyPI rather
  # than the GitHub release URL — the wheel is byte-identical (verified
  # sha256), and PyPI's CDN is far more reliable than github releases
  # from compute nodes.
  # Upstream publishes both x86_64 and aarch64 wheels (see uv.lock).  CI's
  # build-container runner is aarch64 (uv reports aarch64-unknown-linux-gnu),
  # so the marker must include aarch64 — otherwise mooncake is silently
  # excluded from the resolution during the Docker build.
  "mooncake-transfer-engine-cuda13==0.3.11.post1 ; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')",
]

[project.optional-dependencies]
fsdp = [
  # +cu13 wheels from GitHub match torch cu130; PyPI often resolves to +cu12 (libcudart.so.12).
  # https://github.com/Dao-AILab/flash-attention/releases/tag/v2.8.1
  "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_aarch64.whl ; sys_platform == 'linux' and platform_machine == 'aarch64'",
  "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl ; sys_platform == 'linux' and platform_machine == 'x86_64'",
  "flash-attn==2.8.1 ; sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'x86_64')",
  "mamba-ssm",
  "causal-conv1d",
]
automodel = [
  "nemo-automodel[moe]",
  # Flash-attn version should be selected to satisfy both TE + vLLM requirements (xformers in particular)
  # https://github.com/NVIDIA/TransformerEngine/blob/v2.3/transformer_engine/pytorch/attention/dot_product_attention/utils.py#L108
  # https://github.com/facebookresearch/xformers/blob/8354497deb2c04c67fbb2e2ad911e86530da0e90/xformers/ops/fmha/flash.py#L76
  "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_aarch64.whl ; sys_platform == 'linux' and platform_machine == 'aarch64'",
  "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl ; sys_platform == 'linux' and platform_machine == 'x86_64'",
  "flash-attn==2.8.1 ; sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'x86_64')",
  "transformers>=5.5.0,<5.6.0",
  "mistral-common>=1.11.0",
  "mamba-ssm",
  "causal-conv1d",
  "nv-grouped-gemm",
  "transformer-engine[pytorch,core_cu13] @ git+https://github.com/NVIDIA/TransformerEngine.git@v2.14.1",
  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@29d31c095796f3c8ece47ee9cdcc167051bbeed9 ; platform_machine == 'x86_64'",
  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@a48493600c4886c1b297aaa78db0e1ebc2d8dd6c ; platform_machine == 'aarch64'",
]
vllm = [
  "cuda-python",
  "deep_gemm @ git+https://github.com/deepseek-ai/DeepGEMM.git@67fc64863d43521080bf2005e6528d0fceee9510",
  # deep_ep also needs libibverbs-dev
  # sudo apt-get update
  # sudo apt-get install libibverbs-dev
  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@29d31c095796f3c8ece47ee9cdcc167051bbeed9 ; platform_machine == 'x86_64'",
  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@a48493600c4886c1b297aaa78db0e1ebc2d8dd6c ; platform_machine == 'aarch64'",
  # Default wheels on GitHub are cu130. See v0.25.1 release assets:
  # https://github.com/vllm-project/vllm/releases/tag/v0.25.1
  "vllm @ https://github.com/vllm-project/vllm/releases/download/v0.25.1/vllm-0.25.1-cp38-abi3-manylinux_2_28_aarch64.whl ; sys_platform == 'linux' and platform_machine == 'aarch64'",
  "vllm @ https://github.com/vllm-project/vllm/releases/download/v0.25.1/vllm-0.25.1-cp38-abi3-manylinux_2_28_x86_64.whl ; sys_platform == 'linux' and platform_machine == 'x86_64'",
  "vllm==0.25.1 ; sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'x86_64')",
  "num2words>=0.5.14",
  "flashinfer-python==0.6.13",
  "flashinfer-cubin==0.6.13",
  "flashinfer-jit-cache==0.6.13",
  "nvidia-cutlass-dsl[cu13]==4.5.2",
]
sglang = [
  "sglang==0.5.12.post1",
  "sglang-kernel==0.4.2.post2",         # Direct dep for explicit version control
  "sglang-router",                      # Used by nemo_rl.models.generation.sglang.sglang_router for multi-engine routing
  "kernels>=0.12.0,<0.13",
  "flashinfer-python==0.6.11.post1",
  "flashinfer-cubin==0.6.11.post1",
  "flashinfer-jit-cache==0.6.11.post1",
  "transformers==5.6.0",
]
mcore = [
  # also need cudnn (https://developer.nvidia.com/cudnn-downloads?target_os=Linux&target_arch=x86_64&Distribution=Ubuntu&target_version=20.04&target_type=deb_network)
  # wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.1-1_all.deb
  # sudo dpkg -i cuda-keyring_1.1-1_all.deb
  # sudo apt-get update
  # sudo apt-get install cudnn-cuda-13

  # megatron-bridge is installed straight from the submodule's own pyproject.toml, so its
  # dependencies (transformers, transformer-engine, megatron-core[dev,mlm] -> nvidia-modelopt,
  # onnxscript, flash-linear-attention, ...) flow transitively and pick up upstream changes on
  # every submodule bump + `uv lock`. Do not mirror them here. transformer-engine's exact
  # version is still pinned globally via [tool.uv] override-dependencies.
  # [te]/[ssm] activate megatron-bridge's transformer-engine and mamba-ssm/causal-conv1d extras.
  # megatron-core is intentionally NOT listed: it flows from megatron-bridge's own
  # dependency megatron-core[dev,mlm] and its [tool.uv.sources] path mapping, so Megatron-LM
  # pyproject changes propagate on bump without any mirror here.
  "megatron-bridge[te,ssm]",
  # Flash-attn version should be selected to satisfy both TE + vLLM requirements (xformers in particular)
  # https://github.com/NVIDIA/TransformerEngine/blob/v2.3/transformer_engine/pytorch/attention/dot_product_attention/utils.py#L108
  # https://github.com/facebookresearch/xformers/blob/8354497deb2c04c67fbb2e2ad911e86530da0e90/xformers/ops/fmha/flash.py#L76
  "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_aarch64.whl ; sys_platform == 'linux' and platform_machine == 'aarch64'",
  "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl ; sys_platform == 'linux' and platform_machine == 'x86_64'",
  "flash-attn==2.8.1 ; sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'x86_64')",
  # mamba-ssm/causal-conv1d LOOK redundant (megatron-bridge[ssm] pulls them), but do not
  # remove them: [tool.uv.sources] git pins only attach to direct requirements, and the
  # conflicting extras resolve in separate forks — without these direct deps the mcore fork
  # silently falls back to the PyPI releases instead of the pinned git revs.
  "mamba-ssm",
  "causal-conv1d",
  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@29d31c095796f3c8ece47ee9cdcc167051bbeed9 ; platform_machine == 'x86_64'",
  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@a48493600c4886c1b297aaa78db0e1ebc2d8dd6c ; platform_machine == 'aarch64'",
  "flashinfer-python==0.6.8.post1",
  "flashinfer-cubin==0.6.8.post1",
  "flashinfer-jit-cache==0.6.8.post1",
  "nvshmem4py-cu13>=0.2.1",
  "cupy-cuda13x",
]
# TRT-LLM is built from source via 3rdparty/TensorRT-LLM-workspace's custom PEP 517
# backend (_backend.py). The path source lets uv lock resolve the dep graph
# without compiling (prepare_metadata_for_build_wheel returns static metadata).
# The actual ~60-min wheel build is triggered by `uv sync --extra trtllm`.
trtllm = ["tensorrt-llm==1.3.0rc21"]
modelopt = ["nvidia-modelopt"]
nvrx = [
  "nvidia-resiliency-ext==0.6.0",
] # for ft_launcher (fault-tolerant training launcher)
nemo_gym = ["nemo_gym"]

[dependency-groups]

# This is a default group so that we install these even with bare `uv sync`
build = [
  # Build requirement for TE
  "torch==2.11.0",
  # Build requirement for TE
  "setuptools",
  "packaging",
  "einops",
  # Build requirement for nemo_run
  "hatchling",
  # Build requirement for mcore
  "pybind11",
  # Build requirement for flash-attn
  "psutil",
]
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
  "sphinx-design",             # For design components in docs
  "myst_parser",               # For our markdown docs
  "nvidia-sphinx-theme",       # Our NVIDIA theme
  "gitpython>=3.1.45",         # To git-related information
  "python-dotenv",             # For environment variable management
  "sphinxcontrib-mermaid",     # For Mermaid diagram support
  "swagger-plugin-for-sphinx", # For Swagger/OpenAPI documentation
]
dev = [
  "pre-commit>=4.2.0",
  "ruff==0.9.9",
  "types-PyYAML",
  "types-requests",
  "pyrefly==0.24.2",
]
test = [
  "pytest>=8.4.2",
  "pytest-timeout",
  "pytest-cov",
  "pytest-asyncio",
  "pytest-testmon",
  "pytest-shard",
]

[tool.uv.sources]
# tensorrt-llm is built from source via a custom PEP 517 backend in 3rdparty/TensorRT-LLM-workspace/.
# prepare_metadata_for_build_wheel returns static metadata fast (for `uv lock`);
# build_wheel triggers the actual ~60-min compilation (during `uv sync --extra trtllm`).
tensorrt-llm = { path = "3rdparty/TensorRT-LLM-workspace" }
nemo-automodel = { path = "3rdparty/Automodel-workspace/Automodel", editable = true }
megatron-bridge = { path = "3rdparty/Megatron-Bridge-workspace/Megatron-Bridge", editable = true }
nemo_gym = { workspace = true }
nemo_run = { git = "https://github.com/NVIDIA-NeMo/Run", rev = "414f0077c648fde2c71bb1186e97ccbf96d6844c" }
# torch/torchvision/triton all come from the torch index in order to pick up aarch64 wheels
torch = [
  { index = "pytorch-cu130", marker = "sys_platform != 'darwin'" },
  { index = "pypi", marker = "sys_platform == 'darwin'" },
]
torchvision = [
  { index = "pytorch-cu130", marker = "sys_platform != 'darwin'" },
  { index = "pypi", marker = "sys_platform == 'darwin'" },
]
triton = [
  { index = "pytorch-cu130", marker = "sys_platform != 'darwin'" },
  { index = "pypi", marker = "sys_platform == 'darwin'" },
]
causal-conv1d = { git = "https://github.com/Dao-AILab/causal-conv1d", rev = "4f6ae4e26ae5fe8af9372f8d312ab25cc4595223" }
mamba-ssm = { git = "https://github.com/state-spaces/mamba.git", rev = "a14b1dff0454a3bc27d9eb31355dc01e4b2490ec" }
nv-grouped-gemm = { git = "https://github.com/fanshiqing/grouped_gemm", tag = "v1.1.4.post7" }
nvidia-modelopt = { git = "https://github.com/NVIDIA/Model-Optimizer", rev = "c3b913b9cc1d82d5a0af9fa77b4db87829e6f158" }
nvidia-resiliency-ext = { git = "https://github.com/NVIDIA/nvidia-resiliency-ext.git", tag = "v0.6.0" }
flashinfer-jit-cache = { index = "flashinfer-cu130" }

[tool.uv.workspace]
members = [
  "3rdparty/Gym-workspace/Gym",
  # Research projects are also added here in order for them to share the global root level uv.lock.
  # If we don't do this, the research projects do not see the global uv.lock, and may mistakenly
  # install numpy>=2.0 because nemo-rl's core [dependencies] do not pin numpy, but when you inspect
  # nemo-rl's uv.lock you'll see it's 1.X b/c megatron mandates 1.X in the optional dependencies, so
  # globally we must choose 1.X otherwise we run into pickle issues from ray.
  "research/template_project",
]

[[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.uv.index]]
name = "flashinfer-cu130"
url = "https://flashinfer.ai/whl/cu130"
explicit = true

[tool.uv]
managed = true
environments = [
  "python_version == '3.13' and sys_platform == 'linux' and platform_machine == 'x86_64'",
  "python_version == '3.13' and sys_platform == 'linux' and platform_machine == 'aarch64'",
]
preview = true # Enable preview features like extra-build-dependencies
no-build-isolation-package = [
  "transformer-engine-torch",
  "transformer-engine",
  "flash-attn",
  "mamba-ssm",
  "causal-conv1d",
  "fast-hadamard-transform",
  "deep_gemm",
  "deep_ep",
  "nv-grouped-gemm",          # from mlm (added here to make sure it's built no isolation since mlm workspace uses setup.py)
  "tensorrt-llm",             # custom PEP 517 backend calls build-custom-trtllm.sh; needs torch/CUDA from main venv
]
# 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 = ["dev", "build"]
# 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"
# The TE override is needed to pin the TE version across all extras
# The opencv-python-headless override excludes it from the shipped container (sys_platform == 'never')
# because its bundled FFmpeg codec libs (H.264, H.265, AAC) may incur royalties. The CVE floor
# (>=5.0.0) is kept separately in constraint-dependencies so it applies to anyone who does install it.
# Installed at runtime via tools/install_audio_deps.sh (--no-config bypasses the override there).
# vllm defaults to opencv for video IO but falls back to torchcodec.
# The timm override is needed because sglang requires timm==1.0.16.
override-dependencies = [
  "transformer-engine[pytorch,core_cu13] @ git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.15",
  "nvidia-cublas==13.5.1.27; sys_platform != 'darwin'",
  "nvidia-cudnn-cu13==9.20.0.48; sys_platform != 'darwin'",
  "nvidia-cudnn-frontend==1.23.0",
  # NCCL 2.30.4 fixes a context-parallel hang on GB200 when GPUs are non-contiguous.
  # An override is required because torch 2.11.0+cu130 pins nvidia-nccl-cu13==2.28.9.
  "nvidia-nccl-cu13==2.30.4",
  "nvidia-nvshmem-cu13>=3.6.5; sys_platform != 'darwin'",
  "opencv-python-headless; sys_platform == 'never'",
  "timm<=1.0.22",
  "nvidia-modelopt[torch]>=0.44.0a0",
  "torch==2.11.0",
  # sglang has conflicting llguidance versions than vllm, so enforcing vllm's version since it's newer
  "llguidance>=1.3.0,<1.4.0",
  # Override setuptools range in other dependencies to address CVE GHSA-58pv-8j8x-9vj2
  "setuptools>=80.10.2",
  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@29d31c095796f3c8ece47ee9cdcc167051bbeed9 ; platform_machine == 'x86_64'",
  "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@a48493600c4886c1b297aaa78db0e1ebc2d8dd6c ; platform_machine == 'aarch64'",
  # Relax megatron-core workspace member's opentelemetry-api ceiling (<1.34) for protobuf 6.x compat with ray
  "opentelemetry-api>=1.33.1",
  # Override sglang's xgrammar==0.1.32 to address CVE GHSA-7rgv-gqhr-fxg3.
  # Floor raised to 0.2.1 for vLLM 0.25.1, whose tool-choice handling imports
  # xgrammar.normalize_tool_choice (vLLM requires xgrammar>=0.2.1,<1.0.0; the
  # override would otherwise force an incompatible downgrade).
  "xgrammar>=0.2.1,<1.0.0",
  # tensorrt-llm 1.3.0rc21 pins apache-tvm-ffi==0.1.6, but the xgrammar the
  # line above selects needs >=0.1.9, and the xgrammar releases that would
  # accept 0.1.6 ship no cp313 wheels. Take the newer one, as with llguidance
  # above; without this the trtllm extra is unsatisfiable on vLLM 0.25.1.
  "apache-tvm-ffi>=0.1.9",
  # Override dependencies to address CVEs
  "mlflow>=3.13.0",
  # Override langchain-core to address CVE GHSA-qh6h-p6c9-ff54 (path traversal in load_prompt)
  "langchain-core>=1.2.22",
  # Override outlines for Python 3.13 support
  "outlines>=0.2.0",
  # Upgrade pytest to 9.0.3
  "pytest>=9.0.3",
  "langchain>=0.3.28",      # Address CVE-2025-65106
  "langchain-core>=0.3.80", # Address CVE-2025-65106
  # TransferQueue (data-plane extra) pins numpy<2.0.0; megatron-core needs
  # numpy>=2.1.0 via onnx → ml-dtypes. Override globally so the data-plane
  # extra composes with mcore/automodel without version-mirroring TQ's
  # requirements.txt. Forward-compatible across TQ minor bumps.
  "numpy>=2.1.0",
  # av (PyAV) carries CVE-bundled codec libs (libx264, libx265, libopenh264, libmp3lame).
  # It is only needed by megatron-bridge's optional WAN diffusion path, which installs it
  # at test time via scripts/install_diffusion_deps.sh. Exclude from the shipped container.
  "av; sys_platform == 'never'",
  # decord2 is pulled transitively by sglang → qwen-vl-utils[decord] on aarch64 Linux.
  # NeMo-RL no longer calls any decord API (replaced by torchvision/torchaudio).
  "decord2; sys_platform == 'never'",
]

# CVE fixes
constraint-dependencies = [
  # emerging-optimizers / fast-hadamard-transform flow transitively from megatron-core[dev],
  # whose [tool.uv.sources] pins them to git. uv requires such URL deps to also appear as a
  # direct requirement or constraint of the root project. Keep these URLs/revs in sync with
  # Megatron-LM's [tool.uv.sources] (uv errors loudly on mismatch after a submodule bump).
  "emerging-optimizers @ git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.2.0",
  "fast-hadamard-transform @ git+https://github.com/Dao-AILab/fast-hadamard-transform.git@f134af63deb2df17e1171a9ec1ea4a7d8604d5ca",
  # megatron-energon pulls s3fs; without a floor uv picks the 2020-era s3fs 0.4.2
  # (which needs only plain botocore) instead of a modern aiobotocore-based release.
  "s3fs>=2024.6.0",
  "brotli>=1.2.0",            # Address CVE GHSA-2qfp-q593-8484
  "starlette>=0.49.1",        # Address CVE GHSA-7f5h-v6xp-fcq8
  "urllib3>=2.7.0",           # Address CVE GHSA-38jv-5279-wg99
  "aiohttp>=3.13.3",          # Address CVE GHSA-mqqc-3gqh-h2x8
  "pyasn1>=0.6.3",            # Address CVE GHSA-jr27-m4p2-rc6r
  "wheel>=0.46.2",            # Address CVE GHSA-8rrh-rw8j-w5fx
  "protobuf>=6.33.5",         # Address CVE GHSA-7gcm-g887-7qv7
  "python-multipart>=0.0.22", # Address CVE GHSA-wp53-j4wj-2cfg
  "pygments>=2.20.0",         # Address CVE GHSA-5239-wwwm-4pmq
  "cbor2>=5.9.0",             # Address CVE GHSA-3c37-wwvx-h642
  "onnx>=1.21.0rc4",          # Address CVE GHSA-hqmj-h5c6-369m
  "cryptography>=48.0.1",     # Address CVE GHSA-6w46-j5rx-g56g, GHSA-537c-gmf6-5ccf
  "orjson>=3.11.6",           # Address CVE GHSA-hx9q-6w63-j58v
  "pyjwt>=2.13.0",
  "dulwich>=1.2.5",
  "diffusers>=0.38.0",
  # 8.0.1 fixed _IncludedRouter.path crash on fastapi>=0.137 (github.com/trallnag/prometheus-fastapi-instrumentator/issues/370)
  "prometheus-fastapi-instrumentator>=8.0.2",
  "av>=17.1.0",                               # Address CVE-2026-40962 (FFmpeg 8.0.1 bundled in av.libs; 17.1.0+ ships FFmpeg 8.1.x)
  "opencv-python-headless>=5.0.0",            # CVE floor: patches CVE-2025-9951, CVE-2025-1594, CVE-2024-31582 (bundled FFmpeg 5.1.6 in 4.x)
]

# Workaround for https://github.com/NVIDIA/cutlass/issues/3259: libs-base
# and libs-cu13 wheels overlap and race on shared paths. libs-cu13 is a
# functional superset at 4.5.2, so drop libs-base to make installs deterministic.
exclude-dependencies = ["nvidia-cutlass-dsl-libs-base"]

conflicts = [
  [
    { extra = "fsdp" },
    { extra = "sglang" },
  ],
  [
    { extra = "automodel" },
    { extra = "sglang" },
  ],
  [
    { extra = "mcore" },
    { extra = "sglang" },
  ],
  [
    { extra = "vllm" },
    { extra = "sglang" },
  ],
  [
    { extra = "automodel" },
    { extra = "fsdp" },
  ],
  [
    { extra = "automodel" },
    { extra = "mcore" },
  ],
  [
    { extra = "automodel" },
    { extra = "vllm" },
  ],
  [
    { extra = "vllm" },
    { extra = "mcore" },
  ],
  [
    { extra = "trtllm" },
    { extra = "vllm" },
  ],
  [
    { extra = "trtllm" },
    { extra = "sglang" },
  ],
  [
    { extra = "trtllm" },
    { extra = "automodel" },
  ],
  [
    { extra = "trtllm" },
    { extra = "mcore" },
  ],
]


# Augment build dependencies for packages that need torch at build time
[tool.uv.extra-build-dependencies]
flash-attn = [{ requirement = "torch", match-runtime = true }]
# Git-sourced packages CAN use match-runtime = true if we provide dependency-metadata
deep_ep = [{ requirement = "torch", match-runtime = true }]
deep_gemm = [{ requirement = "torch", match-runtime = true }]
transformer-engine = [{ requirement = "torch", match-runtime = true }]
transformer-engine-torch = [{ requirement = "torch", match-runtime = true }]
mamba-ssm = [{ requirement = "torch", match-runtime = true }]
causal-conv1d = [{ requirement = "torch", match-runtime = true }]
fast-hadamard-transform = [{ requirement = "torch", match-runtime = true }]
nv-grouped-gemm = [{ requirement = "torch", match-runtime = true }]

# Needed when building from source
[[tool.uv.dependency-metadata]]
name = "flash-attn"
requires-dist = ["torch", "einops", "setuptools", "psutil", "ninja"]

[[tool.uv.dependency-metadata]]
name = "causal-conv1d"
# This version has to match the version in the commit/rev/tag used
version = "1.5.4"
requires-dist = ["torch", "packaging", "ninja"]

[[tool.uv.dependency-metadata]]
name = "fast-hadamard-transform"
version = "1.1.0"
requires-dist = ["torch", "packaging", "ninja"]

[[tool.uv.dependency-metadata]]
name = "mamba-ssm"
# This version has to match the version in the commit/rev/tag used
version = "2.2.6.post3"
requires-dist = ["torch", "packaging", "ninja", "causal-conv1d"]

[[tool.uv.dependency-metadata]]
name = "deep_ep"
# This version has to match the version in the commit/rev/tag used
version = "v1.2.1+29d31c0"
requires-dist = ["torch", "packaging", "ninja"]

[[tool.uv.dependency-metadata]]
name = "deep_gemm"
# This version has to match the version in the commit/rev/tag used
version = "2.5.0"
requires-dist = ["torch", "packaging", "ninja"]

[[tool.uv.dependency-metadata]]
name = "transformer-engine"
version = "2.15.0+42b8400"
requires-dist = ["torch", "pydantic", "importlib-metadata>=1.0", "packaging"]

[[tool.uv.dependency-metadata]]
name = "transformer-engine-torch"
version = "2.15.0+42b8400"
requires-dist = ["torch", "transformer-engine"]

[[tool.uv.dependency-metadata]]
name = "nv-grouped-gemm"
# This version has to match the version in the commit/rev/tag used
version = "v1.1.4.post7"
requires-dist = ["setuptools", "wheel", "torch", "numpy"]

# TensorRT: wheel_stub packages download multi-GB binaries when built from source.
# Providing static metadata here lets `uv lock` resolve the dep graph without
# triggering the download. At `uv sync --extra trtllm` time (inside Docker where
# there is sufficient disk space), uv installs them normally via wheel_stub.
# Metadata sourced from https://pypi.org/pypi/tensorrt*/10.16.1.11/json.
[[tool.uv.dependency-metadata]]
name = "tensorrt"
version = "10.16.1.11"
requires-dist = ["tensorrt-cu13==10.16.1.11"]

[[tool.uv.dependency-metadata]]
name = "tensorrt-cu13"
version = "10.16.1.11"
requires-dist = [
  "tensorrt-cu13-libs==10.16.1.11",
  "tensorrt-cu13-bindings==10.16.1.11",
]

# wheel_stub — downloads the actual C++ library binary at install time.
[[tool.uv.dependency-metadata]]
name = "tensorrt-cu13-libs"
version = "10.16.1.11"
requires-dist = []

# wheel_stub — downloads the Python bindings binary at install time.
[[tool.uv.dependency-metadata]]
name = "tensorrt-cu13-bindings"
version = "10.16.1.11"
requires-dist = []

# Override logsage metadata to:
# 1. Remove numpy<=2.0.2 and pandas<=2.3.3 upper bounds
#    (numpy cap conflicts with onnx>=1.21.0rc4 CVE fix via nvidia-resiliency-ext)
# 2. Bump langchain ecosystem to 1.x to address CVE GHSA-qh6h-p6c9-ff54
#    (langchain-core path traversal in legacy load_prompt, fix requires >=1.2.22)
# Tracking: https://github.com/NVIDIA/nvidia-resiliency-ext/issues/301
[[tool.uv.dependency-metadata]]
name = "logsage"
version = "0.1.7"
requires-dist = [
  "drain3>=0.9.11,<0.10.0",
  "langchain>=1.2.0,<2.0.0",
  "langchain-core>=1.2.22,<2.0.0",
  "langchain-nvidia-ai-endpoints>=1.0.0,<2.0.0",
  "nh3>=0.3.1,<0.4.0",
  "numpy",
  "pandas",
  "pydantic-settings>=2.11.0,<3.0.0",
  "requests>=2.32.5,<3.0.0",
]

# Override drain3 to relax cachetools==4.2.1 pin (conflicts with mlflow's cachetools>=5.0.0)
# Tracking: https://github.com/logpai/Drain3/issues/119
[[tool.uv.dependency-metadata]]
name = "drain3"
version = "0.9.11"
requires-dist = ["jsonpickle", "cachetools>=4.2.1"]

[tool.black]
line-length = 120
include = '\.pyi?$'
exclude = '''
/(
    \.git
  | \.venv
  | build
)/
'''

[tool.pytest.ini_options]
addopts = "--durations=100 -s -rA -x"
testpaths = ["tests"]
python_files = "test_*.py"
markers = [
  "run_first: marks tests that should run before others",
  "mcore: marks tests that require the mcore extra",
  "hf_gated: marks tests that require HuggingFace token access for gated models",
  "automodel: marks tests that require the automodel extra",
  "vllm: marks tests that require the vllm extra",
  "sglang: marks tests that require the sglang extra",
  "trtllm: marks tests that require the trtllm extra",
  "nemo_gym: marks tests that require the nemo_gym extra",
]

[tool.pyrefly]
project-includes = ["**/*"]
project-excludes = ["**/*venv/**/*"]

[tool.coverage.run]
concurrency = ["thread", "multiprocessing"]
omit = ["/tmp/*"]

[tool.coverage.paths]
source = ["nemo_rl/", "/opt/nemo-rl/nemo_rl/"]

[tool.ruff.lint]
# Enable all `pydocstyle` rules, limiting to those that adhere to the
# Google convention via `convention = "google"`, below.
select = ["D", "F"]

# - On top of the Google convention, disable `D417`, which requires
#   documentation for every function parameter.
# - F841: local variable assigned but never used (exluced to favor readability)
# TODO: Remove D10 once we are about to release to get all the docstrings written
ignore = ["D417", "D10", "F841"]

[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"]