Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion .agents/skills/uv-build/SKILL.md
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ mise run setup && mise run bootstrap-nss cpu
mise run bootstrap-nss dev # dev tools only (no engine/torch)
mise run bootstrap-nss cpu # + engine + CPU PyTorch
mise run bootstrap-nss cu129 # + engine + CUDA 12.9 PyTorch
mise run bootstrap-nss cu130 # + engine + CUDA 13.0 PyTorch
mise run bootstrap-nss cuda # alias for cu129
mise run bootstrap-nss engine # + engine (no torch)
```
Expand All @@ -32,10 +33,11 @@ Under the hood: `uv sync --frozen --extra <extra> [--extra engine] --group dev`
|-------|------------------|
| `cpu` | PyTorch CPU, faiss-cpu, flashinfer (Linux only) |
| `cu129` | PyTorch+CUDA 12.9, faiss-gpu, flashinfer-jit-cache |
| `cu130` | PyTorch+CUDA 13.0, faiss-gpu, flashinfer-jit-cache |
| `engine` | ML pipeline deps (outlines, wandb, tiktoken, etc.) -- no torch |
| `microservices` | `nemo-microservices` from local path |

`cpu` and `cu129` conflict -- you must pick one, never both. Enforced in `[tool.uv] conflicts`.
`cpu`, `cu129`, and `cu130` conflict -- you must pick one, never both. Enforced in `[tool.uv] conflicts`.

## Index Management

Expand Down
1 change: 1 addition & 0 deletions .claude/commands/bootstrap.md
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ Set up the development environment from scratch.
```bash
mise run bootstrap-nss cpu # CPU-only (macOS or Linux without GPU)
mise run bootstrap-nss cuda # CUDA 12.9 (Linux with NVIDIA GPU)
mise run bootstrap-nss cu130 # CUDA 13.0 (Linux with NVIDIA GPU)
mise run bootstrap-nss engine # Engine dependencies only (no torch)
mise run bootstrap-nss dev # Minimal dev dependencies only
```
Expand Down
29 changes: 26 additions & 3 deletions .github/workflows/gpu-tests.yml
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ jobs:
uses: ./.github/actions/detect-changes

gpu-smoke-test:
name: GPU Smoke Tests
name: GPU Smoke Tests (${{ matrix.cuda-extra }})
needs: changes
# `changes` is intentionally skipped on workflow_dispatch. `always()` lets
# manual runs bypass that skipped dependency and run the selected GPU suite.
Expand All @@ -87,10 +87,19 @@ jobs:
}}
timeout-minutes: 30
runs-on: linux-amd64-gpu-a100-latest-1
continue-on-error: ${{ !matrix.required }}
strategy:
fail-fast: false
matrix:
python-version: ["3.13"]
include:
- python-version: "3.13"
cuda-extra: cu129
required: true
- python-version: "3.13"
# Non-required while the CUDA 13 rollout is validated. Promote to
# required once cu130 has container-build coverage (see docs/developer-guide/docker.md).
cuda-extra: cu130
required: false
steps:
- name: checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
Expand All @@ -102,6 +111,7 @@ jobs:
uses: ./.github/actions/setup-gpu-test-env
with:
python-version: ${{ matrix.python-version }}
cuda-extra: ${{ matrix.cuda-extra }}

- name: Run GPU smoke tests - train only
timeout-minutes: 10
Expand All @@ -128,7 +138,7 @@ jobs:
run: mise run test:smoke:gpu:smollm2

gpu-e2e-test:
name: GPU E2E Tests
name: GPU E2E Tests (${{ matrix.cuda-extra }})
needs: changes
# `changes` is intentionally skipped on workflow_dispatch. `always()` lets
# manual runs bypass that skipped dependency and run the selected GPU suite.
Expand All @@ -147,6 +157,17 @@ jobs:
}}
timeout-minutes: 210
runs-on: linux-amd64-gpu-a100-latest-1
continue-on-error: ${{ !matrix.required }}
strategy:
fail-fast: false
matrix:
include:
Comment on lines 157 to +164

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 E2E matrix missing python-version

The gpu-smoke-test matrix explicitly sets python-version: "3.13" on each entry, but the gpu-e2e-test matrix entries contain only cuda-extra and required. If the setup-gpu-test-env action has a required python-version input, the e2e job will silently receive an empty string. The pre-PR e2e job also didn't pass python-version, so the action likely has a working default — but it may be worth aligning the two jobs to make the Python version explicit and avoid future drift.

- cuda-extra: cu129
required: true
# Non-required while the CUDA 13 rollout is validated. Promote to
# required once cu130 has container-build coverage (see docs/developer-guide/docker.md).
- cuda-extra: cu130
required: false
steps:
- name: checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
Expand All @@ -156,6 +177,8 @@ jobs:

- name: Setup GPU test environment
uses: ./.github/actions/setup-gpu-test-env
with:
cuda-extra: ${{ matrix.cuda-extra }}

- name: Run GPU E2E tests
run: mise run test:e2e
Expand Down
9 changes: 6 additions & 3 deletions .mise/tasks/bootstrap-nss
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
#!/usr/bin/env bash
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#MISE description="Install project Python dependencies for a selected profile: dev, engine, cpu, cuda, or cu129. Modifies the project virtualenv."
#MISE description="Install project Python dependencies for a selected profile: dev, engine, cpu, cuda, cu129, or cu130. Modifies the project virtualenv."
#USAGE arg "<extra>" help="Python dependency profile" {
#USAGE choices "dev" "engine" "cpu" "cuda" "cu129"
#USAGE choices "dev" "engine" "cpu" "cuda" "cu129" "cu130"
#USAGE }

set -euo pipefail
Expand All @@ -17,6 +17,9 @@ case "${usage_extra?}" in
cuda|cu129)
uv sync --frozen --extra cu129 --extra engine --group dev
;;
cu130)
uv sync --frozen --extra cu130 --extra engine --group dev
;;
cpu)
uv sync --frozen --extra cpu --extra engine --group dev
;;
Expand All @@ -27,7 +30,7 @@ case "${usage_extra?}" in
uv sync --frozen --group dev
;;
*)
echo "Error: Invalid extra '${usage_extra?}'. Use one of: dev engine cpu cuda cu129" >&2
echo "Error: Invalid extra '${usage_extra?}'. Use one of: dev engine cpu cuda cu129 cu130" >&2
exit 1
;;
esac
1 change: 1 addition & 0 deletions CONTRIBUTING.md
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@ Please read our [Code of Conduct](CODE_OF_CONDUCT.md) before contributing.
# Install Python dependencies (choose one)
mise run bootstrap-nss cpu # CPU-only (macOS or Linux without GPU)
mise run bootstrap-nss cuda # CUDA 12.9 (Linux with NVIDIA GPU)
mise run bootstrap-nss cu130 # CUDA 13.0 (Linux with NVIDIA GPU)
mise run bootstrap-nss engine # Engine dependencies only
mise run bootstrap-nss dev # Minimal dev dependencies only
```
Expand Down
2 changes: 1 addition & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ $(1):
@mise run $(2)
endef

BOOTSTRAP_EXTRAS := dev engine cpu cuda cu129
BOOTSTRAP_EXTRAS := dev engine cpu cuda cu129 cu130
$(BOOTSTRAP_EXTRAS):
@:

Expand Down
6 changes: 4 additions & 2 deletions containers/Dockerfile.cuda
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,10 @@
# Multistage GPU Dockerfile for Safe-Synthesizer.
#
# CUDA support is selected through the project extras instead of a CUDA base
# image. The current distributable variant is cu129; adding cu130 later should
# only require a matching pyproject extra plus another workflow matrix row.
# image. The current distributable variant is cu129; the pyproject extra for
# cu130 exists but its container-build matrix row has not landed yet. Adding
# a variant's container coverage should only require a matching pyproject
# extra plus another workflow matrix row.
#
# Build:
# docker build -f containers/Dockerfile.cuda \
Expand Down
15 changes: 13 additions & 2 deletions cuda_deps.toml
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ torch_runtime_deps = [

cuda_runtime_deps = [
{ name = "flashinfer-jit-cache", version = "0.6.12", local = "{torch_local_version}", sys_platform = "linux", source_kind = "flashinfer" },
{ name = "nvidia-cublas-{cuda_package_suffix}", sys_platform = "linux" },
{ name = "nvidia-cublas{nvidia_package_suffix}", sys_platform = "linux" },
"nvidia-ml-py; sys_platform == 'linux'",
"torch-c-dlpack-ext",
{ name = "triton", specifier = ">=2.0.0", sys_platform = "linux", source_kind = "pytorch" },
Expand All @@ -35,7 +35,7 @@ torch_wheel_deps = [
{ name = "torchao", version = "0.17.0", local = "{torch_local_version}", sys_platform = "linux", arch = "x86_64", source_kind = "pytorch" },
]

managed_extras = ["cpu", "cu129"]
managed_extras = ["cpu", "cu129", "cu130"]
conflict_extras = ["cpu"]
nvidia_cuda_libraries = [
"cublas",
Expand Down Expand Up @@ -109,3 +109,14 @@ cuda_package_suffix = "cu12"
dependencies = [
"vllm==0.24.0+{torch_local_version}; sys_platform == 'linux'",
]

[variants.cu130]
cuda_package_suffix = "cu13"
nvidia_package_suffix = ""
nvidia_index = "nvidia-pypi-public"
flashinfer_index = { name = "flashinfer-jit-cache-cu130", url = "https://flashinfer.ai/whl/cu130", explicit = true }
# No [sources.cu130] vllm route yet: unlike cu129, there is no CUDA-13-specific vllm wheel
# build to pin, so this resolves from the default index. Add one here once that exists.
dependencies = [
"vllm==0.24.0; sys_platform == 'linux'",
]
Comment on lines +120 to +122

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 vllm for cu130 has no dedicated wheel source

cu129 routes vllm through an explicit wheel server entry ([sources.cu129]vllm-v0-24-0-cu129). For cu130, vllm==0.24.0 has no static source override, so uv will resolve it from the default PyPI index. Standard PyPI vllm wheels are typically CPU-only wrappers; if the PyPI wheel doesn't carry a CUDA 13 runtime, the package will install but any GPU operation will fail (or the wheel may not exist at all for 0.24.0). This is mitigated by required: false in CI, but it may cause confusing failures if someone tries mise run bootstrap-nss cu130 manually.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

2 changes: 1 addition & 1 deletion docs/developer-guide/docker.md
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ The variant name is intentionally the same as the CUDA package extra.
| Variant | Extra | Workflow status |
|---------|-------|-----------------|
| `cu129` | `cu129` | Enabled |
| `cu130` | `cu130` | Add when the CUDA 13.0 dependency PR lands |
| `cu130` | `cu130` | Manifest and GPU smoke-test coverage only; no published container image yet |

Adding a new variant should be mechanical:

Expand Down
82 changes: 81 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,7 @@ engine = [
# Source of truth: cuda_deps.toml.
# Regenerate with: uv run --frozen tools/gen_cuda_deps.py cuda_deps.toml --pyproject pyproject.toml
# Manual edits inside this block will be overwritten.
# Generated extras in this block: cpu, cu129.
# Generated extras in this block: cpu, cu129, cu130.
cpu = [
"accelerate>=1.1.0",
"bitsandbytes>=0.46.1",
Expand Down Expand Up @@ -165,6 +165,30 @@ cu129 = [
"torchao==0.17.0+cu129; sys_platform == 'linux' and platform_machine == 'x86_64'",
"vllm==0.24.0+cu129; sys_platform == 'linux'",
]
cu130 = [
"accelerate>=1.1.0",
"bitsandbytes>=0.46.1",
"flashinfer-python==0.6.12; sys_platform == 'linux'",
"flashinfer-cubin==0.6.12; sys_platform == 'linux'",
"gliner",
"kernels>=0.12.1,<0.15",
"opacus",
"peft>=0.18.0",
"sentence-transformers",
"trl>=0.23.0",
"xgrammar>=0.2.0; sys_platform == 'linux'",
"flashinfer-jit-cache==0.6.12+cu130; sys_platform == 'linux'",
"nvidia-cublas; sys_platform == 'linux'",
"nvidia-ml-py; sys_platform == 'linux'",
"torch-c-dlpack-ext",
"triton>=2.0.0; sys_platform == 'linux'",
"transformers>=5.12,<5.12.1",
"torch==2.11.0+cu130; sys_platform == 'linux'",
"torchaudio==2.11.0+cu130; sys_platform == 'linux'",
"torchvision==0.26.0+cu130; sys_platform == 'linux'",
"torchao==0.17.0+cu130; sys_platform == 'linux' and platform_machine == 'x86_64'",
"vllm==0.24.0; sys_platform == 'linux'",
]
# <<< END GENERATED CUDA RUNTIME EXTRAS - DO NOT EDIT >>>

[tool.uv]
Expand All @@ -186,6 +210,7 @@ conflicts = [
[
{extra = "cpu"},
{extra = "cu129"},
{extra = "cu130"},
],
]
# <<< END GENERATED CUDA UV CONFLICTS - DO NOT EDIT >>>
Expand Down Expand Up @@ -241,23 +266,29 @@ vllm = [
torch = [
{index = "pytorch-cpu", extra = "cpu", marker = "sys_platform == 'linux'"},
{index = "pytorch-cu129", extra = "cu129", marker = "sys_platform == 'linux'"},
{index = "pytorch-cu130", extra = "cu130", marker = "sys_platform == 'linux'"},
]
torchaudio = [
{index = "pytorch-cpu", extra = "cpu", marker = "sys_platform == 'linux'"},
{index = "pytorch-cu129", extra = "cu129", marker = "sys_platform == 'linux'"},
{index = "pytorch-cu130", extra = "cu130", marker = "sys_platform == 'linux'"},
]
torchvision = [
{index = "pytorch-cpu", extra = "cpu", marker = "sys_platform == 'linux'"},
{index = "pytorch-cu129", extra = "cu129", marker = "sys_platform == 'linux'"},
{index = "pytorch-cu130", extra = "cu130", marker = "sys_platform == 'linux'"},
]
flashinfer-jit-cache = [
{index = "flashinfer-jit-cache", extra = "cu129", marker = "sys_platform == 'linux'"},
{index = "flashinfer-jit-cache-cu130", extra = "cu130", marker = "sys_platform == 'linux'"},
]
triton = [
{index = "pytorch-cu129", extra = "cu129", marker = "sys_platform == 'linux'"},
{index = "pytorch-cu130", extra = "cu130", marker = "sys_platform == 'linux'"},
]
torchao = [
{index = "pytorch-cu129", extra = "cu129", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'"},
{index = "pytorch-cu130", extra = "cu130", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'"},
]
nvidia-cublas-cu12 = [
{index = "pytorch-cu129"},
Expand Down Expand Up @@ -298,6 +329,45 @@ nvidia-nvshmem-cu12 = [
nvidia-nvtx-cu12 = [
{index = "pytorch-cu129"},
]
nvidia-cublas = [
{index = "nvidia-pypi-public"},
]
nvidia-cuda-cupti = [
{index = "nvidia-pypi-public"},
]
nvidia-cuda-nvrtc = [
{index = "nvidia-pypi-public"},
]
nvidia-cuda-nvprof = [
{index = "nvidia-pypi-public"},
]
nvidia-cuda-runtime = [
{index = "nvidia-pypi-public"},
]
nvidia-cuda-nvml = [
{index = "nvidia-pypi-public"},
]
nvidia-cudnn = [
{index = "nvidia-pypi-public"},
]
nvidia-cusparse = [
{index = "nvidia-pypi-public"},
]
nvidia-cusparselt = [
{index = "nvidia-pypi-public"},
]
nvidia-nccl = [
{index = "nvidia-pypi-public"},
]
nvidia-nvjitlink = [
{index = "nvidia-pypi-public"},
]
nvidia-nvshmem = [
{index = "nvidia-pypi-public"},
]
nvidia-nvtx = [
{index = "nvidia-pypi-public"},
]

[[tool.uv.index]]
name = "pytorch-cu129"
Expand All @@ -309,6 +379,16 @@ name = "flashinfer-jit-cache"
url = "https://flashinfer.ai/whl/cu129"
explicit = true

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

[[tool.uv.index]]
name = "flashinfer-jit-cache-cu130"
url = "https://flashinfer.ai/whl/cu130"
explicit = true

[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
Expand Down
2 changes: 1 addition & 1 deletion tests/TESTING.md
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,7 @@ Tokenizers are function-scoped (expensive to load). Most fixtures are function-s

`ParsedResponse`: `valid_records=[...]`, `invalid_records=[...]`, `errors=[...]`, `prompt_number=int`. Use `fixture_mock_processor` or `fixture_mock_processor_without_valid_records`.

Optional dependencies: use `pytest.importorskip` to gate on packages that require specific extras. E2e tests use this for `sentence_transformers` and `vllm` (require `cu129` extra).
Optional dependencies: use `pytest.importorskip` to gate on packages that require specific extras. E2e tests use this for `sentence_transformers` and `vllm` (require a supported CUDA extra).

Mock Workdir via `mock_workdir(tmp_path)` in `cli/conftest.py`.

Expand Down
5 changes: 3 additions & 2 deletions tests/e2e/test_dataset_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,13 +15,14 @@
logger = get_logger(__name__)

llm = pytest.importorskip(
"vllm", reason="vllm with GPU support is required for these tests (install with: uv sync --extra cu129)"
"vllm",
reason="vllm with GPU support is required for these tests (install with: mise run bootstrap-nss cu129)",
)

try:
importlib.import_module("vllm")
except ImportError:
skip_reason = "vllm with GPU support is required for these tests (install with: uv sync --extra cu129)"
skip_reason = "vllm with GPU support is required for these tests (install with: mise run bootstrap-nss cu129)"
pytest.skip(skip_reason, allow_module_level=True) # ty: ignore[invalid-argument-type,too-many-positional-arguments]


Expand Down
7 changes: 4 additions & 3 deletions tests/e2e/test_safe_synthesizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,18 +18,19 @@
# Skip all tests in this module if sentence_transformers is not available
pytest.importorskip(
"sentence_transformers",
reason="sentence_transformers and a GPU are required for these tests (install with: uv sync --extra cu129)",
reason="sentence_transformers and a GPU are required for these tests (install with: mise run bootstrap-nss cu129)",
)

# Skip all tests in this module if vllm is not properly available.
vllm = pytest.importorskip(
"vllm", reason="vllm with GPU support is required for these tests (install with: uv sync --extra cu129)"
"vllm",
reason="vllm with GPU support is required for these tests (install with: mise run bootstrap-nss cu129)",
)

try:
importlib.import_module("vllm")
except ImportError:
skip_reason = "vllm with GPU support is required for these tests (install with: uv sync --extra cu129)"
skip_reason = "vllm with GPU support is required for these tests (install with: mise run bootstrap-nss cu129)"
pytest.skip(skip_reason, allow_module_level=True) # ty: ignore[invalid-argument-type,too-many-positional-arguments]


Expand Down
Loading
Loading