build: generate CUDA dependency metadata from one source

.agents/skills/uv-build/SKILL.md

@@ -55,7 +55,7 @@ PyTorch wheels come from dedicated indexes, not PyPI:
 | `pytorch-cpu` | `download.pytorch.org/whl/cpu` | torch, torchvision (CPU, Linux) |
 | `pytorch-cu129` | `download.pytorch.org/whl/cu129` | torch, torchvision, triton (CUDA) |
 | `nv-shared-pypi-local` | NVIDIA Artifactory | Internal NVIDIA packages |
-| `flashinfer-jit-cache` | `flashinfer.ai/whl/cu129` | FlashInfer JIT cache |
+| `flashinfer-jit-cache-cu129` | `flashinfer.ai/whl/cu129` | FlashInfer JIT cache |
 | `nvidia-pypi-public` | `pypi.nvidia.com` | Public NVIDIA packages |
 
 All indexes are `explicit = true` (only used when a package is mapped to them in `[tool.uv.sources]`).
@@ -70,12 +70,17 @@ uv add <package>
 uv add --group dev <package>
 uv add --group test <package>
 
-# Add to an optional extra
-# (edit pyproject.toml manually, then lock)
+# Change CPU or CUDA runtime extras
+# Edit cuda_deps.toml, regenerate pyproject.toml, then lock.
+uv run --frozen tools/gen_cuda_deps.py cuda_deps.toml --pyproject pyproject.toml
 uv lock
 ```
 
 After any change: `uv lock` to regenerate `uv.lock`. Pre-commit verifies the lock is up to date.
+The generated CPU/CUDA sections of `pyproject.toml` must not be edited directly;
+`mise run lock-check` verifies that they match `cuda_deps.toml`. The
+generator owns the complete `[tool.uv.sources]` and `[[tool.uv.index]]`
+sections, so add every source or index there through `cuda_deps.toml`.
 
 ## Dependency Groups
 
@@ -135,5 +140,5 @@ Build backend: `hatchling` with wheel target `packages = ["src/nemo_safe_synthes
 2. Use `--frozen` in CI and Make targets to prevent lock updates
 3. Use `uv run` to run tools (pytest, mkdocs, etc.)
 4. uv version is pinned in `.mise.toml`.
-5. Edit extras manually in `pyproject.toml`, then `uv lock`
+5. Edit non-generated `pyproject.toml` sections directly (e.g. dependency groups); CPU/CUDA extras go through `cuda_deps.toml` instead, then `uv lock`
 6. Use `uv add` for base/group deps

.github/workflows/README.md

@@ -118,8 +118,8 @@ The `ci-checks.yml` workflow runs on every push to `main` and on pull requests.
 
 | Job | mise task | What it checks |
 | --- | --- | --- |
-| Format | `format-check` | `ruff format --check` + `ruff check` + SPDX copyright headers |
-| Format (lock) | `lock-check` | `uv.lock` matches `pyproject.toml` |
+| Format | `format-check` | `dprint check` for TOML + `ruff format --check` + `ruff check` + SPDX copyright headers |
+| Format (lock) | `lock-check` | generated CUDA metadata matches `cuda_deps.toml` + `uv.lock` matches `pyproject.toml` |
 | Typecheck | `typecheck` | `ty check` (excludes per `pyproject.toml [tool.ty.src]`) |
 | Unit Tests | `test:ci` | pytest with coverage (excludes slow, e2e, gpu, smoke) |
 | Smoke Tests | `test:smoke` | CPU smoke tests (training/generation hot paths, tiny models) |
@@ -134,7 +134,7 @@ To replicate CI locally:
 
 ```bash
 mise run check        # format-check + typecheck
-mise run lock-check   # verify uv.lock
+mise run lock-check   # verify generated CUDA metadata and uv.lock
 mise run test:ci      # CI unit tests with coverage selectors
 mise run test:smoke   # CPU smoke tests
 ```

.mise.toml

@@ -7,7 +7,6 @@
 # Tools NOT managed here:
 #   - buildctl: only needed for container builds. Install manually or via `brew install buildkit`.
 
-
 # Single source of truth for the pinned mise version. `tools/install-mise.sh`
 # reads `min_version` from this file; an older mise on PATH fails fast here too.
 min_version = '2026.7.5'
@@ -28,7 +27,6 @@ ty = "0.0.44"
 
 "aqua:j178/prek" = "latest"
 
-
 [env]
 # Activate the project virtualenv (equivalent to the old .envrc "layout python3").
 VIRTUAL_ENV = "{{ env.UV_PROJECT_ENVIRONMENT | default(value=config_root ~ '/.venv') }}"
@@ -37,12 +35,12 @@ _.file = [".env", { path = ".env.local", redact = true }]
 
 [task_config]
 includes = [
-  ".mise/tasks",
-  ".mise/tasks/*.toml",
+    ".mise/tasks",
+    ".mise/tasks/*.toml",
 ]
 
 [settings]
 lockfile = true
-install_before = "7d"  # only resolve to versions released more than 7 days ago - can help mitigate supply chain attacks
+install_before = "7d" # only resolve to versions released more than 7 days ago - can help mitigate supply chain attacks
 lockfile_platforms = ["macos-arm64", "linux-x64", "linux-arm64"]
 idiomatic_version_file_enable_tools = []

.mise/tasks/docs.toml

@@ -15,6 +15,6 @@ run = "uv run --frozen --group docs mkdocs build"
 description = "Deploy the latest docs to GitHub Pages with mike. Requires git push permission and docs dependencies."
 alias = "docs-deploy"
 run = [
-  "uv run --frozen --group docs mike deploy --push --update-aliases main latest",
-  "uv run --frozen --group docs mike set-default --push latest",
+    "uv run --frozen --group docs mike deploy --push --update-aliases main latest",
+    "uv run --frozen --group docs mike set-default --push latest",
 ]

.mise/tasks/quality.toml

@@ -2,29 +2,32 @@
 # SPDX-License-Identifier: Apache-2.0
 
 [format]
-description = "Format all tracked source files with ruff and SPDX copyright fixer. Modifies files in place; run before committing."
+description = "Format TOML and Python files with dprint and ruff, then fix SPDX copyright headers. Modifies files in place; run before committing."
 run = [
-  "bash tools/codestyle/format.sh",
-  "uv run --script tools/codestyle/copyright_fixer.py .",
+    "dprint fmt --incremental=false",
+    "bash tools/codestyle/format.sh",
+    "uv run --script tools/codestyle/copyright_fixer.py .",
 ]
 
 ["format-check"]
-description = "Check ruff formatting, lint rules, and SPDX copyright headers without modifying files. Safe for CI and local pre-review checks."
+description = "Check TOML and Python formatting, lint rules, and SPDX copyright headers without modifying source files. Safe for CI and local pre-review checks."
 run = [
-  "bash tools/codestyle/format.sh --check",
-  "bash tools/codestyle/ruff_check.sh",
-  "uv run --script tools/codestyle/copyright_fixer.py --check .",
+    "dprint check --incremental=false",
+    "bash tools/codestyle/format.sh --check",
+    "bash tools/codestyle/ruff_check.sh",
+    "uv run --script tools/codestyle/copyright_fixer.py --check .",
 ]
 
 [typecheck]
 description = "Run ty type checks through the repo wrapper. Requires Python dependencies installed; does not modify files."
 run = "bash tools/codestyle/typecheck.sh"
 
 ["lock-check"]
-description = "Verify uv.lock matches pyproject.toml. Runs uv lock, then fails if uv.lock has a diff."
+description = "Verify generated CUDA metadata and uv.lock match their sources."
 run = [
-  "uv lock",
-  "git diff --exit-code uv.lock",
+    "uv run --offline --frozen tools/gen_cuda_deps.py cuda_deps.toml --pyproject pyproject.toml --check",
+    "uv lock",
+    "git diff --exit-code uv.lock",
 ]
 
 [check]
@@ -36,8 +39,8 @@ run = "echo 'quality checks passed'"
 description = "Run the standard pre-PR validation suite: check, lock-check, then CI unit tests."
 alias = "ci"
 run = [
-  { task = "check" },
-  { task = "lock-check" },
-  { task = "test:ci" },
-  "echo 'validation passed'",
+    { task = "check" },
+    { task = "lock-check" },
+    { task = "test:ci" },
+    "echo 'validation passed'",
 ]

.mise/tasks/setup.toml

@@ -4,10 +4,10 @@
 [setup]
 description = "Install pinned mise-managed tools and create the project virtualenv. Modifies the mise tool cache and .venv or UV_PROJECT_ENVIRONMENT."
 run = [
-  "MISE_YES=1 mise trust",
-  "MISE_YES=1 mise install",
-  { task = "venv" },
-  "echo 'tools and virtual environment installed successfully via mise'",
+    "MISE_YES=1 mise trust",
+    "MISE_YES=1 mise install",
+    { task = "venv" },
+    "echo 'tools and virtual environment installed successfully via mise'",
 ]
 
 ["bootstrap-tools"]

.mise/tasks/tests.toml

@@ -29,12 +29,12 @@ run = 'uv run --frozen pytest -n auto --dist loadscope -vv -m "smoke and not req
 description = "Run all GPU smoke test stages sequentially. Requires CUDA and GPU dependencies; use before GPU-impacting changes."
 alias = "test-smoke-gpu"
 run = [
-  { task = "test:smoke:gpu:train-only" },
-  { task = "test:smoke:gpu:generation" },
-  { task = "test:smoke:gpu:resume" },
-  { task = "test:smoke:gpu:structured-generation" },
-  { task = "test:smoke:gpu:timeseries" },
-  { task = "test:smoke:gpu:smollm2" },
+    { task = "test:smoke:gpu:train-only" },
+    { task = "test:smoke:gpu:generation" },
+    { task = "test:smoke:gpu:resume" },
+    { task = "test:smoke:gpu:structured-generation" },
+    { task = "test:smoke:gpu:timeseries" },
+    { task = "test:smoke:gpu:smollm2" },
 ]
 
 ["test:smoke:gpu:train-only"]
@@ -71,9 +71,9 @@ run = 'uv run --frozen pytest --dist loadscope -vv -n 0 tests/smoke/ -m "require
 description = "Run GPU integration e2e tests for default and DP configs. Requires CUDA; stages are sequential to control GPU load."
 alias = "test-gpu-integration"
 run = [
-  { task = "bootstrap-nss", args = ["cu129"] },
-  "uv run --frozen pytest -n 0 --dist loadscope -vv tests/e2e/test_safe_synthesizer.py -k default",
-  "uv run --frozen pytest -n 0 --dist loadscope -vv tests/e2e/test_safe_synthesizer.py -k dp",
+    { task = "bootstrap-nss", args = ["cu129"] },
+    "uv run --frozen pytest -n 0 --dist loadscope -vv tests/e2e/test_safe_synthesizer.py -k default",
+    "uv run --frozen pytest -n 0 --dist loadscope -vv tests/e2e/test_safe_synthesizer.py -k dp",
 ]
 
 ["test:e2e"]
@@ -85,24 +85,24 @@ run = [{ task = "test:e2e:default" }, { task = "test:e2e:dp" }]
 description = "Run default e2e tests (requires CUDA)"
 alias = "test-e2e-default"
 run = [
-  { task = "bootstrap-nss", args = ["cu129"] },
-  "uv run --frozen pytest --dist loadscope -vv -n 0 tests/e2e/test_safe_synthesizer.py -k default"
+    { task = "bootstrap-nss", args = ["cu129"] },
+    "uv run --frozen pytest --dist loadscope -vv -n 0 tests/e2e/test_safe_synthesizer.py -k default",
 ]
 
 ["test:e2e:dp"]
 description = "Run dp e2e tests (requires CUDA)"
 alias = "test-e2e-dp"
 run = [
-  { task = "bootstrap-nss", args = ["cu129"] },
-  "uv run --frozen pytest --dist loadscope -vv -n 0 tests/e2e/test_safe_synthesizer.py -k dp"
+    { task = "bootstrap-nss", args = ["cu129"] },
+    "uv run --frozen pytest --dist loadscope -vv -n 0 tests/e2e/test_safe_synthesizer.py -k dp",
 ]
 
 ["test:e2e:collect"]
 description = "Dry-run e2e and GPU selectors with pytest collection only. Does not run tests; useful before changing task selectors."
 alias = "test-e2e-collect"
 run = [
-  { task = "bootstrap-nss", args = ["cu129"] },
-  '''
+    { task = "bootstrap-nss", args = ["cu129"] },
+    '''
 #!/usr/bin/env bash
 set -euo pipefail
 
@@ -114,7 +114,7 @@ echo "--- test-gpu-integration (default) ---"
 uv run --frozen pytest -n auto --dist loadscope -vv tests/e2e/test_safe_synthesizer.py -k default -o "addopts=" --collect-only -qq -p no:warnings 2>/dev/null || true
 echo "--- test-gpu-integration (dp) ---"
 uv run --frozen pytest -n auto --dist loadscope -vv tests/e2e/test_safe_synthesizer.py -k dp -o "addopts=" --collect-only -qq -p no:warnings 2>/dev/null || true
-'''
+''',
 ]
 
 ["test:nss-config-dataset"]
@@ -129,8 +129,8 @@ arg "<dataset>" {
 }
 '''
 run = [
-  { task = "bootstrap-nss", args = ["cu129"] },
-  '''
+    { task = "bootstrap-nss", args = ["cu129"] },
+    '''
 #!/usr/bin/env bash
 set -euo pipefail
 

AGENTS.md

@@ -29,6 +29,8 @@ uv sync --frozen --extra cu129 --extra engine --group dev
 
 Bare `uv sync --frozen` (without extras) installs an incomplete environment -- `ty`, import checks, and GPU tests will fail.
 
+The CPU/CUDA optional-dependency and `[tool.uv.sources]`/`[[tool.uv.index]]` sections of `pyproject.toml` are generated from `cuda_deps.toml` by `tools/gen_cuda_deps.py` -- never hand-edit the `# >>> BEGIN GENERATED ... <<<` blocks in `pyproject.toml`. To add or change a CUDA/CPU dependency, edit `cuda_deps.toml`, then run `uv run --frozen tools/gen_cuda_deps.py cuda_deps.toml --pyproject pyproject.toml` followed by `uv lock`. `mise run lock-check` verifies both are in sync with `cuda_deps.toml`.
+
 Feature branches off `main`. Branch names often include an issue number prefix (e.g., `<author>/123-short-name`).
 
 Do not commit unless the user asks for a commit or PR work. When committing, all commits require DCO sign-off and GPG signing. Always use `git commit --signoff --gpg-sign` (or `-s -S`) -- never write the `Signed-off-by` trailer manually, and never pass `--no-gpg-sign`.

CONTRIBUTING.md

@@ -30,7 +30,7 @@ Please read our [Code of Conduct](CODE_OF_CONDUCT.md) before contributing.
 - Python 3.11–3.13 (project supports Python 3.11, 3.12, and 3.13; `.python-version` pins 3.13 for bootstrapping at the repo root. Python 3.14+ is not supported — see [Troubleshooting](docs/user-guide/troubleshooting.md#python-314-is-not-supported))
 - Git 2.34+ (minimum required for SSH commit signing)
 
-> Note: Other tools like [uv](https://docs.astral.sh/uv/), [ruff](https://docs.astral.sh/ruff/), [ty](https://github.com/astral-sh/ty), and [gh](https://cli.github.com/) are installed automatically by `make setup` (via [mise](https://mise.jdx.dev/)). Tool versions are declared in `.mise.toml` and locked in `mise.lock` (committed), ensuring reproducible toolchains across developer systems and CI. These should not interfere with locally installed tools.
+> Note: Other tools like [uv](https://docs.astral.sh/uv/), [dprint](https://dprint.dev/), [ruff](https://docs.astral.sh/ruff/), [ty](https://github.com/astral-sh/ty), and [gh](https://cli.github.com/) are installed automatically by `make setup` (via [mise](https://mise.jdx.dev/)). Tool versions are declared in `.mise.toml` and locked in `mise.lock` (committed), ensuring reproducible toolchains across developer systems and CI. These should not interfere with locally installed tools.
 
 > Note on mise itself: the mise version is pinned in `.mise.toml` (`min_version`). The first run of `make setup` installs exactly that version via `tools/install-mise.sh`, preferring the GPG-verified installer when the full toolchain (`gpg`, `gpg-agent`, and `dirmngr`) is available and falling back to `https://mise.run` otherwise (with a warning). If you already have a different mise version on `PATH`, `make setup` will stop and tell you -- either run `mise self-update <pinned>` or uninstall the existing mise and rerun. It will not silently replace your install.
 
@@ -557,14 +557,14 @@ Although the default development/runtime interpreter is Python 3.13, source code
 Use mise tasks instead of running `ruff` or `ty` directly. The tasks use pinned tool versions from `.mise.[toml|lock]` (installed via `make setup`) and check all tracked files.
 
 ```bash
-mise run format   # auto-fix: ruff format + import sorting + copyright headers
+mise run format   # auto-fix: dprint TOML + ruff format/import sorting + copyright headers
 mise run check    # read-only local quality checks (format + lint + typecheck + copyright)
 mise run test     # unit tests
 # or just
 mise run format && mise run check && mise run test
 ```
 
-We use `ruff` and `ty` for the majority of this work, wrapped with settings for consistency.
+We use `dprint` for TOML, `ruff` for Python formatting and linting, and `ty` for type checking, wrapped with settings for consistency.
 
 CI calls the same tools through atomic read-only mise tasks. Declarative tasks live in `.mise/tasks/*.toml`; bash-heavy tasks are executable file tasks under `.mise/tasks/`. Shared shell helpers live in `.mise/tasks/_lib.sh`, which is sourced by file tasks but is not executable and does not appear in `mise tasks`. `mise run check` replicates format-check + typecheck locally; `mise run validate` runs the broader pre-PR graph (`check`, `lock-check`, and `test:ci`). Pre-commit hooks (`pre-commit install`) provide faster feedback by checking only staged files, but are not a substitute for the mise tasks.
 
@@ -590,10 +590,11 @@ All mise tasks check the entire project. Pre-commit scopes checks to staged file
 
 | Check | CI task | `mise run format` / `mise run check` | Pre-commit |
 |---|---|---|---|
+| dprint TOML format | `mise run format-check` | `format`: auto-fix; `check`: read-only | not run |
 | ruff format + lint | `mise run format-check` | `format`: auto-fix; `check`: read-only | staged files (auto-fix) |
 | ty typecheck | `mise run typecheck` | read-only | all files |
 | copyright headers | `mise run format-check` | `format`: auto-fix; `check`: read-only | staged files (auto-fix) |
-| uv lock drift | `mise run lock-check` | not checked | on `pyproject.toml` changes |
+| generated CUDA metadata and uv lock drift | `mise run lock-check` | not checked | on `pyproject.toml` or `cuda_deps.toml` changes |
 | DCO signoff | branch protection | not checked | commit-msg hook |
 
 ## Documentation

README.md

@@ -40,7 +40,7 @@ make setup # installs pinned mise, pinned tools from mise.lock, and .venv
 mise run bootstrap-nss cuda
 ```
 
-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).
+Development tools (`dprint`, `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).
 
 Project commands run through mise tasks under `.mise/tasks/`: `*.toml` files for declarative tasks, executable scripts for bash-heavy logic.
 

STYLE_GUIDE.md

@@ -825,6 +825,8 @@ readonly OUTPUT_DIR="${1:?Usage: $0 <output-dir>}"
 
 ### TOML
 
+- Format TOML with `mise run format`; dprint configuration lives in `dprint.json`.
+- Use four spaces for multiline arrays and keep inline-table spacing formatter-controlled.
 - Spaces around `=` for key-value pairs
 - Comments: `# comment` with inline comments for dependency pins
 - Section ordering in `pyproject.toml`: `[project]`, `[dependency-groups]`, `[project.optional-dependencies]`, `[tool.uv]`, `[build-system]`, `[tool.*]`