feat: streamline CUDA runtime dependency generation

Collapse runtime dependencies into the generated CUDA extras and update the generator, installer docs, and development workflow around the new CUDA-aware dependency model.

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: binaryaaron

.claude/commands/bootstrap.md

@@ -14,7 +14,8 @@ Set up the development environment from scratch.
    ```bash
    make bootstrap-nss cpu    # CPU-only (macOS or Linux without GPU)
    make bootstrap-nss cuda   # CUDA 12.8 (Linux with NVIDIA GPU)
-   make bootstrap-nss engine # Engine dependencies only (no torch)
+   make bootstrap-nss cu130  # CUDA 13.0 (Linux with NVIDIA GPU)
+   make bootstrap-nss docs   # Documentation dependencies only
    make bootstrap-nss dev    # Minimal dev dependencies only
    ```
 

.cursor/rules/repo-navigation.mdc

@@ -27,7 +27,7 @@ Cursor discovers skills directly from `.agents/skills/` -- no symlinks needed.
 
 | File | Purpose |
 |------|---------|
-| `pyproject.toml` | Package metadata, dependencies, extras (cpu/cu128/engine), uv config |
+| `pyproject.toml` | Package metadata, dependencies, runtime extras (cpu/cu128/cu130), uv config |
 | `pytest.ini` | Test markers, pytest options, timeout, parallelism |
 | `ruff.toml` | Ruff linting and formatting rules |
 | `mkdocs.yml` | Documentation site config (MkDocs Material) |

.cursor/setup-worktree.sh

@@ -16,10 +16,10 @@ fi
 
 # Bare --frozen installs the base environment. For GPU dev work (ty, import
 # checks, GPU tests) run the full command manually after setup:
-#   uv sync --frozen --extra cu128 --extra engine --group dev
+#   uv sync --frozen --extra cu128 --group dev
 uv sync --frozen
 echo "Venv ready: $(pwd)/.venv"
-echo "Note: for GPU extras run: uv sync --frozen --extra cu128 --extra engine --group dev"
+echo "Note: for GPU extras run: uv sync --frozen --extra cu128 --group dev"
 
 for _envfile in .env .env.local mise.local.toml .local.envrc; do
     if [ -f "$ROOT_WORKTREE_PATH/$_envfile" ]; then

AGENTS.md

@@ -24,7 +24,7 @@ Common commands: `make test` (unit tests), `make format` (auto-fix formatting +
 The canonical `uv sync` command for a full GPU/dev environment is:
 
 ```bash
-uv sync --frozen --extra cu128 --extra engine --group dev
+uv sync --frozen --extra cu128 --group dev
 ```
 
 Bare `uv sync --frozen` (without extras) installs an incomplete environment -- `ty`, import checks, and GPU tests will fail.

CONTRIBUTING.md

@@ -60,7 +60,8 @@ Please read our [Code of Conduct](CODE_OF_CONDUCT.md) before contributing.
    # Install Python dependencies (choose one)
    make bootstrap-nss cpu    # CPU-only (macOS or Linux without GPU)
    make bootstrap-nss cuda   # CUDA 12.8 (Linux with NVIDIA GPU)
-   make bootstrap-nss engine # Engine dependencies only
+   make bootstrap-nss cu130  # CUDA 13.0 (Linux with NVIDIA GPU)
+   make bootstrap-nss docs   # Documentation dependencies only
    make bootstrap-nss dev    # Minimal dev dependencies only
   ```
 

Makefile

@@ -83,27 +83,29 @@ verify-python-version: ## Verify Python version and install if necessary
 
 .PHONY: bootstrap-python
 bootstrap-python: .venv ## Bootstrap Python dependencies. Set PYTORCH_DEPS to 'cpu' or 'cu128'. Here mostly for legacy usage.
-	uv sync --frozen --extra ${PYTORCH_DEPS} --extra engine --group dev
+	uv sync --frozen --extra ${PYTORCH_DEPS} --group dev
 
 # Dynamic targets for bootstrap-nss
-# Usage: make bootstrap-nss {dev,engine,cpu,cuda}
-BOOTSTRAP_EXTRAS := dev engine cpu cuda cu128
+# Usage: make bootstrap-nss {dev,docs,cpu,cuda,cu128,cu130}
+BOOTSTRAP_EXTRAS := dev docs cpu cuda cu128 cu130
 $(BOOTSTRAP_EXTRAS):
 	@:
 
 .PHONY: bootstrap-nss
-bootstrap-nss: .venv ## Bootstrap Python dependencies. Usage: make bootstrap-nss {dev,engine,cpu,cuda}
+bootstrap-nss: .venv ## Bootstrap Python dependencies. Usage: make bootstrap-nss {dev,docs,cpu,cuda,cu128,cu130}
 	$(eval EXTRA := $(filter-out $@, $(MAKECMDGOALS)))
 	@echo "~~~~~~"
 	@echo "attempting to install nss package with primary extra: $(EXTRA)"
 	@if [ "$(EXTRA)" = "cuda" ]; then \
-		uv sync --frozen --extra cu128 --extra engine --group dev; \
+		uv sync --frozen --extra cu128 --group dev; \
 	elif [ "$(EXTRA)" = "cu128" ]; then \
-		uv sync --frozen --extra cu128 --extra engine --group dev; \
+		uv sync --frozen --extra cu128 --group dev; \
 	elif [ "$(EXTRA)" = "cpu" ]; then \
-		uv sync --frozen --extra cpu --extra engine --group dev; \
-	elif [ "$(EXTRA)" = "engine" ]; then \
-		uv sync --frozen --extra engine --group dev; \
+		uv sync --frozen --extra cpu --group dev; \
+	elif [ "$(EXTRA)" = "cu130" ]; then \
+		uv sync --frozen --extra cu130 --group dev; \
+	elif [ "$(EXTRA)" = "docs" ]; then \
+		uv sync --frozen --group docs; \
 	elif [ "$(EXTRA)" = "dev" ]; then \
 		uv sync --frozen --group dev; \
 	else \

README.md

@@ -36,23 +36,23 @@ Raw commands:
 CONSTRAINTS_URL="https://raw.githubusercontent.com/NVIDIA-NeMo/Safe-Synthesizer/main/constraints.txt"
 
 # CUDA 12.8 with uv (recommended):
-uv pip install "nemo-safe-synthesizer[cu128,engine]" \
+uv pip install "nemo-safe-synthesizer[cu128]" \
   -c "$CONSTRAINTS_URL" \
   --index https://flashinfer.ai/whl/cu128 \
   --index https://download.pytorch.org/whl/cu128 \
   --index https://pypi.nvidia.com \
   --index-strategy unsafe-best-match
 
 # CUDA 13.0 with uv (requires NVIDIA Linux driver 580.65.06+):
-uv pip install "nemo-safe-synthesizer[cu130,engine]" \
+uv pip install "nemo-safe-synthesizer[cu130]" \
   -c "$CONSTRAINTS_URL" \
   --index https://flashinfer.ai/whl/cu130 \
   --index https://download.pytorch.org/whl/cu130 \
   --index https://pypi.nvidia.com \
   --index-strategy unsafe-best-match
 
 # CUDA 12.8 with pip:
-pip install "nemo-safe-synthesizer[cu128,engine]" \
+pip install "nemo-safe-synthesizer[cu128]" \
   -c "$CONSTRAINTS_URL" \
   --extra-index-url https://download.pytorch.org/whl/cu128 \
   --extra-index-url https://flashinfer.ai/whl/cu128 \
@@ -65,9 +65,9 @@ Or install from source:
 git clone https://github.com/NVIDIA-NeMo/Safe-Synthesizer.git
 cd Safe-Synthesizer
 make setup # installs the pinned mise version (if missing) + pinned tool versions from mise.lock
-uv sync --frozen --extra cu128 --extra engine --group dev
+uv sync --frozen --extra cu128 --group dev
 # or, for CUDA 13.0:
-uv sync --frozen --extra cu130 --extra engine --group dev
+uv sync --frozen --extra cu130 --group dev
 ```
 
 Development tools (`ruff`, `ty`, `yq`, `gh`, etc.) are managed via [mise](https://mise.jdx.dev/). Tool versions are declared in `.mise.toml` and locked in `mise.lock` (committed). mise also manages environment variables -- place project-local secrets or overrides in `.env` or `.env.local` (both git-ignored, auto-loaded by mise).

containers/Dockerfile.cuda

@@ -5,7 +5,7 @@
 #
 # Stages:
 #   tools   -- install mise and all dev tools (uv, ruff, etc.)
-#   deps    -- install Python, uv, and all cu128+engine dependencies
+#   deps    -- install Python, uv, and all cu128 dependencies
 #   runtime -- minimal CLI wrapper (ENTRYPOINT entrypoint.sh -> safe-synthesizer)
 #   dev     -- extends runtime with dev tools, tests, and interactive shell
 #
@@ -142,14 +142,14 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # not when source code changes.
 COPY pyproject.toml uv.lock ./
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv sync --no-install-project --extra cu128 --extra engine --no-group dev
+    uv sync --no-install-project --extra cu128 --no-group dev
 
 # -- Project layer: install the project (non-editable so the venv is
 # self-contained and source code is not needed in the runtime image).
 COPY README.md ./
 COPY src/ src/
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv sync --no-editable --extra cu128 --extra engine --no-group dev
+    uv sync --no-editable --extra cu128 --no-group dev
 
 # ---------------------------------------------------------------------------
 # Stage 3: runtime -- minimal image wrapping the safe-synthesizer CLI
@@ -233,7 +233,7 @@ COPY . .
 
 # Install dev dependencies into the existing venv.
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv sync --extra cu128 --extra engine --group dev
+    uv sync --extra cu128 --group dev
 
 RUN git config --global --add safe.directory /workspace
 

cuda_deps.toml

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-base_dependencies = [
+requires_torch = [
   "accelerate",
   "bitsandbytes==0.49.1",
   { name = "flashinfer-python", version = "0.6.6", sys_platform = "linux", source_kind = "flashinfer", source_marker = "sys_platform=='linux'" },
@@ -17,6 +17,37 @@ base_dependencies = [
   "vllm==0.18.0; sys_platform == 'linux'",
 ]
 
+runtime_deps = [
+  "anyascii",
+  "pycountry",
+  "betterproto",
+  "flashtext",
+  "cached-property",
+  "category-encoders",
+  "dython",
+  "dateparser",
+  "faker",
+  "datasets>=4.8.4",
+  "huggingface-hub>=0.34.4,<1",
+  "json-repair",
+  "matplotlib",
+  "outlines>=1.0.0",
+  "prv-accountant",
+  "smart-open==7.0.5",
+  "python-stdnum",
+  "pandas>=2.1.3, <3",
+  "plotly",
+  "ratelimit",
+  "scikit-learn",
+  "range_regex>=0.1.0",
+  "tenacity==9.1.4",
+  "tiktoken>=0.7.0,<1.0",
+  "tldextract",
+  "tqdm>=4.67.1",
+  "urllib3>=2.6.1",
+  "wandb==0.26.1",
+]
+
 cuda_dependencies = [
   { name = "flashinfer-jit-cache", version = "0.6.6", local = "{torch_local_version}", sys_platform = "linux", source_kind = "flashinfer", source_marker = "sys_platform=='linux'", variants = ["cu128"] },
   { name = "nvidia-cublas", sys_platform = "linux" },
@@ -28,7 +59,7 @@ cuda_dependencies = [
   { name = "xformers", version = "v0.0.34", sys_platform = "linux", arch = "x86_64", source_kind = "pytorch" },
 ]
 
-generated_extras = ["cpu", "cu128", "cu130"]
+generated_extras = ["engine", "cpu", "cu128", "cu130"]
 
 # CUDA Python packaging context:
 # - Starting with CUDA Python 12.8, cuda-python is a metapackage that currently

docs/developer-guide/docker.md

@@ -24,7 +24,7 @@ flowchart TD
 
     subgraph stages [Build Stages]
         tools["tools\nInstalls mise + all dev tools\n(.mise.toml is single source of truth)"]
-        deps["deps\nInstalls Python 3.11 via uv\nuv sync cu128+engine"]
+        deps["deps\nInstalls Python 3.11 via uv\nuv sync cu128"]
         runtime["runtime\nCopies venv + Python\nNon-root appuser\ntini + entrypoint.sh"]
         dev["dev\nExtends runtime\nCopies mise tree from tools\nRoot user"]
     end
@@ -44,7 +44,7 @@ flowchart TD
   Uses the [mise Docker cookbook](https://mise.jdx.dev/mise-cookbook/docker.html)
   pattern with `MISE_DATA_DIR=/mise` for stable, copyable paths.
 - deps: copies the uv binary from `tools`, then installs Python and all
-  cu128+engine dependencies. Uses `--mount=type=cache` to avoid
+  cu128 dependencies. Uses `--mount=type=cache` to avoid
   re-downloading ~10 GB of PyTorch/CUDA wheels.
 - runtime: copies the venv and uv-managed Python into a fresh CUDA runtime
   base. Runs as non-root `appuser` (uid 1000). GPU access is declared via
@@ -263,7 +263,7 @@ To reduce size:
 | Aspect | `Dockerfile.cuda` | `Dockerfile.test_ci` |
 |--------|-------------------|----------------------|
 | Base | `nvidia/cuda:12.8.1-runtime-ubuntu22.04` | `python:3.11-slim` |
-| Extras | `cu128` + `engine` | `cpu` + `engine` |
+| Extras | `cu128` | `cpu` |
 | GPU | Required | Not needed |
 | Stages | `tools` / `deps` / `runtime` / `dev` | `setup` / `install-deps` |
 | Use case | Training, generation, evaluation | CPU-only unit tests and CI checks |