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.14,<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
  # Deliberately a range, not an exact pin: each backend extra pins the torch its
  # wheels were built against (vllm/mcore/automodel/trtllm: 2.13.0; sglang:
  # 2.11.0) and resolves in its own uv conflict fork, so the base spec must admit
  # both. The ceiling keeps the driver venv from drifting to a newer torch on a
  # relock; it resolves to the highest allowed (2.13.0), matching vLLM.
  "torch>=2.11,<2.14",
  "triton; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')",
  "colored==2.2.3",
  "ray[default]>=2.58.0",                                                                                # CVE mitigation
  "wandb>=0.30.0",                                                                                       # Address dependency vulnerabilities in bundled wandb-core
  "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.29",                                                                             # tracks the torch range above (0.26 <-> 2.11, 0.28 <-> 2.13)
  # 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 compatible with backend-specific pins while selecting 5.12.1
  # for the base and compatible backend environments in uv.lock.
  "transformers>=5.5.0,<=5.12.1",
  "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>=0.2.0",                                                                                         # Rust-backed BPE tokenizer (~10x faster encode); enable at runtime with NRL_USE_FASTOKENS=1
  # OpenTelemetry instrumentation. Base (not extra) so it reaches every worker
  # venv; `telemetry.enabled` in the config is the on/off switch. Sourced from
  # git in [tool.uv.sources] — see the note there; the floor is above the
  # newest PyPI release so a dropped pin fails to resolve instead of silently
  # falling back to 0.1.0. The aiohttp extra is what stamps a traceparent on
  # the HTTP calls NeMo-Gym's client makes, so a Gym rollout continues the
  # run's trace rather than starting its own.
  "nemo-lens[sdk,aiohttp]>=0.2.0",
  # Headers for transformer-engine's no-build-isolation install, and the runtime floor for
  # TE 2.18 fused-attention backward with 256-wide heads on SM 10.x (cuDNN >= 9.23.0 for
  # bshd/sbhd, >= 9.25.0 for thd; see TE fused_attn.cpp). Keep equal to the override below.
  "nvidia-cudnn-cu13==9.26.0.51; sys_platform != 'darwin'",
  # tilelang — required by fused DSA TileLang kernels on Blackwell and by the
  # replacement Triton kernel mamba-ssm requires when Triton >= 3.4.0 on Hopper;
  # see https://github.com/state-spaces/mamba/issues/640.
  # Without it, qwen3.5 / nano-v3 / moonlight Megatron recipes crash at first
  # gated-chunk backward with a RuntimeError pointing at this exact pip install.
  # Keep the base range broad for both supported Linux architectures; backend
  # extras tighten the version as required.
  "tilelang>=0.1.9,<0.1.13; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')",
  # 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",
  # TransferQueue v0.1.11 adds SimpleStorage SSD offload on top of v0.1.9's
  # full-system checkpoint save/load APIs and exact-key MooncakeStore clear
  # behavior required by DAPO.
  # Pin the immutable release commit rather than the mutable version tag.
  "TransferQueue @ git+https://github.com/Ascend/TransferQueue.git@05d4cc3377bee9b0f2897bbb90914ee20c0aea39",
  # 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.
  # NOTE: 0.3.13.post1 is the latest release on PyPI as of this bump, but still
  # vendors golang.org/x/net v0.55.0 (confirmed via strings on libetcd_wrapper.so) --
  # one version behind the fix for a known dependency vulnerability. This bump
  # narrows the gap (v0.48.0 -> v0.55.0) but does not fully close it; re-check
  # when a newer release ships.
  "mooncake-transfer-engine-cuda13==0.3.13.post1 ; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')",
]


[project.optional-dependencies]
automodel = [
  # Colocated with vLLM engines (CUDA-IPC refit); must share vLLM's torch ABI.
  "torch==2.13.0",
  "torchvision==0.28.0",
  "nemo-automodel[moe]",
  # Automodel's FLA kernels require this TileLang version on Hopper.
  # Keep this scoped to the automodel fork; sglang intentionally pins TileLang 0.1.8.
  "tilelang==0.1.11; sys_platform == 'linux' and platform_machine == 'x86_64'",
  # 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==2.8.1",
  # Mirrors nemo-automodel's own `transformers==5.12.1` pin (Automodel r0.6.0).
  "transformers==5.12.1",
  "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@17cfb817bccec3a9c247013360cc550c2bac441e",
]
vllm = [
  # ABI of the vLLM 0.29.0 cu130 wheels.
  "torch==2.13.0",
  "torchvision==0.28.0",
  "transformers==5.12.1",
  "tilelang==0.1.12; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')",
  "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@17cfb817bccec3a9c247013360cc550c2bac441e",
  # PyPI publishes the default cu130 wheels for both supported architectures.
  # https://github.com/vllm-project/vllm/releases/tag/v0.29.0
  "vllm==0.29.0",
  "num2words>=0.5.14",
  "flashinfer-python==0.6.18",
  "flashinfer-cubin==0.6.18",
  "flashinfer-jit-cache==0.6.18",
  "nvidia-cutlass-dsl[cu13]==4.6.2",
]
sglang = [
  # sglang-miles (3003d70f) and sglang-kernel 0.4.5 are built against torch 2.11;
  # on 2.13 sgl_kernel fails with `undefined symbol: _ZNK2at10TensorBase14const_data_ptr…`.
  "torch==2.11.0",
  "torchvision==0.26.0",
  # sgl-project/sglang branch `sglang-miles` (based on v0.5.16). Needed for the
  # begin_weight_update / end_weight_update refit session: quantized weights
  # carry derived scales (swizzled blockscales, g1/g2 alphas) that have to be
  # re-materialized once, after the last bucket lands. No tagged release has
  # those endpoints, and re-running the finalize per bucket is not idempotent.
  "sglang",
  "sglang-kernel==0.4.5",                                                       # Direct dep for explicit version control; required by sglang-miles
  "sglang-router",                                                              # Used by nemo_rl.models.generation.sglang.sglang_router for multi-engine routing
  "kernels>=0.14.1,<0.15",
  "flashinfer-python==0.6.14",
  "flashinfer-cubin==0.6.14",
  "flashinfer-jit-cache==0.6.14",
  "transformers==5.12.1",
  "tilelang==0.1.11; sys_platform == 'linux' and platform_machine == 'x86_64'",
]
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]",
  # Colocated with vLLM engines (CUDA-IPC refit); must share vLLM's torch ABI.
  "torch==2.13.0",
  "torchvision==0.28.0",
  "transformers==5.12.1",
  "tilelang==0.1.9; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')",
  # 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==2.8.1",
  # FlashAttention-4 (CuTe DSL). Required by megatron_cfg.flash_attention_version=4, which
  # zero_train_gen_mismatch / batch_invariant_mode pins on both train and generation.
  # Pinned to b19: b20+ need apache-tvm-ffi>=0.1.12, but the shared override keeps 0.1.11.
  "flash-attn-4==4.0.0b19 ; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')",
  # 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@17cfb817bccec3a9c247013360cc550c2bac441e",
  "flashinfer-python==0.6.18.post1",
  "flashinfer-cubin==0.6.18.post1",
  "flashinfer-jit-cache==0.6.18.post1",
  "flash_mla ; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')",
  "nvidia-cutlass-dsl[cu13]==4.6.2 ; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')",
  "torch-memory-saver>=0.0.9.post1",
  "nvshmem4py-cu13>=0.2.1",
  "cupy-cuda13x",
  # megatron-energon already flows in transitively (megatron-core[dev,mlm] declares it), so
  # this is not a mirror in the sense forbidden above -- it is a deliberately tighter floor:
  # Megatron-LM pins ~=7.0, and the energon SFT path requires 7.3+.
  # tests/unit/test_dependency_pins.py fails if Megatron-LM's own pin moves, so a
  # Megatron-Bridge bump forces a re-check of this floor.
  "megatron-energon[av-decode]~=7.3", # for energon backend
]
# 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 = [
  # feat/rl_rc's requirements.txt needs torch>=2.13.0,<=2.14.0a0.
  "torch==2.13.0",
  "torchvision==0.28.0",
  "tensorrt-llm==1.3.0rc28",
  # tensorrt-llm pins nvidia-cutlass-dsl[cu13]==4.8.0.dev0; uv only admits pre-releases
  # that are requested explicitly, so name its pre-release sub-packages here too.
  "nvidia-cutlass-dsl-libs-base==4.8.0.dev0",
  "nvidia-cutlass-dsl-libs-core==4.8.0.dev0",
  "nvidia-cutlass-dsl-libs-cu12==4.8.0.dev0",
  "nvidia-cutlass-dsl-libs-cu13==4.8.0.dev0",
  "transformers==5.5.4",
  "tilelang==0.1.9; sys_platform == 'linux' and platform_machine == 'x86_64'",
]
modelopt = ["nvidia-modelopt"]
nvrx = [
  "nvidia-resiliency-ext==0.7.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 (version comes from the extra being synced)
  "torch",
  # 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
  "sphinx-reredirects",        # Preserves URLs when documentation pages move
  "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 }
# Keep Gym as an editable path dependency rather than a workspace member. Both
# projects define a ``vllm`` extra with intentionally different server stacks;
# uv applies a requested extra name to every selected workspace member, which
# makes the RL vLLM + base Gym worker environment impossible to select.
nemo_gym = { path = "3rdparty/Gym-workspace/Gym", editable = true }
nemo_run = { git = "https://github.com/NVIDIA-NeMo/Run", rev = "414f0077c648fde2c71bb1186e97ccbf96d6844c" }
# Deliberately ahead of the v0.2.0 rev megatron-core pins: this rev carries the
# SpanRegistry that `span_groups` registers into and the metric registry the
# `rl` group is declared through, neither of which exists in v0.2.0. Keep in
# sync with the matching [tool.uv] override-dependencies entry, which is what
# actually reconciles the two revs; the floor in [dependencies] stays above the
# newest PyPI release so dropping either one fails to resolve instead of
# falling back to 0.1.0.
nemo-lens = { git = "https://github.com/NVIDIA-NeMo/Lens.git", rev = "b0f977d414b2f89938604a0b7eaa78ee08bc8700" }
sglang = { git = "https://github.com/sgl-project/sglang.git", rev = "3003d70f680d41c59d1b7acbf65cb47795dfd19e", subdirectory = "python" }
# 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" }
# FlashMLA derives its version from the Git revision in setup.py. Keep this pin
# synchronized with the static dependency metadata below.
flash_mla = { git = "https://github.com/deepseek-ai/FlashMLA", rev = "b7643bd54521f563b839b98289b5cd048c062ba2" }
# Pinned to 43fd41a58 (2026-08-15, release cherry-pick branch). It carries #2054 ("Fix vLLM
# 0.24+ compatibility": the vLLM plugin registers on `RoutedExperts`; the older plugin
# registered on `fused_moe.layer.FusedMoE`, which vLLM >= 0.28 removed, and NeMo-RL's explicit
# re-import then failed with "RowParallelLinear already registered") but NOT 2ded5aba3
# (#1550, 2026-08-07). Do not move to a revision that contains #1550 without porting
# nemo_rl/modelopt: its per-expert `GroupedQuantizer` modules break the Megatron fakequant
# DDP hook bookkeeping (`KeyError: GroupedQuantizer(...)` in `disable_forward_pre_hook`) and
# real-quant dist-checkpoint sharding validation (`Invalid access pattern for
# ShardedTensor(... experts.N.linear_fc1.weight_quantizer._amax`).
nvidia-modelopt = { git = "https://github.com/NVIDIA/Model-Optimizer", rev = "43fd41a58d52c4e6e5dec1d1ff5989ecc737ae1a" }
nvidia-resiliency-ext = { git = "https://github.com/NVIDIA/nvidia-resiliency-ext.git", tag = "v0.7.0" }
flashinfer-cubin = { index = "flashinfer" }
flashinfer-jit-cache = { index = "flashinfer-cu130" }

[tool.uv.workspace]
members = [
  # 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",
  "research/flow_grpo",
]

[[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"
# Keep this spelling (trailing slash included) identical to Megatron-Bridge's
# flashinfer [[tool.uv.index]] entry: uv treats differently spelled URLs for the
# same package as conflicting indexes and refuses to lock.
url = "https://flashinfer.ai/whl/"
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",
  "flash-attn",
  "mamba-ssm",
  "causal-conv1d",
  "fast-hadamard-transform",
  "deep_gemm",
  "deep_ep",
  "flash_mla",
  "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]==2.18.0",
  "nvidia-cublas==13.5.1.27; sys_platform != 'darwin'",
  "nvidia-cudnn-cu13==9.26.0.51; sys_platform != 'darwin'",
  "nvidia-cudnn-frontend==1.25.0 ; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')",
  # NCCL 2.30.4 fixes a context-parallel hang on GB200 when GPUs are non-contiguous;
  # 2.30.7 additionally fixes hangs in the NCCL m2n reshard (<link to NCCL release note / internal bug>).
  # Keep NCCL shared with the isolated Dynamo environment across torch upgrades.
  "nvidia-nccl-cu13==2.30.7",
  "nvidia-nvshmem-cu13>=3.6.5; sys_platform != 'darwin'",
  "opencv-python-headless; sys_platform == 'never'",
  "timm<=1.0.22",
  "fastapi[standard]>=0.133.0,<0.137.0",
  "nvidia-modelopt[torch]>=0.44.0a0",
  # vLLM 0.29 lists instanttensor (the optional `load_format="instanttensor"`
  # safetensors loader) as an unconditional dependency. NeMo-RL never selects
  # that load format and vLLM imports it lazily, so it is kept on x86_64 only;
  # on aarch64 (GB200) the dependency is dropped rather than installed. 0.2.0
  # ships aarch64 wheels, so this is a choice, not a wheel-availability limit.
  "instanttensor>=0.1.9; platform_machine == 'x86_64'",
  "llguidance>=1.7.0,<1.8.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@17cfb817bccec3a9c247013360cc550c2bac441e",
  # 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.0rc28 pins apache-tvm-ffi==0.1.13.post2, while xgrammar needs
  # >=0.1.9. The binding ceiling is nemo-automodel[cuda] (apache-tvm-ffi<=0.1.11)
  # together with tilelang 0.1.12 (<=0.1.11); tilelang 0.1.11 itself accepts
  # ~=0.1.0 and would admit 0.1.13.post2. Keep the shared override at 0.1.11 so
  # the automodel venv stays resolvable; the trtllm venv inherits it.
  "apache-tvm-ffi==0.1.11",
  # 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
  # Forces numpy past tensorrt-llm's `numpy>=2.0.0,<2.4` cap. The original
  # driver — TransferQueue pinning `numpy<2.0.0` — is gone as of TQ v0.1.9;
  # drop this override once the trtllm cap lifts. The `<2.5` ceiling keeps the
  # override inside numba's metadata bound: vLLM 0.29 pins `numba==0.65.0`,
  # which requires `numpy<2.5` and refuses to import against 2.5.x
  # ("Numba needs NumPy 2.4 or less"), leaving a broken numba in the vLLM venv.
  "numpy>=2.1.0,<2.5",
  # av (PyAV) bundles FFmpeg-derived codecs that may incur royalties. Exclude it from the
  # shipped container; audio/video tests install it via tools/install_audio_deps.sh.
  "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'",
  # megatron-bridge[ssm] pins mamba-ssm==2.3.1 for its patched Docker sdist build; we build
  # mamba-ssm from the git rev in [tool.uv.sources] instead.
  "mamba-ssm==2.3.2.post1",
  # megatron-core still pins lens v0.2.0 in its own [tool.uv.sources]; without an
  # override the two git URLs are a hard conflict, since the workspace resolves a
  # single nemo-lens. We need the newer rev for `SpanRegistry` (which replaced the
  # shipped `SpanGroup` class) and the consumer-driven metric registry
  # (`MetricSpec`/`register_metric_group`), both of which landed after v0.2.0. Drop
  # this once Megatron-LM bumps to the same or newer rev.
  "nemo-lens[sdk,aiohttp] @ git+https://github.com/NVIDIA-NeMo/Lens.git@b0f977d414b2f89938604a0b7eaa78ee08bc8700",
]

# 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.3.0",
  "fast-hadamard-transform @ git+https://github.com/Dao-AILab/fast-hadamard-transform.git@f134af63deb2df17e1171a9ec1ea4a7d8604d5ca",
  # Same story for megatron-fsdp, which nemo-automodel (r0.6.0) resolves from a
  # Megatron-LM fork via its own [tool.uv.sources] instead of PyPI. Keep this URL/rev
  # in sync with 3rdparty/Automodel-workspace/Automodel/pyproject.toml.
  "megatron-fsdp @ git+https://github.com/yuhezhang-ai/Megatron-LM.git@455389c480af6b3acdca74c7830c68b3274eb083#subdirectory=megatron/core/distributed/fsdp/src",
  # 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)
]

# NOTE: do not exclude `nvidia-cutlass-dsl-libs-base` again. The 4.5.x workaround for
# https://github.com/NVIDIA/cutlass/issues/3259 (libs-base and libs-cu13 wrote the same
# paths) is obsolete: 4.6 split the wheels so libs-base ships the whole Python `cutlass`
# DSL (`cutlass.cute`, `cutlass._mlir.dialects`, ...) and libs-cu13 only the CUDA-13
# `_cutlass_ir` extension, with no overlapping files. flashinfer >= 0.6.16 imports
# `cutlass.cute` at module load (`gdn_prefill` -> `delta_rule_dsl.varlen_helper`), and
# vLLM 0.28 imports `flashinfer.gdn_prefill` eagerly, so excluding libs-base breaks every
# Qwen3.5 model and every Blackwell vLLM worker with `ModuleNotFoundError: cutlass.cute`.

conflicts = [
  [{ extra = "automodel" }, { extra = "sglang" }],
  [{ extra = "mcore" }, { extra = "sglang" }],
  [{ extra = "vllm" }, { extra = "sglang" }],
  [{ 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" }],
  [{ extra = "trtllm" }, { extra = "modelopt" }],
  [{ extra = "trtllm" }, { extra = "nvrx" }],
  [{ extra = "trtllm" }, { extra = "nemo_gym" }],
  # Gym pins openai==2.44.0 (merged PR #2872 tree) while sglang==0.5.12.post1
  # pins openai==2.6.1; no environment installs both extras together.
  # Spelled nemo_gym (the raw extras-table key): uv < 0.12 compares conflict
  # entries against unnormalized extra names and rejects "nemo-gym".
  [{ extra = "nemo_gym" }, { extra = "sglang" }],
]


# 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-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 }]
flash_mla = [{ requirement = "torch", match-runtime = true }]

# torch-memory-saver's build backend reads TMS_CUDA_MAJOR at build time. Keep this in
# sync with docker/Dockerfile's BASE_IMAGE CUDA major version (currently cuda13.x).
[tool.uv.extra-build-variables]
torch-memory-saver = { TMS_CUDA_MAJOR = "13" }
# flash-attn 2.8.1 has no prebuilt wheel for torch >= 2.11, so it compiles from the sdist.
# Its setup.py sizes MAX_JOBS from psutil's *host* free memory (not the container cgroup)
# and cpu_count()//2, then runs each nvcc with --threads 4; on nemo-ci's build runners that
# fan-out gets the CUDA compiler OOM-killed (`ResourceExhausted`, dozens of
# `Killed "$CICC_PATH/cicc"`) on both x86 and aarch64. 4 jobs x 4 threads built cleanly on
# both architectures (2h47m aarch64, ~2h x86 alongside TE/flash-mla).
flash-attn = { MAX_JOBS = "4" }

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

[[tool.uv.dependency-metadata]]
name = "flash-mla"
version = "1.0.0+b7643bd"
requires-dist = []

[[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.3.2.post1"
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 = "1.2.1+17cfb81"
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 = "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"]

# 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"]
"research/*/tests/**" = ["D"]
# Ignore all files that end in `_test.py`.
"*_test.py" = ["D"]
# Ignore F401 (import but unused) in __init__.py
"__init__.py" = ["F401"]