Automate Puzzletron GPU lifecycle tests

.github/workflows/puzzletron_gpu_tests.yml

@@ -0,0 +1,101 @@
+name: Puzzletron GPU tests
+
+"on":
+  push:
+    branches: ["pull-request/[0-9]+"]
+  schedule:
+    - cron: "30 1 * * *"
+  workflow_dispatch:
+    # On-demand
+
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ startsWith(github.ref, 'refs/heads/pull-request/') && github.ref || github.sha }}
+  cancel-in-progress: true
+
+jobs:
+  pr-gate:
+    uses: ./.github/workflows/_pr_gate.yml
+    permissions:
+      checks: read
+    with:
+      files: |
+        .github/workflows/_pr_gate.yml
+        .github/workflows/puzzletron_gpu_tests.yml
+        .github/actions/cache-extensions/**
+        examples/puzzletron/**/*.py
+        examples/puzzletron/**/*.sh
+        examples/puzzletron/**/*.yaml
+        examples/puzzletron/ci/**
+        examples/puzzletron/ci_environment.json
+        examples/puzzletron/requirements.txt
+        modelopt/torch/puzzletron/**
+        noxfile.py
+        puzzletron_orchestrator/**
+        puzzletron_setup/**
+        pyproject.toml
+        tests/conftest.py
+        tests/_test_utils/torch/puzzletron/**
+        tests/_test_utils/torch/transformers_models.py
+        tests/gpu/torch/puzzletron/**
+
+  resolve-image:
+    needs: [pr-gate]
+    if: needs.pr-gate.outputs.run_tests == 'true'
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+    outputs:
+      image: ${{ steps.image.outputs.image }}
+      cache_key: ${{ steps.image.outputs.cache_key }}
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          persist-credentials: false
+      - name: Resolve immutable Puzzletron image
+        id: image
+        env:
+          PUZZLETRON_GPU_CI_IMAGE: ${{ vars.PUZZLETRON_GPU_CI_IMAGE }}
+        run: python examples/puzzletron/ci/resolve_ci_image.py >> "${GITHUB_OUTPUT}"
+
+  gpu-puzzletron:
+    needs: [resolve-image]
+    runs-on: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
+    timeout-minutes: 50
+    container:
+      image: ${{ needs.resolve-image.outputs.image }}
+      options: --shm-size=16gb
+      env:
+        GIT_DEPTH: 1000
+        PIP_CONSTRAINT: ""
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          persist-credentials: false
+      - uses: nv-gha-runners/setup-proxy-cache@main
+      - uses: ./.github/actions/cache-extensions
+        with:
+          cache-key: rtxpro6000-puzzletron-${{ needs.resolve-image.outputs.cache_key }}
+      - name: Run the Puzzletron lifecycle gate
+        run: nox -s gpu_puzzletron
+
+  gpu-puzzletron-required-check:
+    if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }}
+    needs: [pr-gate, resolve-image, gpu-puzzletron]
+    runs-on: ubuntu-latest
+    steps:
+      - name: Report intentionally scoped Puzzletron GPU tests
+        if: needs.pr-gate.outputs.run_tests != 'true'
+        run: |
+          echo "## Puzzletron GPU tests were not required" >> "${GITHUB_STEP_SUMMARY}"
+          echo >> "${GITHUB_STEP_SUMMARY}"
+          echo "No Puzzletron lifecycle path changed in this pull request." >> "${GITHUB_STEP_SUMMARY}"
+      - name: Required Puzzletron GPU tests did not succeed
+        if: >-
+          ${{ needs.pr-gate.result != 'success' ||
+            (needs.pr-gate.outputs.run_tests == 'true' &&
+            (needs.resolve-image.result != 'success' ||
+            needs.gpu-puzzletron.result != 'success')) }}
+        run: exit 1

.github/workflows/unit_tests.yml

@@ -7,6 +7,7 @@ on:
     branches: [main, release/*, feature/*]
     paths:
       - ".github/workflows/unit_tests.yml"
+      - ".github/workflows/puzzletron_gpu_tests.yml"
       - "examples/puzzletron/**"
       - "modelopt/**"
       - "noxfile.py"
@@ -83,6 +84,7 @@ jobs:
         with:
           files: |
             .github/workflows/unit_tests.yml
+            .github/workflows/puzzletron_gpu_tests.yml
             examples/puzzletron/**/*.html
             examples/puzzletron/**/*.json
             examples/puzzletron/**/*.md

examples/puzzletron/ci/Dockerfile

@@ -0,0 +1,31 @@
+FROM nvidia/cuda:12.9.2-cudnn-devel-ubuntu24.04@sha256:b4db213759eb86d55a7271909bdad891fab300ec5700fb4f4656463b2f51980f
+
+SHELL ["/bin/bash", "-o", "pipefail", "-c"]
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+ENV VIRTUAL_ENV=/venv
+ENV PATH=/venv/bin/:$PATH
+ENV PIP_NO_CACHE_DIR=1
+ENV PUZZLETRON_CI_ENVIRONMENT=/opt/puzzletron/ci_environment.json
+ENV PUZZLETRON_REQUIREMENTS=/opt/puzzletron/requirements.txt
+ENV PYTHONUNBUFFERED=1
+
+COPY examples/puzzletron/ci_environment.json /opt/puzzletron/ci_environment.json
+COPY examples/puzzletron/requirements.txt /opt/puzzletron/requirements.txt
+COPY examples/puzzletron/ci/setup_env.sh /opt/puzzletron/setup_env.sh
+COPY pyproject.toml LICENSE_HEADER /opt/modelopt-dependencies/
+
+RUN bash /opt/puzzletron/setup_env.sh --deps
+RUN mkdir -p /opt/modelopt-dependencies/modelopt && \
+    touch /opt/modelopt-dependencies/modelopt/__init__.py && \
+    python -m pip install "/opt/modelopt-dependencies[hf,puzzletron,dev-test]" && \
+    python -m pip uninstall -y nvidia-modelopt && \
+    python -m pip check && \
+    bash /opt/puzzletron/setup_env.sh --verify
+
+LABEL org.opencontainers.image.source="https://github.com/NVIDIA/Model-Optimizer" \
+    com.nvidia.modelopt.puzzletron.environment="examples/puzzletron/ci_environment.json" \
+    com.nvidia.modelopt.puzzletron.environment-recipe="examples/puzzletron/ci/setup_env.sh"
+
+WORKDIR /workspace/modelopt

examples/puzzletron/ci/README.md

@@ -0,0 +1,41 @@
+# Puzzletron GPU lifecycle CI image
+
+This directory defines the execution image for the hermetic one-GPU Puzzletron
+lifecycle gate. The image supplies the pinned CUDA and Python dependencies;
+the `gpu_puzzletron` Nox session installs the checked-out ModelOpt revision and
+runs the sole reusable tiny-Qwen lifecycle test.
+
+The base image is pinned by OCI digest in both
+[`Dockerfile`](Dockerfile) and [`ci_environment.json`](../ci_environment.json).
+The environment file also owns the exact Torch, Transformers, LMMS-Eval,
+AutoModel, AIPerf, and Nox versions. The image preinstalls ModelOpt's declared
+runtime and test dependencies. The setup script verifies the pins when the
+image is built and again after installing only the checked-out ModelOpt source,
+without resolving dependencies or using build isolation at runtime.
+
+Build the image from the repository root:
+
+```bash
+docker build \
+  --platform linux/amd64 \
+  --file examples/puzzletron/ci/Dockerfile \
+  --tag modelopt-puzzletron-ci:local \
+  .
+```
+
+Run the same entry point used by CI with one visible GPU:
+
+```bash
+docker run --gpus device=0 --ipc=host --rm \
+  --volume "${PWD}:/workspace/modelopt" \
+  --workdir /workspace/modelopt \
+  modelopt-puzzletron-ci:local \
+  nox -s gpu_puzzletron
+```
+
+Publishing the reviewed image is a separate registry operation. Publish it with
+anonymous pull access, then configure `PUZZLETRON_GPU_CI_IMAGE` with the
+complete immutable `nvcr.io/...@sha256:...` reference. The workflow
+intentionally does not expose registry credentials to copied pull-request
+branches, and rejects tags and non-NVCR references before allocating a GPU
+runner.

examples/puzzletron/ci/resolve_ci_image.py

@@ -0,0 +1,79 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Validate and resolve the immutable image used by Puzzletron GPU CI."""
+
+import json
+import os
+import re
+import sys
+from pathlib import Path
+
+__all__ = ["resolve_image_reference", "validate_repository_contract"]
+
+_NVCR_IMAGE = re.compile(
+    r"nvcr\.io/(?:[A-Za-z0-9._-]+/)*[A-Za-z0-9._-]+@sha256:(?P<digest>[0-9a-f]{64})"
+)
+_CUDA_BASE_IMAGE = re.compile(r"nvidia/cuda:[A-Za-z0-9._-]+@sha256:[0-9a-f]{64}")
+
+
+def resolve_image_reference(image: str) -> tuple[str, str]:
+    """Return an immutable nvcr.io image and its digest cache key."""
+    match = _NVCR_IMAGE.fullmatch(image)
+    if match is None:
+        raise ValueError("PUZZLETRON_GPU_CI_IMAGE must be an immutable nvcr.io digest")
+    return image, match.group("digest")
+
+
+def validate_repository_contract(repository_root: Path) -> None:
+    """Verify the checked-out image recipe agrees with its recorded environment."""
+    ci_root = repository_root / "examples/puzzletron"
+    environment = json.loads((ci_root / "ci_environment.json").read_text())
+    dockerfile = (ci_root / "ci/Dockerfile").read_text()
+    base_image = environment["gpu_image"]["base_image"]
+
+    if _CUDA_BASE_IMAGE.fullmatch(base_image) is None:
+        raise ValueError("gpu_image.base_image must use a full lowercase SHA-256 digest")
+
+    required_lines = (
+        f"FROM {base_image}",
+        "ENV PUZZLETRON_CI_ENVIRONMENT=/opt/puzzletron/ci_environment.json",
+        "ENV PUZZLETRON_REQUIREMENTS=/opt/puzzletron/requirements.txt",
+        "COPY pyproject.toml LICENSE_HEADER /opt/modelopt-dependencies/",
+        "RUN bash /opt/puzzletron/setup_env.sh --deps",
+        'python -m pip install "/opt/modelopt-dependencies[hf,puzzletron,dev-test]"',
+        "bash /opt/puzzletron/setup_env.sh --verify",
+    )
+    missing = [line for line in required_lines if line not in dockerfile]
+    if missing:
+        raise ValueError(f"Dockerfile is missing recorded contract lines: {missing}")
+
+
+def main() -> int:
+    """Write validated values in GitHub output format."""
+    try:
+        validate_repository_contract(Path.cwd())
+        image, cache_key = resolve_image_reference(os.environ.get("PUZZLETRON_GPU_CI_IMAGE", ""))
+    except (KeyError, OSError, ValueError, json.JSONDecodeError) as error:
+        print(f"::error::{error}", file=sys.stderr)
+        return 1
+
+    print(f"image={image}")
+    print(f"cache_key={cache_key}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())