ci: test GPU workflows against unpinned wheel

.github/actions/setup-gpu-test-env/action.yml

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 name: "Setup GPU Test Environment"
-description: "Common setup for GPU test workflows: mise, Python, CUDA dependencies, and GPU availability"
+description: "Build and install the wheel with unpinned GPU dependencies, test tools, and GPU validation"
 
 inputs:
   python-version:
@@ -25,7 +25,7 @@ inputs:
     required: false
     default: "true"
   cuda-extra:
-    description: "CUDA dependency extra to bootstrap"
+    description: "CUDA dependency extra to install from the built wheel"
     required: false
     default: "cu129"
 
@@ -39,7 +39,6 @@ runs:
         cache-dependency-glob: |
           .python-version
           pyproject.toml
-          uv.lock
 
     - name: Setup Python environment
       uses: ./.github/actions/setup-python-env
@@ -48,11 +47,79 @@ runs:
         python-version: ${{ inputs.python-version }}
         bootstrap-tools: ${{ inputs.bootstrap-tools }}
 
-    - name: Bootstrap CUDA environment
+    - name: Build wheel
       shell: bash
-      run: mise run bootstrap-nss ${{ inputs.cuda-extra }}
+      run: mise run build-wheel
 
-    - name: Check GPU availability
+    - name: Install wheel with unpinned GPU dependencies
       shell: bash
+      env:
+        CUDA_EXTRA: ${{ inputs.cuda-extra }}
       run: |
-        uv run python -c "import torch; print('cuda available:', torch.cuda.is_available()); print('device count:', torch.cuda.device_count())"
+        set -euo pipefail
+        shopt -s nullglob
+        wheels=("${GITHUB_WORKSPACE}"/dist/*.whl)
+        if (( ${#wheels[@]} != 1 )); then
+          echo "Error: expected exactly one wheel under ${GITHUB_WORKSPACE}/dist, found ${#wheels[@]}" >&2
+          exit 1
+        fi
+
+        wheel="${wheels[0]}"
+        venv_path="${UV_PROJECT_ENVIRONMENT:-${GITHUB_WORKSPACE}/.venv}"
+        uv --no-config venv --clear --python "${PYTHON_VERSION}" "${venv_path}"
+        python_bin="${venv_path}/bin/python"
+
+        index_args=()
+        index_count=0
+        while IFS= read -r url; do
+          [[ -n "${url}" ]] || continue
+          index_args+=(--index "${url}")
+          index_count=$((index_count + 1))
+          echo "CUDA index: ${url}"
+        done < <("${python_bin}" - "${GITHUB_WORKSPACE}/pyproject.toml" "${CUDA_EXTRA}" <<'PY'
+        import sys
+        import tomllib
+
+        with open(sys.argv[1], "rb") as handle:
+            indexes = tomllib.load(handle)["tool"]["uv"]["index"]
+
+        cuda_extra = sys.argv[2]
+        print(
+            "\n".join(
+                index["url"]
+                for index in indexes
+                if index["name"].endswith(f"-{cuda_extra}")
+                or f"/{cuda_extra}" in index["url"]
+            )
+        )
+        PY
+        )
+
+        if (( index_count < 3 )); then
+          echo "Error: expected 3+ ${CUDA_EXTRA} indexes in pyproject.toml, found ${index_count}" >&2
+          exit 1
+        fi
+
+        uv --no-config pip install \
+          --python "${python_bin}" \
+          --no-sources \
+          --default-index https://pypi.org/simple \
+          "${index_args[@]}" \
+          --index-strategy unsafe-best-match \
+          "nemo-safe-synthesizer[${CUDA_EXTRA},engine] @ file://${wheel}" \
+          "pytest>=9.0.3" \
+          "pytest-asyncio>=0.24.0" \
+          "pytest-cov>=6.0.0" \
+          "pytest-env>=1.1.0" \
+          "pytest-subtests>=0.13.1" \
+          "pytest-timeout>=2.3.1" \
+          "pytest-xdist>=3.5.0"
+
+        uv --no-config pip check --python "${python_bin}"
+        "${python_bin}" -c \
+          "from importlib.metadata import version; from nemo_safe_synthesizer.package_info import __version__; assert __version__ == version('nemo-safe-synthesizer')"
+        "${python_bin}" -c \
+          "import torch; print('cuda available:', torch.cuda.is_available()); print('device count:', torch.cuda.device_count()); assert torch.cuda.is_available()"
+
+        echo "UV_PROJECT_ENVIRONMENT=${venv_path}" >> "${GITHUB_ENV}"
+        echo "UV_NO_SYNC=1" >> "${GITHUB_ENV}"

.github/workflows/README.md

@@ -145,13 +145,13 @@ All jobs run on `ubuntu-latest` (GitHub-hosted).
 
 The `gpu-tests.yml` workflow runs nightly at 02:00 UTC, and can also be triggered manually via `workflow_dispatch`. Manual dispatch includes a `suite` dropdown with `all`, `smoke`, and `e2e` options. The `push` trigger for `pull-request/*` branches is currently commented out due to internal blockers, so PRs do not automatically produce GPU status checks. We expect to re-enable that path as soon as those blockers are resolved. There are several key jobs:
 
-- GPU Smoke Tests: staged smoke tests on a gpu runner with a 30-minute job timeout. The train-only, generation, resume, structured generation, timeseries, and SmolLM2 lanes run as separate workflow steps. Required for merge when the workflow is part of branch protection.
-- GPU E2E Tests: End-to-end tests on a gpu runner with a 60-minute job timeout and 45-minute step timeout. Informational -- failures produce a warning but don't block merge.
+- GPU Smoke Tests: runs GPU-marked unit tests, followed by staged train-only, generation, resume, structured generation, timeseries, and SmolLM2 smoke tests. Required for merge when the workflow is part of branch protection.
+- GPU E2E Tests: End-to-end tests on a gpu runner with a 210-minute job timeout and 190-minute step timeout. Informational -- failures produce a warning but don't block merge.
 - GPU CI Status: Aggregation job for the GPU workflow. It is not currently a live branch-protection requirement while PR GPU runs are disabled; when re-enabled, it is intended to be the required GPU check. It fails if smoke tests fail and warns if E2E tests fail.
 
 The `changes` (Detect Changes) job is skipped on `workflow_dispatch`. GPU jobs use `always()` in their job conditions so manual runs can bypass the skipped dependency and run the selected suite. On scheduled runs, `changes` gates GPU jobs with the `src_test_deps` output, which is true for source, test, `pytest.ini`, dependency, or CI workflow/action changes.
 
-GPU jobs use `.github/actions/setup-gpu-test-env` for shared GPU setup: enabling the `uv` cache, setting up Python from `.python-version`, bootstrapping CUDA dependencies with mise, and checking GPU availability.
+GPU jobs use `.github/actions/setup-gpu-test-env` for shared GPU setup. Following the release wheel verification's clean-room pattern, the action builds the wheel and installs `[cu129,engine]` plus test tooling with uv configuration and project sources disabled. Dependencies resolve from PyPI and the required CUDA wheel indexes without consulting `uv.lock` before GPU availability is checked.
 
 To trigger manually from the CLI (produces a run but not a PR status check):
 

.github/workflows/gpu-tests.yml

@@ -85,7 +85,7 @@ jobs:
           inputs.suite == 'smoke'
         )
       }}
-    timeout-minutes: 30
+    timeout-minutes: 60
     runs-on: linux-amd64-gpu-a100-latest-1
     strategy:
       fail-fast: false
@@ -103,6 +103,10 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
+      - name: Run GPU unit tests
+        timeout-minutes: 30
+        run: mise run test:unit:gpu
+
       - name: Run GPU smoke tests - train only
         timeout-minutes: 10
         run: mise run test:smoke:gpu:train-only
@@ -158,7 +162,7 @@ jobs:
         uses: ./.github/actions/setup-gpu-test-env
 
       - name: Run GPU E2E tests
-        run: mise run test:e2e
+        run: mise run test:e2e:prepared
         timeout-minutes: 190
 
   # ---------------------------------------------------------------------------

.mise/tasks/tests.toml

@@ -10,6 +10,11 @@ description = "Run all unit tests including slow tests, excluding e2e and smoke.
 alias = "test-unit-slow"
 run = 'uv run --frozen pytest -n auto --dist loadscope -vv tests -m "unit"'
 
+["test:unit:gpu"]
+description = "Run non-slow GPU-marked unit tests outside the smoke and e2e suites. Requires CUDA and GPU dependencies."
+alias = "test-unit-gpu"
+run = 'uv run --frozen pytest --dist loadscope -vv -n 0 tests -m "requires_gpu and not slow and not smoke and not e2e"'
+
 ["test:ci"]
 description = "Run CI unit tests with coverage, excluding slow, e2e, GPU, and smoke tests. Produces coverage.json."
 alias = "test-ci"
@@ -81,22 +86,34 @@ description = "Run all GPU e2e tests: default then DP. Requires CUDA; stages are
 alias = "test-e2e"
 run = [{ task = "test:e2e:default" }, { task = "test:e2e:dp" }]
 
+["test:e2e:prepared"]
+description = "Run all GPU e2e tests without bootstrapping dependencies. Requires a prepared CUDA environment."
+run = [{ task = "test:e2e:default:run" }, { task = "test:e2e:dp:run" }]
+
 ["test:e2e:default"]
 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 = "test:e2e:default:run" },
 ]
 
+["test:e2e:default:run"]
+hide = true
+run = "uv run --frozen --no-sync 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 = "test:e2e:dp:run" },
 ]
 
+["test:e2e:dp:run"]
+hide = true
+run = "uv run --frozen --no-sync 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"

tests/TESTING.md

@@ -19,6 +19,7 @@ All mise test tasks, grouped by scope:
 ```bash
 mise run test                              # Unit (excludes slow, e2e, and smoke)
 mise run test:unit-slow                    # Unit tests including slow (excludes e2e and smoke)
+mise run test:unit:gpu                     # Non-slow GPU-marked unit tests outside smoke and e2e
 mise run test:smoke                        # CPU smoke tests (~few min, no GPU required)
 mise run test:smoke:gpu                    # All staged GPU smoke tests (requires CUDA)
 mise run test:smoke:gpu:train-only
@@ -28,6 +29,7 @@ mise run test:smoke:gpu:structured-generation
 mise run test:smoke:gpu:timeseries
 mise run test:smoke:gpu:smollm2
 mise run test:e2e                          # All e2e (requires CUDA) -- runs default + dp
+mise run test:e2e:prepared                 # All e2e without dependency bootstrap
 mise run test:e2e:default                  # e2e default (no-DP) tests only
 mise run test:e2e:dp                       # e2e DP tests only
 mise run test:ci                           # CI unit tests with coverage (excludes slow, e2e, gpu, smoke)