pytest-gpu-proof 0.1.0__py3-none-any.whl

pytest_gpu_proof/plugin.py

@@ -0,0 +1,331 @@
+"""
+pytest plugin entry point.
+
+Registers CLI options, markers, and the gpu_proof_check fixture.
+At session end, builds a signed receipt for every test that used the fixture.
+"""
+
+import datetime
+import warnings
+from typing import Any, Dict, List
+
+import pytest
+
+from .compare import run_comparison
+from .config import GpuProofConfig, load_config
+
+
+def _utcnow() -> str:
+    return datetime.datetime.now(datetime.UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
+
+
+def _has_gpu() -> bool:
+    import subprocess
+    try:
+        r = subprocess.run(["nvidia-smi"], capture_output=True, timeout=5)
+        if r.returncode == 0:
+            return True
+    except (FileNotFoundError, subprocess.TimeoutExpired):
+        pass
+    try:
+        import torch
+        return torch.cuda.is_available()
+    except ImportError:
+        pass
+    return False
+
+
+class GpuProofPlugin:
+    def __init__(self, pytest_config):
+        self.gpu_proof_config: GpuProofConfig = load_config(pytest_config)
+        self.test_results: List[Dict[str, Any]] = []
+        self.skipped_required: List[str] = []
+        self.started_at: str = ""
+
+    # ------------------------------------------------------------------
+    # session lifecycle
+    # ------------------------------------------------------------------
+
+    def pytest_sessionstart(self, session):
+        self.started_at = _utcnow()
+
+    def pytest_sessionfinish(self, session, exitstatus):
+        if not self.gpu_proof_config.enabled:
+            return
+
+        skipped_marked = [
+            t["node_id"] for t in self.test_results if t.get("outcome") == "skipped"
+        ]
+        if self.gpu_proof_config.fail_on_skip and (skipped_marked or self.skipped_required):
+            names = sorted(set(skipped_marked + self.skipped_required))
+            print(
+                "\n[gpu-proof] --gpu-proof-fail-on-skip: "
+                f"{len(names)} marked test(s) were skipped:\n"
+                + "".join(f"            - {n}\n" for n in names)
+                + "            No receipt was written; session marked as failed."
+            )
+            session.exitstatus = 1
+            return
+
+        if not self.test_results:
+            print(
+                "\n[gpu-proof] --gpu-proof-enable is set but no gpu_proof tests were found.\n"
+                "            Add @pytest.mark.gpu_proof to your tests or use the gpu_proof_check fixture.\n"
+                "            No receipt was written.\n"
+                "            Tip: try 'pytest examples/minimal_python_only/test_minimal.py --gpu-proof-enable -v'"
+            )
+            return
+        self._emit_receipt()
+
+    # ------------------------------------------------------------------
+    # result collection
+    # ------------------------------------------------------------------
+
+    def _is_marked(self, item) -> bool:
+        return bool(
+            item.get_closest_marker(self.gpu_proof_config.required_marker)
+            or item.get_closest_marker("gpu_equivalence")
+        )
+
+    @pytest.hookimpl(hookwrapper=True)
+    def pytest_runtest_makereport(self, item, call):
+        outcome = yield
+        report = outcome.get_result()
+
+        if call.when == "setup" and report.skipped:
+            # Skipped tests never reach the "call" phase — record them here so
+            # they are visible in the receipt (and to --gpu-proof-fail-on-skip)
+            # instead of being silently dropped.
+            if item.get_closest_marker("gpu_required") and not self._is_marked(item):
+                self.skipped_required.append(item.nodeid)
+            if self._is_marked(item):
+                self.test_results.append(
+                    {
+                        "node_id": item.nodeid,
+                        "outcome": "skipped",
+                        "duration_s": round(call.duration, 4),
+                        "checks": [],
+                    }
+                )
+            return
+
+        if call.when == "teardown":
+            # The gpu_proof_check fixture records its checks during fixture
+            # teardown, which happens *after* the call-phase report. Patch the
+            # already-recorded entry here so checks are not silently dropped.
+            checks = getattr(item, "_gpu_proof_checks", None)
+            if checks is not None:
+                for t in reversed(self.test_results):
+                    if t["node_id"] == item.nodeid:
+                        t["checks"] = checks
+                        break
+            return
+
+        if call.when != "call":
+            return
+
+        uses_fixture = "gpu_proof_check" in getattr(item, "fixturenames", ())
+        if not uses_fixture and not self._is_marked(item):
+            return
+
+        if report.passed:
+            outcome_str = "passed"
+        elif report.skipped:
+            outcome_str = "skipped"
+        else:
+            outcome_str = "failed"
+
+        self.test_results.append(
+            {
+                "node_id": item.nodeid,
+                "outcome": outcome_str,
+                "duration_s": round(call.duration, 4),
+                "checks": [],  # filled in at teardown once the fixture finalizes
+            }
+        )
+
+    # ------------------------------------------------------------------
+    # receipt emission
+    # ------------------------------------------------------------------
+
+    def _emit_receipt(self):
+        from .receipt import build_receipt_payload, finalize_receipt, write_receipt
+        from .signers.ed25519 import SSHSigner
+        from .signers.base import VerifierError
+
+        ended_at = _utcnow()
+        cfg = self.gpu_proof_config
+
+        if cfg.signing_backend == "none":
+            try:
+                payload = build_receipt_payload(cfg, self.test_results, self.started_at, ended_at)
+                receipt = dict(payload)
+                receipt["signature"] = None
+                write_receipt(receipt, cfg.output)
+                print(f"\n[gpu-proof] Receipt written to {cfg.output}")
+                print(
+                    "[gpu-proof] WARNING: signing backend is 'none' — the receipt is UNSIGNED\n"
+                    "            and will fail verification unless --allow-unsigned is passed."
+                )
+            except Exception as e:
+                warnings.warn(f"[gpu-proof] Failed to write receipt: {e}", stacklevel=1)
+            return
+
+        try:
+            signer = SSHSigner(key_path=cfg.key_path)
+        except VerifierError as e:
+            warnings.warn(f"[gpu-proof] Signing skipped: {e}", stacklevel=1)
+            return
+
+        try:
+            payload = build_receipt_payload(cfg, self.test_results, self.started_at, ended_at)
+            receipt = finalize_receipt(payload, signer)
+            write_receipt(receipt, cfg.output)
+            print(f"\n[gpu-proof] Receipt written to {cfg.output}")
+            print(f"[gpu-proof] Signed with key {signer.key_fingerprint()}")
+        except Exception as e:
+            warnings.warn(f"[gpu-proof] Failed to write receipt: {e}", stacklevel=1)
+
+
+# ------------------------------------------------------------------
+# plugin registration hooks (module-level, always active)
+# ------------------------------------------------------------------
+
+def pytest_addoption(parser):
+    group = parser.getgroup("gpu-proof", "GPU proof receipt generation")
+    group.addoption(
+        "--gpu-proof-enable",
+        action="store_true",
+        default=False,
+        help="Enable GPU proof receipt generation",
+    )
+    group.addoption(
+        "--gpu-proof-mode",
+        default="local",
+        choices=["local", "ci-gpu"],
+        help="Execution mode: local (default) or ci-gpu",
+    )
+    group.addoption(
+        "--gpu-proof-out",
+        default="gpu-proof.json",
+        metavar="PATH",
+        help="Output path for the receipt JSON (default: gpu-proof.json)",
+    )
+    group.addoption(
+        "--gpu-proof-key",
+        default=None,
+        metavar="PATH",
+        help="Path to SSH private key (default: auto-discover from git config or ~/.ssh/)",
+    )
+    group.addoption(
+        "--gpu-proof-signing-backend",
+        default="ed25519",
+        choices=["ed25519", "none"],
+        help="Signing backend (default: ed25519 via SSH key)",
+    )
+    group.addoption(
+        "--gpu-proof-policy",
+        default=None,
+        metavar="PATH",
+        help="Path to verification policy YAML",
+    )
+    group.addoption(
+        "--gpu-proof-required-marker",
+        default="gpu_proof",
+        help="Marker name that flags a test for the receipt (default: gpu_proof)",
+    )
+    group.addoption(
+        "--gpu-proof-fail-on-skip",
+        action="store_true",
+        default=False,
+        help="Fail the session if any gpu_required test is skipped",
+    )
+    group.addoption(
+        "--gpu-proof-fingerprint-paths",
+        default="src,tests",
+        metavar="PATHS",
+        help="Comma-separated paths to fingerprint (default: src,tests)",
+    )
+    group.addoption(
+        "--gpu-proof-github-user",
+        default=None,
+        metavar="USERNAME",
+        help="GitHub username of the signer (default: auto-detect from git remote)",
+    )
+
+
+def pytest_configure(config):
+    config.addinivalue_line("markers", "gpu_proof: include test in GPU proof receipt")
+    config.addinivalue_line("markers", "gpu_required: skip test if no GPU is available")
+    config.addinivalue_line(
+        "markers",
+        "gpu_equivalence: alias for gpu_proof; marks a GPU equivalence check",
+    )
+
+    try:
+        enabled = config.getoption("--gpu-proof-enable")
+    except ValueError:
+        enabled = False
+
+    if enabled:
+        plugin = GpuProofPlugin(config)
+        config.pluginmanager.register(plugin, "gpu-proof-plugin")
+
+
+def pytest_collection_modifyitems(config, items):
+    try:
+        enabled = config.getoption("--gpu-proof-enable")
+    except ValueError:
+        enabled = False
+
+    skip_no_gpu = pytest.mark.skip(reason="No GPU available (gpu_required marker)")
+    has_gpu = None  # lazy
+
+    for item in items:
+        if item.get_closest_marker("gpu_required"):
+            if has_gpu is None:
+                has_gpu = _has_gpu()
+            if not has_gpu:
+                item.add_marker(skip_no_gpu)
+
+
+# ------------------------------------------------------------------
+# fixture
+# ------------------------------------------------------------------
+
+@pytest.fixture
+def gpu_proof_check(request):
+    """
+    Fixture for GPU equivalence checking.
+
+    Usage::
+
+        def test_rnea(gpu_proof_check):
+            gpu_proof_check(
+                name="rnea",
+                reference=python_ref,
+                candidate=cuda_fn,
+                args=(q, qd, qdd),
+                compare=my_compare_fn,   # optional; default: numpy.allclose / ==
+                metadata={"robot": "go2"},
+            )
+    """
+    checks = []
+
+    def check(name, reference, candidate, args=(), kwargs=None, compare=None, metadata=None):
+        outcome, _ref, _cand, error = run_comparison(
+            reference, candidate, args or (), kwargs or {}, compare
+        )
+        checks.append(
+            {
+                "name": name,
+                "outcome": outcome,
+                "metadata": metadata or {},
+            }
+        )
+        if outcome != "passed":
+            pytest.fail(f"gpu_proof_check {name!r}: {error}")
+
+    yield check
+
+    request.node._gpu_proof_checks = checks

pytest_gpu_proof/receipt.py

@@ -0,0 +1,143 @@
+"""
+Build, sign, and write the JSON receipt artifact.
+
+Signing covers the canonical (compact, sorted-key) JSON of the receipt without
+the signature field. The signature is then embedded as receipt["signature"].
+"""
+
+import base64
+import datetime
+import json
+import platform
+import subprocess
+import sys
+from typing import Any, Dict, List, Optional
+
+from .gitutils import (
+    get_branch,
+    get_commit_sha,
+    get_gh_cli_login,
+    get_github_username,
+    get_remote_url,
+    is_dirty,
+)
+
+
+def _utcnow() -> str:
+    return datetime.datetime.now(datetime.UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
+
+
+def _gpu_info() -> Optional[Dict[str, Any]]:
+    try:
+        result = subprocess.run(
+            [
+                "nvidia-smi",
+                "--query-gpu=name,driver_version,memory.total",
+                "--format=csv,noheader",
+            ],
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+        if result.returncode == 0 and result.stdout.strip():
+            parts = [p.strip() for p in result.stdout.strip().split(",")]
+            return {
+                "name": parts[0] if len(parts) > 0 else None,
+                "driver_version": parts[1] if len(parts) > 1 else None,
+                "memory": parts[2] if len(parts) > 2 else None,
+            }
+    except (FileNotFoundError, subprocess.TimeoutExpired):
+        pass
+    return None
+
+
+def _env_info() -> Dict[str, Any]:
+    try:
+        import pytest
+        pytest_version = pytest.__version__
+    except ImportError:
+        pytest_version = None
+
+    from pytest_gpu_proof import __version__
+
+    return {
+        "python_version": sys.version.split()[0],
+        "platform": platform.system().lower(),
+        "pytest_version": pytest_version,
+        "plugin_version": __version__,
+        "gpu_info": _gpu_info(),
+    }
+
+
+def canonicalize(receipt_dict: dict) -> bytes:
+    return json.dumps(receipt_dict, sort_keys=True, separators=(",", ":")).encode()
+
+
+def build_receipt_payload(
+    config,
+    test_results: List[dict],
+    started_at: str,
+    ended_at: str,
+    override_github_username: Optional[str] = None,
+) -> dict:
+    from .fingerprint import compute_fingerprint
+
+    remote_url = get_remote_url()
+    # Signer resolution: explicit override > [tool.gpu_proof]/flag config >
+    # authenticated gh CLI login (the actual keyholder) > origin-remote owner.
+    # The last is only a guess — for org-owned repos it yields the ORG, which
+    # has no SSH keys, so verification would fail; warn when we land there.
+    github_username = override_github_username or config.github_username
+    if not github_username:
+        github_username = get_gh_cli_login()
+    if not github_username:
+        github_username = get_github_username()
+        if github_username:
+            print(
+                f"[gpu-proof] WARNING: signer '{github_username}' was derived "
+                "from the origin remote owner, which may be an org with no SSH "
+                "keys. If verification fails, set --gpu-proof-github-user or "
+                "github_username in [tool.gpu_proof] to the keyholder."
+            )
+    fingerprint = compute_fingerprint(config.fingerprint_paths)
+
+    return {
+        "schema_version": "1",
+        "mode": config.mode,
+        "repo": {
+            "remote_url": remote_url,
+            "github_username": github_username,
+            "commit_sha": get_commit_sha(),
+            "branch": get_branch(),
+            "dirty": is_dirty(),
+        },
+        "fingerprint": fingerprint,
+        "session": {
+            "started_at": started_at,
+            "ended_at": ended_at,
+            "node_ids": [t["node_id"] for t in test_results],
+        },
+        "tests": test_results,
+        "environment": _env_info(),
+    }
+
+
+def finalize_receipt(payload: dict, signer) -> dict:
+    data = canonicalize(payload)
+    sig_bytes = signer.sign(data)
+
+    receipt = dict(payload)
+    receipt["signature"] = {
+        "algorithm": signer.algorithm(),
+        "backend": "ssh-local",
+        "signer": payload["repo"].get("github_username") or "unknown",
+        "key_fingerprint": signer.key_fingerprint(),
+        "value": base64.b64encode(sig_bytes).decode(),
+    }
+    return receipt
+
+
+def write_receipt(receipt: dict, path: str) -> None:
+    with open(path, "w") as f:
+        json.dump(receipt, f, indent=2, sort_keys=True)
+        f.write("\n")

pytest_gpu_proof/signers/base.py

@@ -0,0 +1,20 @@
+from abc import ABC, abstractmethod
+
+
+class SignerBase(ABC):
+    @abstractmethod
+    def sign(self, data: bytes) -> bytes:
+        ...
+
+    @abstractmethod
+    def key_fingerprint(self) -> str:
+        ...
+
+    @abstractmethod
+    def algorithm(self) -> str:
+        """Name of the signature algorithm actually used (e.g. 'ed25519')."""
+        ...
+
+
+class VerifierError(Exception):
+    pass

pytest_gpu_proof/signers/ed25519.py

@@ -0,0 +1,178 @@
+"""
+SSH-key-based signing and GitHub-key-based verification.
+
+Signing uses the developer's existing SSH private key (the same key they use
+to push to GitHub). Verification fetches the signer's public keys from
+https://github.com/{username}.keys — no separate key distribution step needed.
+"""
+
+import base64
+import hashlib
+import os
+from pathlib import Path
+from typing import List, Optional
+from urllib.request import urlopen
+
+from cryptography.exceptions import InvalidSignature
+from cryptography.hazmat.primitives.asymmetric.ec import (
+    ECDSA,
+    EllipticCurvePrivateKey,
+    EllipticCurvePublicKey,
+)
+from cryptography.hazmat.primitives.asymmetric.ed25519 import (
+    Ed25519PrivateKey,
+    Ed25519PublicKey,
+)
+from cryptography.hazmat.primitives.asymmetric.padding import MGF1, PSS
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey
+from cryptography.hazmat.primitives.hashes import SHA256
+from cryptography.hazmat.primitives.serialization import (
+    Encoding,
+    PublicFormat,
+    load_ssh_private_key,
+    load_ssh_public_key,
+)
+
+from .base import SignerBase, VerifierError
+
+
+def _discover_ssh_key() -> Optional[str]:
+    from pytest_gpu_proof.gitutils import get_git_signing_key
+
+    signing_key = get_git_signing_key()
+    if signing_key and os.path.exists(signing_key):
+        return signing_key
+
+    home = Path.home()
+    for candidate in ["id_ed25519", "id_ecdsa", "id_rsa"]:
+        path = home / ".ssh" / candidate
+        if path.exists():
+            return str(path)
+
+    return None
+
+
+def _public_key_fingerprint(public_key) -> str:
+    """SHA256 fingerprint matching `ssh-keygen -l -E sha256` output."""
+    ssh_line = public_key.public_bytes(Encoding.OpenSSH, PublicFormat.OpenSSH)
+    b64_part = ssh_line.split()[1]
+    raw = base64.b64decode(b64_part)
+    digest = hashlib.sha256(raw).digest()
+    return "SHA256:" + base64.b64encode(digest).decode().rstrip("=")
+
+
+def _sign_with_key(private_key, data: bytes) -> bytes:
+    if isinstance(private_key, Ed25519PrivateKey):
+        return private_key.sign(data)
+    if isinstance(private_key, EllipticCurvePrivateKey):
+        return private_key.sign(data, ECDSA(SHA256()))
+    if isinstance(private_key, RSAPrivateKey):
+        return private_key.sign(
+            data,
+            PSS(mgf=MGF1(SHA256()), salt_length=PSS.MAX_LENGTH),
+            SHA256(),
+        )
+    raise VerifierError(f"Unsupported private key type: {type(private_key).__name__}")
+
+
+def _verify_with_key(public_key, signature: bytes, data: bytes) -> bool:
+    try:
+        if isinstance(public_key, Ed25519PublicKey):
+            public_key.verify(signature, data)
+        elif isinstance(public_key, EllipticCurvePublicKey):
+            public_key.verify(signature, data, ECDSA(SHA256()))
+        elif isinstance(public_key, RSAPublicKey):
+            public_key.verify(
+                signature,
+                data,
+                PSS(mgf=MGF1(SHA256()), salt_length=PSS.MAX_LENGTH),
+                SHA256(),
+            )
+        else:
+            return False
+        return True
+    except InvalidSignature:
+        return False
+
+
+def _parse_pubkey_line(line: str):
+    parts = line.strip().split()
+    if len(parts) < 2:
+        return None
+    try:
+        return load_ssh_public_key(f"{parts[0]} {parts[1]}".encode())
+    except Exception:
+        return None
+
+
+def fetch_github_public_keys(username: str) -> List:
+    url = f"https://github.com/{username}.keys"
+    try:
+        with urlopen(url, timeout=10) as resp:
+            content = resp.read().decode()
+    except Exception as e:
+        raise VerifierError(
+            f"Could not fetch public keys for GitHub user {username!r}: {e}"
+        )
+
+    keys = [_parse_pubkey_line(line) for line in content.splitlines()]
+    keys = [k for k in keys if k is not None]
+
+    if not keys:
+        raise VerifierError(
+            f"No usable public keys found for GitHub user {username!r}. "
+            "Ensure they have SSH keys registered at github.com/settings/keys."
+        )
+    return keys
+
+
+def verify_with_github_keys(data: bytes, signature: bytes, github_username: str) -> bool:
+    """Return True if signature was made by any SSH key the user has on GitHub."""
+    keys = fetch_github_public_keys(github_username)
+    return any(_verify_with_key(k, signature, data) for k in keys)
+
+
+class SSHSigner(SignerBase):
+    """Signs with the developer's SSH private key (same key used to push to GitHub)."""
+
+    def __init__(self, key_path: Optional[str] = None):
+        if key_path is None:
+            key_path = _discover_ssh_key()
+        if key_path is None:
+            raise VerifierError(
+                "No SSH private key found. Tried git config user.signingKey and "
+                "~/.ssh/id_ed25519, ~/.ssh/id_ecdsa, ~/.ssh/id_rsa. "
+                "Pass --gpu-proof-key=PATH to specify one explicitly."
+            )
+
+        self._key_path = key_path
+        try:
+            key_data = Path(key_path).read_bytes()
+        except FileNotFoundError:
+            raise VerifierError(f"No SSH private key found at {key_path}")
+
+        try:
+            self._private_key = load_ssh_private_key(key_data, password=None)
+        except TypeError:
+            import getpass
+            pw = getpass.getpass(f"Passphrase for {key_path}: ").encode()
+            self._private_key = load_ssh_private_key(key_data, password=pw)
+
+        self._public_key = self._private_key.public_key()
+
+    def sign(self, data: bytes) -> bytes:
+        return _sign_with_key(self._private_key, data)
+
+    def key_fingerprint(self) -> str:
+        return _public_key_fingerprint(self._public_key)
+
+    def algorithm(self) -> str:
+        if isinstance(self._private_key, Ed25519PrivateKey):
+            return "ed25519"
+        if isinstance(self._private_key, EllipticCurvePrivateKey):
+            return "ecdsa-sha256"
+        if isinstance(self._private_key, RSAPrivateKey):
+            return "rsa-pss-sha256"
+        raise VerifierError(
+            f"Unsupported private key type: {type(self._private_key).__name__}"
+        )