From f3a6e8b09ca76933897d3b7c6a850680e5ca3328 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 23 Sep 2026 21:31:35 -0700 Subject: [PATCH 1/3] Compare resident pipeline and separate stage processes Add a reproducible matched-API benchmark with explicit startup and warm batch timing, file-mediated stage boundaries, and scientific parity checks for every image and round. Signed-off-by: Trent Nelson --- docs/components/pipeline-stage-benchmark.md | 133 +++++ docs/components/xscan.md | 5 + .../xscan/pipeline_benchmark/__init__.py | 5 + .../xscan/pipeline_benchmark/__main__.py | 518 ++++++++++++++++ .../xscan/pipeline_benchmark/fixture.py | 326 ++++++++++ .../xscan/pipeline_benchmark/stages.py | 558 ++++++++++++++++++ tests/xscan/test_pipeline_benchmark.py | 184 ++++++ tests/xscan/test_pipeline_benchmark_stages.py | 260 ++++++++ 8 files changed, 1989 insertions(+) create mode 100644 docs/components/pipeline-stage-benchmark.md create mode 100644 src/cuphoton/xscan/pipeline_benchmark/__init__.py create mode 100644 src/cuphoton/xscan/pipeline_benchmark/__main__.py create mode 100644 src/cuphoton/xscan/pipeline_benchmark/fixture.py create mode 100644 src/cuphoton/xscan/pipeline_benchmark/stages.py create mode 100644 tests/xscan/test_pipeline_benchmark.py create mode 100644 tests/xscan/test_pipeline_benchmark_stages.py diff --git a/docs/components/pipeline-stage-benchmark.md b/docs/components/pipeline-stage-benchmark.md new file mode 100644 index 0000000..e1f1f76 --- /dev/null +++ b/docs/components/pipeline-stage-benchmark.md @@ -0,0 +1,133 @@ +# Pipeline versus separate-stage benchmark + +This benchmark compares one persistent GPU worker running xPOIS, xFit and +XScan with three fresh processes that exchange intermediate arrays through +files. Both treatments use the same device numerical APIs, inputs, +checkpoint and per-image candidate batches. The comparison measures the +combined cost of process lifetime, transfers and intermediate artifacts. + +The separate-stage treatment uses benchmark-specific entry points. It does +not invoke the stock component CLI commands: those use different host +materialization, coefficient-solving and probability-calculation paths. + +## Run the comparison + +Use a CUDA 13 GPU with both CuPy and PyTorch available. From the repository +root: + +```bash +uv sync --locked --extra gpu + +uv run --locked --extra gpu python -m cuphoton.xscan.pipeline_benchmark \ + --output /tmp/cuphoton-pipeline-forward \ + --images 4 --image-size 256 --candidates 9 --stamp-size 17 \ + --seed 2026 --device cuda:0 --warmup 1 --repeat 3 \ + --order pipeline-first +``` + +Each output directory must be new. The default fixture contains four +independent image pairs and nine candidates per image, so each measured +round processes four pairs and 36 candidate stamps. Candidates are supplied +positions; this benchmark does not include source detection. All images +run serially on the selected GPU. + +Reverse the treatment order using the *same generated fixture*: + +```bash +uv run --locked --extra gpu python -m cuphoton.xscan.pipeline_benchmark \ + --output /tmp/cuphoton-pipeline-reverse \ + --config /tmp/cuphoton-pipeline-forward/input/config.json \ + --items /tmp/cuphoton-pipeline-forward/input/items.json \ + --warmup 1 --repeat 3 --order staged-first +``` + +`--config` and `--items` must be supplied together. They contain +`DevicePipelineConfig.to_payload()` and an ordered list of +`DevicePipelineItem.to_payload()` values, including absolute file paths and +hashes. With these options, the manifests determine the device and workload; +fixture-generation options do not replace them. Preserve the referenced +files when running the reversed comparison. Use the same CPU affinity, +thread settings and GPU, without another workload running concurrently. +`--timeout` bounds each child invocation; its default is 600 seconds. + +## What each treatment runs + +The pipeline initializes `DeviceWorkerContext` once, performs its warmup +rounds, then retains the model and CUDA context for measured rounds. Each +item flows through constant-kernel subtraction, stamp extraction, Gaussian +difference fitting, xFit feature conversion and model inference. Device +arrays pass to PyTorch through DLPack. The result contains compact scientific +evidence and predictions. + +The separate-stage treatment starts one xPOIS child, one xFit child and one +XScan child per complete round. Each child processes all image items in +manifest order, keeping the candidate batch for each image unchanged. It +writes lossless, uncompressed NPY arrays between stages. Those files include +full subtraction images and fit residuals in addition to the compact +scientific outputs shared with the pipeline. Their transfers, hashing and +I/O are part of this treatment's cost. Its preliminary rounds prepare +caches; subsequent measured rounds still create fresh processes. + +Both treatments use float64 subtraction and xFit inputs, unweighted xFit +stamps, float32 triplets/features, and the same GPU sigmoid. Inference forces +AMP, TF32, compilation and cuDNN benchmarking off. An xPOIS variance plane +does not become an xFit variance plane. + +## Read the timers + +| Measurement | Boundary | +| --- | --- | +| Pipeline `setup_seconds` / `context_load_seconds` | Worker setup and context/model initialization, recorded separately from numerical warmup and measured batches. The external invocation also includes interpreter startup. | +| Pipeline `batch_seconds` | One ordered image batch through completed device work and per-item result JSON writes. Measured batches reuse the initialized, warmed worker. | +| Pipeline `invocation_external_seconds` | Parent-observed process duration, including imports, setup, all warmup and measured rounds, final summary and process exit. | +| Staged `batch_seconds` | Parent clock before launching xPOIS through successful XScan process exit, including all three process lifetimes and intermediate files. | +| Staged per-process `external_seconds` | Parent-observed duration of that individual child, including imports and shutdown. | + +Stage summaries also record internal setup, reads, uploads, computation, +downloads and writes. Their GPU operations synchronize explicitly; the +pipeline's existing component timers are asynchronous host elapsed times, +with pending work completed at the terminal copy. Comparing those component +timers as isolated GPU kernel durations would be misleading. Use the +completed batch timers for the main workflow comparison. + +The report's `fresh_stages_over_warm_pipeline_ratio` divides the median +staged batch time by the median warm pipeline batch time. It includes +repeated startup and file costs in the staged treatment. Keep setup and +whole-invocation measurements alongside that ratio when discussing a service +that may process only a few batches. + +The benchmark preserves existing filesystem and CUDA caches. Fresh +processes can benefit from both. Writes close files without `fsync`, so +elapsed time does not measure durable storage completion. Fixture preparation +and the final parity audit run outside the treatment timers. + +## Acceptance and artifacts + +The benchmark checks every warmup and measured item against the other +treatment. Acceptance requires exact equality of all 22 compact scientific +arrays, including shapes, dtypes and NaN locations, plus matching candidate +identity/order, fit metadata, subtraction diagnostics and predictions. It +verifies retained file hashes and rechecks the original inputs. This checks +equivalence between two compositions of the same algorithms; it does not +independently establish their astronomical accuracy. Full intermediate +images are retained by the staged treatment but are outside the pipeline's +compact parity contract. + +`manifest.json` records the configuration, ordered input descriptors and +runtime settings. Treatment subdirectories retain each attempted round and +child log. `parity.json` records numerical comparisons; a successful +`report.json` contains all round times and the measured medians. A failure +stops the run and leaves its logs and `failure.json` for inspection. A +successful process exit alone does not satisfy the parity check. + +The default inputs are synthetic textured images with planted dipoles, a +known convolution kernel, noise and an exclusion mask around the candidates. +The fitted 15×15 kernel uses Gaussian sigmas 1.5/3/6 and degrees 2/1/0, +a constant background and no flux constraint. +A small, randomly initialized triplet model has a nonzero xFit fusion branch. +A separate CPU fit creates the canonical feature-schema artifacts before +measurement. Input arrays and model weights repeat for the same seed; +provenance paths and timestamps can change artifact hashes across newly +prepared fixtures. This fixture exercises the workflow and numerical +boundaries. It provides no trained-classifier accuracy result or claim of +representative production throughput. diff --git a/docs/components/xscan.md b/docs/components/xscan.md index 4f0c90e..cc1bcb5 100644 --- a/docs/components/xscan.md +++ b/docs/components/xscan.md @@ -296,6 +296,11 @@ outputs, and merging and validation occur after the timed worker phase. ## Persistent XPOIS, xFit and XScan pipeline +For a reproducible comparison with separately launched stages and intermediate +files, see the [pipeline stage benchmark](pipeline-stage-benchmark.md). It +checks the same scientific outputs while reporting startup and warm execution +separately. + The Python API in `cuphoton.xscan.device_pipeline` runs complete image pairs through constant-kernel XPOIS, stamp extraction, Gaussian difference-mode xFit, feature conversion and triplet XScan inference. A `DeviceWorkerContext` diff --git a/src/cuphoton/xscan/pipeline_benchmark/__init__.py b/src/cuphoton/xscan/pipeline_benchmark/__init__.py new file mode 100644 index 0000000..d321f15 --- /dev/null +++ b/src/cuphoton/xscan/pipeline_benchmark/__init__.py @@ -0,0 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Reproducible pipeline versus file-mediated stage measurement helpers.""" diff --git a/src/cuphoton/xscan/pipeline_benchmark/__main__.py b/src/cuphoton/xscan/pipeline_benchmark/__main__.py new file mode 100644 index 0000000..4942261 --- /dev/null +++ b/src/cuphoton/xscan/pipeline_benchmark/__main__.py @@ -0,0 +1,518 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Compare a persistent device pipeline with three file-mediated processes.""" + +from __future__ import annotations + +import argparse +import importlib.metadata +import json +import os +import platform +import signal +import statistics +import subprocess +import sys +import time +from pathlib import Path +from typing import Any + +import numpy as np + +MODULE = "cuphoton.xscan.pipeline_benchmark" + + +def write_json(path: Path, value: Any) -> None: + path.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n") + + +def read_inputs(config_path: Path, items_path: Path) -> tuple[Any, list[Any]]: + from cuphoton.xscan.device_pipeline import ( + DevicePipelineConfig, + DevicePipelineItem, + ) + + config = DevicePipelineConfig.from_payload( + json.loads(config_path.read_text()) + ) + items = [ + DevicePipelineItem.from_payload(value) + for value in json.loads(items_path.read_text()) + ] + if not items or len({item.item_id for item in items}) != len(items): + raise ValueError("items must be nonempty with unique item IDs") + return config, items + + +def compare_arrays( + expected: dict[str, np.ndarray], actual: dict[str, np.ndarray] +) -> dict[str, Any]: + """Require exact values, shapes, dtypes and NaN locations.""" + if expected.keys() != actual.keys(): + raise ValueError("scientific evidence array names differ") + checks = {} + for name, left in expected.items(): + right = actual[name] + if left.shape != right.shape or left.dtype != right.dtype: + raise ValueError(f"{name}: scientific shape or dtype differs") + equal = np.array_equal(left, right, equal_nan=True) + finite = np.isfinite(left) & np.isfinite(right) + error = ( + float( + np.max( + np.abs( + left[finite].astype(np.float64) + - right[finite].astype(np.float64) + ) + ) + ) + if np.any(finite) + else 0.0 + ) + checks[name] = {"equal": bool(equal), "max_absolute_error": error} + return { + "passed": all(v["equal"] for v in checks.values()), + "arrays": checks, + } + + +def run_child(command: list[str], log: Path, *, timeout: float) -> float: + """Wait for an actual process exit; preserve its output on failure.""" + started = time.perf_counter() + with log.open("w") as stream: + process = subprocess.Popen( + command, + stdout=stream, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + try: + returncode = process.wait(timeout=timeout) + if returncode: + raise subprocess.CalledProcessError(returncode, command) + except BaseException: + # The session contains only this benchmark child and its children. + # Clean it up on interrupts as well as ordinary timeouts/failures. + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait() + raise + return time.perf_counter() - started + + +def pipeline_worker(args: argparse.Namespace) -> None: + started = time.perf_counter() + import cupy as cp + import torch + + from cuphoton.xscan.device_pipeline import ( + DeviceWorkerContext, + run_device_pipeline_item, + ) + + config, items = read_inputs(args.config, args.items) + ordinal = int(config.device.split(":")[1]) + torch.set_num_threads(config.inference_policy["worker_cpu_threads"]) + torch.cuda.set_device(ordinal) + cp.cuda.Device(ordinal).use() + context = DeviceWorkerContext.initialize(config) + cp.cuda.runtime.deviceSynchronize() + setup_seconds = time.perf_counter() - started + rounds = [] + for index in range(args.warmup + args.repeat): + output = args.output / f"round-{index:03d}" + output.mkdir() + phase = "warmup" if index < args.warmup else "measured" + completed_items = 0 + round_started = time.perf_counter() + try: + for item_index, item in enumerate(items): + result = run_device_pipeline_item(item, context) + write_json( + output / f"item-{item_index:04d}.json", + result.to_payload(), + ) + completed_items += 1 + cp.cuda.runtime.deviceSynchronize() + torch.cuda.synchronize(ordinal) + except BaseException as exc: + record = { + "index": index, + "phase": phase, + "status": "failed", + "batch_seconds": time.perf_counter() - round_started, + "completed_items": completed_items, + "failure": {"type": type(exc).__name__, "message": str(exc)}, + } + try: + write_json(output / "timing.json", record) + except OSError as write_error: + exc.add_note( + f"Could not preserve round timing: {write_error}" + ) + raise + record = { + "index": index, + "phase": phase, + "status": "success", + "batch_seconds": time.perf_counter() - round_started, + "completed_items": completed_items, + } + rounds.append(record) + write_json(output / "timing.json", record) + print(json.dumps(record), flush=True) + write_json( + args.output / "summary.json", + { + "setup_seconds": setup_seconds, + "context_load_seconds": context.load_seconds, + "worker_seconds": time.perf_counter() - started, + "configuration_sha256": config.configuration_sha256, + "rounds": rounds, + "gpu_name": torch.cuda.get_device_name(ordinal), + }, + ) + + +def measure_pipeline(args: argparse.Namespace) -> dict[str, Any]: + output = args.output / "pipeline" + output.mkdir() + elapsed = run_child( + [ + sys.executable, + "-m", + MODULE, + "--pipeline-worker", + "--config", + str(args.config), + "--items", + str(args.items), + "--output", + str(output), + "--warmup", + str(args.warmup), + "--repeat", + str(args.repeat), + ], + output / "process.log", + timeout=args.timeout, + ) + result = json.loads((output / "summary.json").read_text()) + result["invocation_external_seconds"] = elapsed + return result + + +def measure_stages(args: argparse.Namespace) -> dict[str, Any]: + output = args.output / "staged" + output.mkdir() + rounds = [] + for index in range(args.warmup + args.repeat): + root = output / f"round-{index:03d}" + root.mkdir() + phase = "cache_preparation" if index < args.warmup else "measured" + record: dict[str, Any] = { + "index": index, + "phase": phase, + "stages": {}, + "completed_stages": [], + } + started = time.perf_counter() + try: + for stage in ("xpois", "xfit", "xscan"): + stage_started = time.perf_counter() + elapsed = run_child( + [ + sys.executable, + "-m", + f"{MODULE}.stages", + "--stage", + stage, + "--config", + str(args.config), + "--items", + str(args.items), + "--output-dir", + str(root), + ], + root / f"{stage}.log", + timeout=args.timeout, + ) + record["stages"][stage] = { + "external_seconds": elapsed, + "status": "success", + } + record["completed_stages"].append(stage) + except BaseException as exc: + ended = time.perf_counter() + record.update( + status="failed", + batch_seconds=ended - started, + failure={"type": type(exc).__name__, "message": str(exc)}, + ) + record["stages"][stage] = { + "external_seconds": ended - stage_started, + "status": "failed", + } + try: + write_json(root / "timing.json", record) + except OSError as write_error: + exc.add_note( + f"Could not preserve round timing: {write_error}" + ) + raise + record["batch_seconds"] = time.perf_counter() - started + record["status"] = "success" + rounds.append(record) + write_json(root / "timing.json", record) + print(json.dumps({"treatment": "staged", **record}), flush=True) + result = {"rounds": rounds} + write_json(output / "summary.json", result) + return result + + +def audit(args: argparse.Namespace) -> dict[str, Any]: + from cuphoton.core.artifacts import file_sha256 + from cuphoton.xscan.device_pipeline import ( + _verify_item_hashes, + decode_device_pipeline_evidence, + ) + + from .stages import load_science_arrays + + config, items = read_inputs(args.config, args.items) + checks = [] + for index in range(args.warmup + args.repeat): + root = args.output / "staged" / f"round-{index:03d}" + stages = { + stage: json.loads((root / stage / "summary.json").read_text()) + for stage in ("xpois", "xfit", "xscan") + } + for summary in stages.values(): + if summary["configuration_sha256"] != config.configuration_sha256: + raise ValueError("staged configuration differs") + if [entry["item"] for entry in summary["items"]] != [ + item.to_payload() for item in items + ]: + raise ValueError("staged candidate/item identity differs") + for item_index, item in enumerate(items): + pipeline_path = ( + args.output + / "pipeline" + / f"round-{index:03d}" + / f"item-{item_index:04d}.json" + ) + result = json.loads(pipeline_path.read_text()) + if result["item_id"] != item.item_id: + raise ValueError("pipeline item identity differs") + if result["configuration_sha256"] != config.configuration_sha256: + raise ValueError("pipeline configuration differs") + for candidate, prediction in zip( + item.candidates, result["predictions"], strict=True + ): + for key, value in candidate.to_payload().items(): + if prediction[key] != value: + raise ValueError( + "pipeline candidate identity differs" + ) + if prediction["decision"] != ( + prediction["probability"] >= config.decision_threshold + ): + raise ValueError("pipeline decision differs") + expected = decode_device_pipeline_evidence( + result["scientific_evidence"], config=config + ) + metadata = { + stage: summary["items"][item_index]["metadata"] + for stage, summary in stages.items() + } + for name, value in result["xpois"].items(): + if metadata["xpois"][f"xpois_{name}"] != value: + raise ValueError(f"xPOIS {name} differs") + evidence = result["scientific_evidence"] + if ( + metadata["xpois"]["basis_terms"] + != evidence["xpois"]["basis_terms"] + ): + raise ValueError("xPOIS basis metadata differs") + for name in ("parameter_names", "feature_names"): + if metadata["xfit"][name] != evidence[name]: + raise ValueError(f"xFit {name} differs") + if metadata["xscan"]["predictions"] != result["predictions"]: + raise ValueError("xScan predictions or identity differs") + actual = load_science_arrays(root, item_index) + checks.append( + { + "round": index, + "item_id": item.item_id, + **compare_arrays(expected, actual), + } + ) + for item in items: + _verify_item_hashes(item, when="after benchmark") + for path, expected_digest in ( + ( + Path(config.checkpoint_dir) / "checkpoint.pt", + config.checkpoint_sha256, + ), + (Path(config.feature_schema_path), config.feature_schema_sha256), + ): + if file_sha256(path) != expected_digest: + raise ValueError( + "checkpoint or feature schema changed after execution" + ) + return {"passed": all(v["passed"] for v in checks), "checks": checks} + + +def provenance() -> dict[str, Any]: + from cuphoton.core.artifacts import file_sha256 + + package_root = Path(__file__).resolve().parents[2] + numerical_sources = ( + "xscan/device_pipeline.py", + "xpois/ois.py", + "xfit/api.py", + "xfit/backend.py", + "xfit/models.py", + "xfit/solver.py", + "xscan/xfit_features.py", + "xscan/training.py", + "xscan/model.py", + ) + packages = {} + for name in ("cuphoton", "numpy", "cupy-cuda13x", "torch"): + try: + packages[name] = importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + packages[name] = None + policy_keys = ( + "CUDA_VISIBLE_DEVICES", + "OMP_NUM_THREADS", + "MKL_NUM_THREADS", + "OPENBLAS_NUM_THREADS", + "CUPY_CACHE_DIR", + "CUDA_CACHE_PATH", + ) + return { + "python": platform.python_version(), + "platform": platform.platform(), + "packages": packages, + "environment": {key: os.environ.get(key) for key in policy_keys}, + "cpu_affinity": sorted(os.sched_getaffinity(0)), + "benchmark_source_sha256": { + path.name: file_sha256(path) + for path in sorted(Path(__file__).parent.glob("*.py")) + }, + "numerical_source_sha256": { + name: file_sha256(package_root / name) + for name in numerical_sources + }, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--config", type=Path) + parser.add_argument("--items", type=Path) + parser.add_argument("--images", type=int, default=4) + parser.add_argument("--image-size", type=int, default=256) + parser.add_argument("--candidates", type=int, default=9) + parser.add_argument("--stamp-size", type=int, default=17) + parser.add_argument("--seed", type=int, default=2026) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--repeat", type=int, default=3) + parser.add_argument("--timeout", type=float, default=600) + parser.add_argument( + "--order", + choices=("pipeline-first", "staged-first"), + default="pipeline-first", + ) + parser.add_argument( + "--pipeline-worker", action="store_true", help=argparse.SUPPRESS + ) + args = parser.parse_args() + if args.repeat < 1 or args.warmup < 1 or args.timeout <= 0: + parser.error("repeat, warmup and timeout must be positive") + if (args.config is None) != (args.items is None): + parser.error("provide both --config and --items, or neither") + args.output = args.output.resolve() + if args.pipeline_worker: + pipeline_worker(args) + return + args.output.mkdir(parents=True, exist_ok=False) + try: + if args.config is None: + from .fixture import prepare_fixture + + args.config, args.items = prepare_fixture( + args.output / "input", + images=args.images, + image_size=args.image_size, + candidates=args.candidates, + stamp_size=args.stamp_size, + seed=args.seed, + device=args.device, + ) + args.config, args.items = args.config.resolve(), args.items.resolve() + config, items = read_inputs(args.config, args.items) + report: dict[str, Any] = { + "schema": "cuphoton.pipeline-stage-benchmark/v1", + "provenance": provenance(), + "configuration": config.to_payload(), + "items": [item.to_payload() for item in items], + "order": args.order, + "warmup": args.warmup, + "repeat": args.repeat, + "fsync": False, + "cache_policy": "preserve existing caches", + } + write_json(args.output / "manifest.json", report) + treatments = ( + [("pipeline", measure_pipeline), ("staged", measure_stages)] + if args.order == "pipeline-first" + else [("staged", measure_stages), ("pipeline", measure_pipeline)] + ) + for name, run in treatments: + report[name] = run(args) + acceptance = audit(args) + write_json(args.output / "parity.json", acceptance) + if not acceptance["passed"]: + raise ValueError("scientific parity failed; see parity.json") + report["parity_passed"] = True + for name in ("pipeline", "staged"): + values = [ + row["batch_seconds"] + for row in report[name]["rounds"] + if row["phase"] == "measured" + ] + report[name]["measured_batch_median_seconds"] = statistics.median( + values + ) + report["fresh_stages_over_warm_pipeline_ratio"] = ( + report["staged"]["measured_batch_median_seconds"] + / report["pipeline"]["measured_batch_median_seconds"] + ) + write_json(args.output / "report.json", report) + print( + json.dumps( + {"report": str(args.output / "report.json"), "parity": True} + ) + ) + except BaseException as exc: + write_json( + args.output / "failure.json", + { + "type": type(exc).__name__, + "message": str(exc), + }, + ) + raise + + +if __name__ == "__main__": + main() diff --git a/src/cuphoton/xscan/pipeline_benchmark/fixture.py b/src/cuphoton/xscan/pipeline_benchmark/fixture.py new file mode 100644 index 0000000..93d8a89 --- /dev/null +++ b/src/cuphoton/xscan/pipeline_benchmark/fixture.py @@ -0,0 +1,326 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Deterministic CPU inputs and an untrained benchmark fusion model.""" + +from __future__ import annotations + +import json +import math +from dataclasses import asdict +from pathlib import Path +from typing import Any + +import numpy as np +from scipy.signal import fftconvolve + +from cuphoton.core.artifacts import file_sha256 +from cuphoton.xfit import GaussianDipoleModel, LMConfig, fit_dipoles +from cuphoton.xfit.io import load_xfit_dataset, write_fit_artifacts +from cuphoton.xpois import ( + GaussianBasisComponent, + build_gaussian_polynomial_basis, +) +from cuphoton.xscan.config import ModelConfig, PerformanceConfig +from cuphoton.xscan.device_pipeline import ( + DevicePipelineCandidate, + DevicePipelineConfig, + DevicePipelineItem, + DeviceXFitPipelineConfig, + DeviceXPOISPipelineConfig, + NpyArrayDescriptor, +) +from cuphoton.xscan.xfit_features import ( + FEATURE_NAMES, + build_xfit_feature_bundle, +) + + +def _write_json(path: Path, payload: Any) -> None: + path.write_text( + json.dumps(payload, indent=2, sort_keys=True, allow_nan=False) + "\n", + encoding="utf-8", + ) + + +def _save_array(path: Path, values: np.ndarray) -> NpyArrayDescriptor: + values = np.ascontiguousarray(values) + np.save(path, values, allow_pickle=False) + return NpyArrayDescriptor( + path=str(path.resolve()), + sha256=file_sha256(path), + shape=tuple(values.shape), + dtype=values.dtype.str, + ) + + +def _feature_schema(root: Path, stamp_size: int) -> Path: + """Produce the canonical schema and artifacts from a real CPU fit.""" + + calibration = root / "schema-calibration" + calibration.mkdir() + truth = np.array([[1.0, 1.6, 1.2, 0.2, 1.5, 0.6, -1.5, -0.6]]) + stamps = GaussianDipoleModel((stamp_size, stamp_size)).evaluate(truth) + stamps += np.random.default_rng(0).normal(0.0, 0.001, stamps.shape) + np.save(calibration / "difference.npy", stamps, allow_pickle=False) + (calibration / "metadata.jsonl").write_text( + '{"candidate_id":"schema-example"}\n', encoding="utf-8" + ) + input_path = calibration / "input.npz" + np.savez( + input_path, + candidate_id=np.array(["schema-example"]), + images=stamps, + ) + dataset = load_xfit_dataset( + input_path, model="gaussian", mode="difference" + ) + solver = LMConfig(max_evaluations=100) + result = fit_dipoles( + stamps, + model="gaussian", + initial=truth, + backend="numpy", + config=solver, + ) + fit_root = calibration / "fit" + write_fit_artifacts( + fit_root, + dataset=dataset, + result=result, + effective_config={ + "backend": "numpy", + "model": "gaussian", + "mode": "difference", + "solver": asdict(solver), + "purpose": "Synthetic schema example, not benchmark inference", + }, + ) + feature_root = calibration / "features" + build_xfit_feature_bundle( + dataset_dir=calibration, + xfit_run_dir=fit_root, + output_dir=feature_root, + ) + return feature_root / "schema.json" + + +def _checkpoint( + root: Path, stamp_size: int, seed: int, schema_path: Path +) -> Path: + import torch + + from cuphoton.xscan.model import build_model + + model_config = ModelConfig( + input_mode="triplet", + image_size=stamp_size, + depths=[1], + num_heads=[1], + embed_dims=[8], + decoder_embedding_dim=4, + pos_dim=8, + output_nc=1, + drop_rate=0.0, + attn_drop=0.0, + xfit_feature_names=list(FEATURE_NAMES), + xfit_hidden_dim=8, + ) + # Fork only the CPU generator; construction never initializes CUDA. + with torch.random.fork_rng(devices=[]), torch.device("cpu"): + torch.random.default_generator.manual_seed(seed) + model = build_model(**asdict(model_config)) + assert model.xfit_head is not None + # The model constructor zeroes this layer. Exercise fusion here. + with torch.no_grad(): + model.xfit_head[-1].weight.normal_(mean=0.0, std=0.1) + model.xfit_head[-1].bias.fill_(0.03) + schema = json.loads(schema_path.read_text(encoding="utf-8")) + checkpoint_dir = root / "checkpoint" + checkpoint_dir.mkdir() + torch.save( + { + "model_config": asdict(model_config), + "model_state": model.state_dict(), + "train_config": {"performance": asdict(PerformanceConfig())}, + "xfit_feature_bundle": { + "schema_sha256": file_sha256(schema_path), + "feature_sha256": schema["artifacts"]["features"]["sha256"], + "source_artifacts": schema["source_artifacts"], + }, + "benchmark_fixture": { + "seed": seed, + "trained": False, + "description": "Random synthetic model; no accuracy claim", + }, + }, + checkpoint_dir / "checkpoint.pt", + ) + return checkpoint_dir + + +def prepare_fixture( + root: Path, + *, + images: int, + image_size: int, + candidates: int, + stamp_size: int, + seed: int, + device: str, +) -> tuple[Path, Path]: + """Write config/items manifests and all inputs below a new directory. + + Each image contains an independently seeded textured reference, a matched + target with planted dipoles, a variance plane and an exclusion mask around + candidate stamps. Candidate order is a seeded permutation of spatial grid + positions. Inputs and model weights repeat for the same arguments; + provenance paths/timestamps and their artifact hashes are specific to this + preparation. The small untrained model is a workflow fixture, not a model + accuracy or representative-production-throughput benchmark. + """ + + for name, value in ( + ("images", images), + ("image_size", image_size), + ("candidates", candidates), + ("stamp_size", stamp_size), + ): + if ( + isinstance(value, bool) + or not isinstance(value, int) + or value <= 0 + ): + raise ValueError(f"{name} must be a positive integer") + if stamp_size < 9 or stamp_size % 2 == 0: + raise ValueError("stamp_size must be odd and at least 9") + if ( + isinstance(seed, bool) + or not isinstance(seed, int) + or not 0 <= seed < 2**63 + ): + raise ValueError("seed must be an integer in [0, 2**63)") + if not isinstance(device, str) or not device.startswith("cuda:"): + raise ValueError("device must be an explicit cuda:N device") + ordinal = device.removeprefix("cuda:") + if ( + not ordinal.isascii() + or not ordinal.isdecimal() + or str(int(ordinal)) != ordinal + ): + raise ValueError("device must be an explicit cuda:N device") + half = stamp_size // 2 + margin = half + 7 + side = math.isqrt(candidates - 1) + 1 + low, high = margin, image_size - 1 - margin + if high - low < 2 or ( + side > 1 and (high - low) // (side - 1) < stamp_size + 2 + ): + raise ValueError( + "image_size is too small for nonoverlapping candidate stamps" + ) + coordinates = ( + np.array([image_size // 2]) + if side == 1 + else np.linspace(low, high, side, dtype=int) + ) + centers = [(int(y), int(x)) for y in coordinates for x in coordinates][ + :candidates + ] + root = root.expanduser().resolve() + root.mkdir(parents=True, exist_ok=False) + schema_path = _feature_schema(root, stamp_size) + checkpoint_dir = _checkpoint(root, stamp_size, seed, schema_path) + xpois_config = DeviceXPOISPipelineConfig( + kernel_shape=(15, 15), + basis_sigmas=(1.5, 3.0, 6.0), + basis_degrees=(2, 1, 0), + flux_conserve=False, + ) + config = DevicePipelineConfig( + device=device, + checkpoint_dir=str(checkpoint_dir), + checkpoint_sha256=file_sha256(checkpoint_dir / "checkpoint.pt"), + feature_schema_path=str(schema_path), + feature_schema_sha256=file_sha256(schema_path), + stamp_shape=(stamp_size, stamp_size), + decision_threshold=0.5, + xpois=xpois_config, + xfit=DeviceXFitPipelineConfig(max_evaluations=100), + ) + kernels, _ = build_gaussian_polynomial_basis( + xpois_config.kernel_shape, + [ + GaussianBasisComponent(sigma=sigma, degree=degree) + for sigma, degree in zip( + xpois_config.basis_sigmas, + xpois_config.basis_degrees, + strict=True, + ) + ], + flux_conserve=xpois_config.flux_conserve, + ) + truth_kernel = 0.88 * kernels[0] + 0.12 * kernels[-1] + model = GaussianDipoleModel(config.stamp_shape) + items = [] + for image_index in range(images): + rng = np.random.default_rng( + np.random.SeedSequence([seed, image_index]) + ) + reference = rng.normal(0.2, 0.05, (image_size, image_size)) + target = fftconvolve(reference, truth_kernel, mode="same") + 0.1 + target += rng.normal(0.0, 0.001, target.shape) + fit_mask = np.ones(target.shape, dtype=bool) + descriptors = [] + for spatial_index in rng.permutation(candidates).tolist(): + cy, cx = centers[spatial_index] + region = ( + slice(cy - half, cy + half + 1), + slice(cx - half, cx + half + 1), + ) + truth = np.array( + [ + [ + rng.uniform(0.5, 1.5), + 1.7, + 1.2, + rng.uniform(-0.5, 0.5), + 1.5, + 0.7, + -1.5, + -0.7, + ] + ], + dtype=np.float64, + ) + target[region] += model.evaluate(truth)[0] + fit_mask[region] = False + descriptors.append( + DevicePipelineCandidate( + candidate_id=f"image-{image_index:04d}-candidate-{spatial_index:04d}", + center_x=cx, + center_y=cy, + source_index=spatial_index, + ) + ) + image_root = root / f"image-{image_index:04d}" + image_root.mkdir() + items.append( + DevicePipelineItem( + item_id=f"image-{image_index:04d}", + reference=_save_array( + image_root / "reference.npy", reference + ), + target=_save_array(image_root / "target.npy", target), + variance=_save_array( + image_root / "variance.npy", np.full(target.shape, 1.0e-6) + ), + fit_mask=_save_array(image_root / "fit-mask.npy", fit_mask), + candidates=tuple(descriptors), + ) + ) + config_path, items_path = root / "config.json", root / "items.json" + _write_json(config_path, config.to_payload()) + _write_json(items_path, [item.to_payload() for item in items]) + return config_path, items_path diff --git a/src/cuphoton/xscan/pipeline_benchmark/stages.py b/src/cuphoton/xscan/pipeline_benchmark/stages.py new file mode 100644 index 0000000..6a373c3 --- /dev/null +++ b/src/cuphoton/xscan/pipeline_benchmark/stages.py @@ -0,0 +1,558 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Matched numerical stages with real process/file boundaries. + +This benchmark reuses private device-pipeline validation and stamp helpers; +those helpers are not new public APIs. Each invocation processes the complete +ordered item manifest. Stage timers explicitly synchronize; they must not be +compared as GPU-only timings with the pipeline's asynchronous host timers. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import time +from dataclasses import asdict +from pathlib import Path +from typing import Any, Callable + +import numpy as np + +from cuphoton.core.artifacts import file_sha256 +from cuphoton.xscan import device_pipeline as pipeline + +SCHEMA = "cuphoton.xscan.pipeline-benchmark.stage/v1" +STAGES = ("xpois", "xfit", "xscan") +XFIT_FIELDS = ( + *pipeline._XFIT_SCALAR_EVIDENCE_FIELDS, + "parameters", + "standard_errors", + "covariance", +) +SCIENCE_KEYS = ( + "xpois.kernel_coefficients", + "xpois.background_coefficients", + *(f"xfit.{name}" for name in XFIT_FIELDS), + "xfit.features", + "xscan.logits", + "xscan.probabilities", +) + + +def _json_hash(value: Any) -> str: + return hashlib.sha256( + json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + + +def _write_json(path: Path, value: Any) -> None: + with path.open("x", encoding="utf-8") as stream: + json.dump(value, stream, indent=2, allow_nan=False) + stream.write("\n") + + +def _measure( + timings: dict[str, float], + name: str, + function: Callable[[], Any], + synchronize: Callable[[], None] | None = None, +) -> Any: + if synchronize is not None: + synchronize() + started = time.perf_counter() + value = function() + if synchronize is not None: + synchronize() + timings[name] = time.perf_counter() - started + return value + + +def _read_summary(root: Path, stage: str) -> dict[str, Any]: + value = json.loads((root / stage / "summary.json").read_text()) + if ( + value.get("schema") != SCHEMA + or value.get("stage") != stage + or value.get("status") != "success" + ): + raise ValueError(f"invalid completed {stage} summary") + expected_upstream = set(STAGES[: STAGES.index(stage)]) + if set(value.get("upstream_summary_sha256", {})) != expected_upstream: + raise ValueError(f"{stage} upstream summary links are incomplete") + records = value.get("items") + if not isinstance(records, list) or not records: + raise ValueError(f"{stage} item manifest is empty or invalid") + if _json_hash([record["item"] for record in records]) != value.get( + "items_sha256" + ): + raise ValueError(f"{stage} item metadata or order changed") + return value + + +def _read_array(root: Path, descriptor: dict[str, Any]) -> np.ndarray: + path = (root / descriptor["path"]).resolve() + if not path.is_relative_to(root.resolve()): + raise ValueError("intermediate artifact escapes its round directory") + if file_sha256(path) != descriptor["sha256"]: + raise ValueError(f"intermediate artifact hash changed: {path.name}") + array = np.load(path, allow_pickle=False) + if ( + not isinstance(array, np.ndarray) + or list(array.shape) != descriptor["shape"] + or array.dtype.str != descriptor["dtype"] + or array.nbytes != descriptor["nbytes"] + or not array.flags.c_contiguous + ): + raise ValueError( + f"intermediate artifact contract changed: {path.name}" + ) + if file_sha256(path) != descriptor["sha256"]: + raise ValueError( + f"intermediate artifact changed while reading: {path.name}" + ) + return array + + +def _write_arrays( + root: Path, directory: Path, arrays: dict[str, np.ndarray] +) -> dict[str, Any]: + directory.mkdir() + artifacts = {} + for name, values in arrays.items(): + array = np.ascontiguousarray(values) + if array.dtype.kind not in "biuf": + raise TypeError(f"unsupported numeric artifact dtype: {name}") + path = directory / f"{name}.npy" + with path.open("xb") as stream: + np.save(stream, array, allow_pickle=False) + artifacts[name] = { + "path": str(path.relative_to(root)), + "sha256": file_sha256(path), + "dtype": array.dtype.str, + "shape": list(array.shape), + "nbytes": array.nbytes, + "file_bytes": path.stat().st_size, + } + return artifacts + + +def _setup( + stage: str, config: pipeline.DevicePipelineConfig +) -> dict[str, Any]: + from cuphoton.xscan.xfit_features import FEATURE_NAMES + + schema = pipeline._load_feature_schema_contract( + Path(config.feature_schema_path), + expected_sha256=config.feature_schema_sha256, + ) + pipeline._validate_feature_schema_source( + schema, config=config, feature_names=FEATURE_NAMES + ) + checkpoint_path = Path(config.checkpoint_dir) / "checkpoint.pt" + if file_sha256(checkpoint_path) != config.checkpoint_sha256: + raise ValueError("checkpoint hash changed") + runtime: dict[str, Any] = {"cp": None, "torch": None} + if stage == "xscan": + import torch + + from cuphoton.xscan.config import PerformanceConfig + from cuphoton.xscan.training import load_model_from_checkpoint + + torch.set_num_threads(config.inference_policy["worker_cpu_threads"]) + torch.cuda.set_device(config.device_id) + runtime.update(torch=torch, device=torch.device(config.device)) + model, checkpoint, performance = load_model_from_checkpoint( + Path(config.checkpoint_dir), + device=runtime["device"], + performance_override=PerformanceConfig(**config.inference_policy), + ) + if asdict(performance) != config.inference_policy: + raise ValueError("inference policy changed during normalization") + pipeline._validate_checkpoint_contract( + checkpoint, + stamp_shape=config.stamp_shape, + feature_names=FEATURE_NAMES, + feature_schema_sha256=config.feature_schema_sha256, + feature_schema=schema, + ) + pipeline._validate_loaded_model_contract( + model, feature_names=FEATURE_NAMES + ) + runtime.update(model=model, performance=performance) + torch.cuda.synchronize(runtime["device"]) + else: + import cupy as cp + + cp.cuda.Device(config.device_id).use() + cp.cuda.Device(config.device_id).synchronize() + runtime.update(cp=cp) + if file_sha256(checkpoint_path) != config.checkpoint_sha256: + raise ValueError("checkpoint changed during stage initialization") + if ( + file_sha256(config.feature_schema_path) + != config.feature_schema_sha256 + ): + raise ValueError("feature schema changed during stage initialization") + return runtime + + +def _run_item( + stage: str, + item: pipeline.DevicePipelineItem, + config: pipeline.DevicePipelineConfig, + runtime: dict[str, Any], + root: Path, + priors: dict[str, dict[str, Any]], + index: int, +) -> tuple[dict[str, np.ndarray], dict[str, Any], dict[str, float]]: + cp, torch = runtime["cp"], runtime["torch"] + + def sync() -> None: + if stage == "xscan": + torch.cuda.synchronize(runtime["device"]) + else: + cp.cuda.Device().synchronize() + + timings: dict[str, float] = {} + + def previous(which: str, name: str) -> np.ndarray: + return _read_array( + root, priors[which]["items"][index]["artifacts"][name] + ) + + if stage == "xpois": + from cuphoton.xpois import ( + GaussianBasisComponent, + solve_constant_kernel_device, + ) + + host = _measure( + timings, "read_seconds", lambda: pipeline._read_item_inputs(item) + ) + device = _measure( + timings, + "upload_seconds", + lambda: {name: cp.asarray(value) for name, value in host.items()}, + sync, + ) + + def compute() -> tuple[dict[str, Any], dict[str, Any]]: + result = solve_constant_kernel_device( + device["reference"], + device["target"], + [ + GaussianBasisComponent(sigma=s, degree=d) + for s, d in zip( + config.xpois.basis_sigmas, + config.xpois.basis_degrees, + strict=True, + ) + ], + kernel_shape=config.xpois.kernel_shape, + variance=device.get("variance"), + fit_mask=device.get("fit_mask"), + background_degree=config.xpois.background_degree, + flux_conserve=config.xpois.flux_conserve, + flux_reference_index=config.xpois.flux_reference_index, + ) + stamps, difference = pipeline._extract_stamps( + cp=cp, + target=device["target"], + reference=device["reference"], + residual=result.residual, + candidates=item.candidates, + stamp_shape=config.stamp_shape, + ) + arrays = { + f"xpois.{name}": getattr(result, name) + for name in ( + "kernel", + "matched", + "residual", + "fit_mask", + "background", + "kernel_coefficients", + "background_coefficients", + "basis_kernels", + ) + } + arrays.update(stamps=stamps, difference=difference) + return arrays, { + "xpois_chi2": float(result.chi2), + "xpois_dof": result.dof, + "xpois_fit_pixel_count": result.fit_pixel_count, + "basis_terms": [asdict(term) for term in result.basis_terms], + "solver": result.solver, + "backend": result.backend, + "flux_conserve": result.flux_conserve, + } + + arrays, metadata = _measure(timings, "compute_seconds", compute, sync) + elif stage == "xfit": + from cuphoton.xfit import LMConfig, fit_dipoles_device + from cuphoton.xscan.xfit_features import ( + transform_xfit_result_features_device, + ) + + host = _measure( + timings, "read_seconds", lambda: previous("xpois", "difference") + ) + difference = _measure( + timings, "upload_seconds", lambda: cp.asarray(host), sync + ) + + def compute() -> tuple[dict[str, Any], dict[str, Any]]: + result = fit_dipoles_device( + difference, + model=config.xfit.model, + mode=config.xfit.mode, + config=LMConfig(**config.xfit.solver_payload()), + ) + features = transform_xfit_result_features_device( + result, + image_shape=config.stamp_shape, + variance_present=False, + ) + return { + **{ + f"xfit.{name}": getattr(result, name) + for name in (*XFIT_FIELDS, "residuals") + }, + "xfit.features": features.values, + }, { + "parameter_names": list(result.parameter_names), + "feature_names": list(features.feature_names), + "model": result.model, + "mode": result.mode, + "backend": result.backend, + "solver": result.solver, + "dtype": result.dtype, + "variance_present": False, + } + + arrays, metadata = _measure(timings, "compute_seconds", compute, sync) + else: + from cuphoton.xscan.training import predict_tensors + + host = _measure( + timings, + "read_seconds", + lambda: { + "images": previous("xpois", "stamps"), + "xfit_features": previous("xfit", "xfit.features"), + }, + ) + device = _measure( + timings, + "upload_seconds", + lambda: { + name: torch.from_numpy(value).to(runtime["device"]) + for name, value in host.items() + }, + sync, + ) + prediction = _measure( + timings, + "compute_seconds", + lambda: predict_tensors( + model=runtime["model"], + **device, + device=runtime["device"], + performance=runtime["performance"], + ), + sync, + ) + arrays = { + f"xscan.{name}": value for name, value in prediction.items() + } + metadata = {"inference_policy": config.inference_policy} + + host_arrays = _measure( + timings, + "download_seconds", + lambda: { + name: ( + value.detach().cpu().numpy() + if stage == "xscan" + else cp.asnumpy(value) + ) + for name, value in arrays.items() + }, + sync, + ) + if stage == "xscan": + if not all( + np.isfinite(value).all() for value in host_arrays.values() + ): + raise ValueError("non-finite xScan predictions") + metadata["predictions"] = [ + { + **candidate.to_payload(), + "logit": float(host_arrays["xscan.logits"][position]), + "probability": float( + host_arrays["xscan.probabilities"][position] + ), + "decision": bool( + host_arrays["xscan.probabilities"][position] + >= config.decision_threshold + ), + } + for position, candidate in enumerate(item.candidates) + ] + return host_arrays, metadata, timings + + +def run_stage( + stage: str, + config: pipeline.DevicePipelineConfig, + items: list[pipeline.DevicePipelineItem], + root: Path, +) -> dict[str, Any]: + """Run one stage, refusing stale output or mismatched inputs.""" + started = time.perf_counter() + if stage not in STAGES or not items: + raise ValueError("a known stage and nonempty item list are required") + if len({item.item_id for item in items}) != len(items): + raise ValueError("item IDs must be unique within a round") + root = Path(root).resolve() + stage_dir = root / stage + stage_dir.mkdir(parents=True, exist_ok=False) + item_payloads = [item.to_payload() for item in items] + items_sha256 = _json_hash(item_payloads) + identity = { + "configuration_sha256": config.configuration_sha256, + "checkpoint_sha256": config.checkpoint_sha256, + "feature_schema_sha256": config.feature_schema_sha256, + "items_sha256": items_sha256, + } + priors = {} + for previous in STAGES[: STAGES.index(stage)]: + prior = _read_summary(root, previous) + if any(prior.get(key) != value for key, value in identity.items()): + raise ValueError( + f"{previous} input/configuration identity mismatch" + ) + if [entry["item"] for entry in prior["items"]] != item_payloads: + raise ValueError(f"{previous} item order or metadata mismatch") + priors[previous] = prior + for item in items: + pipeline._validate_item_contract(item, config=config) + runtime = _setup(stage, config) + setup_seconds = time.perf_counter() - started + records = [] + for index, item in enumerate(items): + item_started = time.perf_counter() + arrays, metadata, timings = _run_item( + stage, + item, + config, + runtime, + root, + priors, + index, + ) + artifacts = _measure( + timings, + "write_seconds", + lambda: _write_arrays( + root, + stage_dir / f"item-{index:04d}", + arrays, + ), + ) + _measure( + timings, + "input_recheck_seconds", + lambda: pipeline._verify_item_hashes( + item, + when=f"during {stage} execution", + ), + ) + timings["wall_seconds"] = time.perf_counter() - item_started + records.append( + { + "item_id": item.item_id, + "item": item.to_payload(), + "metadata": metadata, + "artifacts": artifacts, + "timings_seconds": timings, + } + ) + summary = { + "schema": SCHEMA, + "status": "success", + "stage": stage, + **identity, + "device": config.device, + "setup_seconds": setup_seconds, + "wall_seconds": time.perf_counter() - started, + "upstream_summary_sha256": { + name: file_sha256(root / name / "summary.json") for name in priors + }, + "items": records, + "timing_note": ( + "Setup starts inside run_stage; parent process wall includes " + "imports and shutdown. Compute/upload/download synchronize " + "explicitly; these are host elapsed, not GPU-only. Wall excludes " + "this final summary write. NPY writes close without fsync." + ), + } + _write_json(stage_dir / "summary.json", summary) + return summary + + +def load_science_arrays(root: Path, index: int) -> dict[str, np.ndarray]: + """Read all 22 arrays in the pipeline decoder's float64 form.""" + root = Path(root).resolve() + summaries = {stage: _read_summary(root, stage) for stage in STAGES} + first = summaries["xpois"] + for stage, summary in summaries.items(): + for key in ( + "configuration_sha256", + "checkpoint_sha256", + "feature_schema_sha256", + "items_sha256", + ): + if summary[key] != first[key]: + raise ValueError("scientific stage identities disagree") + if [entry["item"] for entry in summary["items"]] != [ + entry["item"] for entry in first["items"] + ]: + raise ValueError("scientific stage item identities disagree") + for upstream, digest in summary["upstream_summary_sha256"].items(): + if file_sha256(root / upstream / "summary.json") != digest: + raise ValueError(f"{stage} upstream summary changed") + return { + name: _read_array( + root, + summaries[name.split(".")[0]]["items"][index]["artifacts"][name], + ).astype(np.float64) + for name in SCIENCE_KEYS + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--stage", required=True, choices=STAGES) + parser.add_argument("--config", required=True, type=Path) + parser.add_argument("--items", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + args = parser.parse_args() + config = pipeline.DevicePipelineConfig.from_payload( + json.loads(args.config.read_text()) + ) + items = [ + pipeline.DevicePipelineItem.from_payload(item) + for item in json.loads(args.items.read_text()) + ] + summary = run_stage(args.stage, config, items, args.output_dir) + print( + json.dumps({"stage": summary["stage"], "status": summary["status"]}) + ) + + +if __name__ == "__main__": + main() diff --git a/tests/xscan/test_pipeline_benchmark.py b/tests/xscan/test_pipeline_benchmark.py new file mode 100644 index 0000000..1466f4b --- /dev/null +++ b/tests/xscan/test_pipeline_benchmark.py @@ -0,0 +1,184 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Scientific acceptance and process failure gates for the benchmark.""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest + +from cuphoton.xscan.pipeline_benchmark import __main__ as benchmark +from cuphoton.xscan.pipeline_benchmark.__main__ import ( + compare_arrays, + run_child, +) + + +def test_science_acceptance_preserves_nan_and_row_identity() -> None: + expected = {"fit": np.array([[1.0, np.nan], [2.0, 3.0]])} + assert compare_arrays(expected, expected)["passed"] + reordered = {"fit": expected["fit"][::-1]} + assert not compare_arrays(expected, reordered)["passed"] + changed = {"fit": np.array([[1.0, 0.0], [2.0, 3.0]])} + assert not compare_arrays(expected, changed)["passed"] + rounded = {"fit": np.array([[1.0, np.nan], [2.0, 3.0 + 1e-12]])} + assert not compare_arrays(expected, rounded)["passed"] + + +@pytest.mark.parametrize( + "actual", + [ + {"wrong": np.array([1.0])}, + {"fit": np.array([[1.0]])}, + {"fit": np.array([1.0], dtype=np.float32)}, + ], +) +def test_science_rejects_schema_drift(actual: dict[str, np.ndarray]) -> None: + with pytest.raises(ValueError): + compare_arrays({"fit": np.array([1.0])}, actual) + + +def test_science_discrete_diagnostics_are_exact() -> None: + assert not compare_arrays( + {"valid": np.array([True])}, {"valid": np.array([False])} + )["passed"] + + +def test_child_failure_retains_diagnostics(tmp_path: Path) -> None: + log = tmp_path / "child.log" + with pytest.raises(subprocess.CalledProcessError) as raised: + run_child( + [ + sys.executable, + "-c", + "print('stage failed'); raise SystemExit(7)", + ], + log, + timeout=10, + ) + assert raised.value.returncode == 7 + assert "stage failed" in log.read_text() + + +def test_child_timeout_is_not_success(tmp_path: Path) -> None: + with pytest.raises(subprocess.TimeoutExpired): + run_child( + [sys.executable, "-c", "import time; time.sleep(20)"], + tmp_path / "timeout.log", + timeout=0.2, + ) + + +def test_failed_stage_round_retains_completed_work( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + failure = subprocess.CalledProcessError(7, ["xfit"]) + outcomes = iter([0.1, 0.2, 0.3, 0.4, failure]) + + def child(*args: object, **kwargs: object) -> float: + result = next(outcomes) + if isinstance(result, BaseException): + raise result + return result + + monkeypatch.setattr(benchmark, "run_child", child) + args = argparse.Namespace( + output=tmp_path, + config=tmp_path / "config.json", + items=tmp_path / "items.json", + warmup=1, + repeat=2, + timeout=10, + ) + with pytest.raises(subprocess.CalledProcessError) as raised: + benchmark.measure_stages(args) + assert raised.value is failure + root = tmp_path / "staged" + previous = json.loads((root / "round-000/timing.json").read_text()) + failed = json.loads((root / "round-001/timing.json").read_text()) + assert previous["status"] == "success" + assert failed["index"] == 1 and failed["phase"] == "measured" + assert failed["status"] == "failed" + assert failed["completed_stages"] == ["xpois"] + assert failed["stages"]["xpois"]["external_seconds"] == 0.4 + assert failed["stages"]["xfit"]["status"] == "failed" + assert ( + failed["batch_seconds"] + >= failed["stages"]["xfit"]["external_seconds"] + >= 0 + ) + assert failed["failure"]["type"] == "CalledProcessError" + assert not (root / "round-002").exists() + assert not (root / "summary.json").exists() + + +def test_failed_pipeline_round_retains_completed_items( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from cuphoton.xscan import device_pipeline + + # Replace CUDA bindings so the failure path can be checked on a CPU. + monkeypatch.setitem( + sys.modules, + "cupy", + SimpleNamespace( + cuda=SimpleNamespace( + Device=lambda _: SimpleNamespace(use=lambda: None), + runtime=SimpleNamespace(deviceSynchronize=lambda: None), + ) + ), + ) + monkeypatch.setitem( + sys.modules, + "torch", + SimpleNamespace( + set_num_threads=lambda _: None, + cuda=SimpleNamespace( + set_device=lambda _: None, synchronize=lambda _: None + ), + ), + ) + config = SimpleNamespace( + device="cuda:0", inference_policy={"worker_cpu_threads": 1} + ) + monkeypatch.setattr(benchmark, "read_inputs", lambda *_: (config, [0, 1])) + monkeypatch.setattr( + device_pipeline, + "DeviceWorkerContext", + SimpleNamespace(initialize=lambda _: object()), + ) + result = SimpleNamespace(to_payload=lambda: {"completed": True}) + failure = RuntimeError("injected item failure") + outcomes = iter([result, result, result, failure]) + + def run_item(*_: object) -> object: + value = next(outcomes) + if isinstance(value, BaseException): + raise value + return value + + monkeypatch.setattr(device_pipeline, "run_device_pipeline_item", run_item) + args = argparse.Namespace( + output=tmp_path, config=None, items=None, warmup=1, repeat=2 + ) + with pytest.raises(RuntimeError, match="injected item failure") as raised: + benchmark.pipeline_worker(args) + assert raised.value is failure + failed = json.loads((tmp_path / "round-001/timing.json").read_text()) + assert failed["index"] == 1 and failed["phase"] == "measured" + assert failed["status"] == "failed" and failed["completed_items"] == 1 + assert failed["batch_seconds"] >= 0 + assert failed["failure"]["type"] == "RuntimeError" + assert (tmp_path / "round-001/item-0000.json").is_file() + assert not (tmp_path / "round-001/item-0001.json").exists() + assert not (tmp_path / "round-002").exists() + assert not (tmp_path / "summary.json").exists() diff --git a/tests/xscan/test_pipeline_benchmark_stages.py b/tests/xscan/test_pipeline_benchmark_stages.py new file mode 100644 index 0000000..e6144f9 --- /dev/null +++ b/tests/xscan/test_pipeline_benchmark_stages.py @@ -0,0 +1,260 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""CPU rejection tests for file-mediated benchmark stage artifacts.""" + +from __future__ import annotations + +import builtins +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest + +from cuphoton.core.artifacts import file_sha256 +from cuphoton.xscan.pipeline_benchmark import stages + + +@pytest.fixture +def stage_files(tmp_path: Path) -> Path: + """Make two tiny typed records, not a scientific solver fixture.""" + items = [ + { + "item_id": f"pair-{index}", + "candidate_ids": [index * 2, index * 2 + 1], + } + for index in range(2) + ] + for stage in stages.STAGES: + directory = tmp_path / stage + directory.mkdir() + records = [] + for index, item in enumerate(items): + arrays = {} + for name in stages.SCIENCE_KEYS: + if not name.startswith(f"{stage}."): + continue + if name.endswith("converged"): + array = np.array([True, False], dtype=np.bool_) + elif name.endswith("evaluations"): + array = np.array([1, 7], dtype=np.int32) + elif name.endswith(("features", "logits", "probabilities")): + array = np.array([index + 0.25, 0.75], dtype=np.float32) + else: + array = np.array([index + 0.5, np.nan], dtype=np.float64) + arrays[name] = array + artifacts = stages._write_arrays( + tmp_path, directory / f"item-{index:04d}", arrays + ) + records.append({"item": item, "artifacts": artifacts}) + stages._write_json( + directory / "summary.json", + { + "schema": stages.SCHEMA, + "status": "success", + "stage": stage, + "configuration_sha256": "config-fixture", + "checkpoint_sha256": "checkpoint-fixture", + "feature_schema_sha256": "schema-fixture", + "items_sha256": stages._json_hash(items), + "items": records, + "upstream_summary_sha256": { + previous: file_sha256( + tmp_path / previous / "summary.json" + ) + for previous in stages.STAGES[ + : stages.STAGES.index(stage) + ] + }, + }, + ) + return tmp_path + + +def _edit_summary(root: Path, stage: str, change) -> None: + path = root / stage / "summary.json" + summary = json.loads(path.read_text()) + change(summary) + path.write_text(json.dumps(summary)) + + +def test_science_loads_all_arrays_with_float64_and_nan_masks( + stage_files: Path, +) -> None: + for index in range(2): + arrays = stages.load_science_arrays(stage_files, index) + assert len(arrays) == 22 + assert set(arrays) == set(stages.SCIENCE_KEYS) + for name, values in arrays.items(): + stage = name.split(".")[0] + summary = stages._read_summary(stage_files, stage) + artifact = summary["items"][index]["artifacts"][name] + source = stages._read_array(stage_files, artifact) + assert source.dtype.str == artifact["dtype"] + assert values.dtype == np.dtype("float64") + np.testing.assert_array_equal(values, source.astype(np.float64)) + np.testing.assert_array_equal(np.isnan(values), np.isnan(source)) + + +def test_last_item_prediction_tamper_is_rejected(stage_files: Path) -> None: + path = stage_files / "xscan/item-0001/xscan.probabilities.npy" + np.save(path, np.array([0.3, 0.4], dtype=np.float32)) + # Earlier occurrences remain valid; checking only the first hides this. + stages.load_science_arrays(stage_files, 0) + with pytest.raises(ValueError, match="artifact hash changed"): + stages.load_science_arrays(stage_files, 1) + + +@pytest.mark.parametrize("kind", ["order", "candidate", "config"]) +def test_changed_stage_identity_is_rejected( + stage_files: Path, kind: str +) -> None: + def change(summary): + if kind == "order": + summary["items"].reverse() + elif kind == "candidate": + summary["items"][1]["item"]["candidate_ids"].reverse() + else: + summary["configuration_sha256"] = "different-config" + + _edit_summary(stage_files, "xfit", change) + with pytest.raises(ValueError, match="metadata or order|identities"): + stages.load_science_arrays(stage_files, 0) + + +def test_completed_upstream_summary_change_is_rejected( + stage_files: Path, +) -> None: + _edit_summary( + stage_files, + "xpois", + lambda summary: summary.update(extra_changed_metadata=1), + ) + with pytest.raises(ValueError, match="upstream summary changed"): + stages.load_science_arrays(stage_files, 0) + + +def test_missing_upstream_link_cannot_skip_chain_check( + stage_files: Path, +) -> None: + _edit_summary( + stage_files, + "xscan", + lambda summary: summary["upstream_summary_sha256"].pop("xfit"), + ) + with pytest.raises(ValueError, match="links are incomplete"): + stages.load_science_arrays(stage_files, 0) + + +@pytest.mark.parametrize("field,value", [("dtype", " None: + summary = stages._read_summary(stage_files, "xscan") + artifact = summary["items"][0]["artifacts"]["xscan.logits"] + artifact[field] = value + with pytest.raises(ValueError, match="artifact contract changed"): + stages._read_array(stage_files, artifact) + + +def test_synchronized_timer_waits_for_completion(monkeypatch) -> None: + calls = [] + ticks = iter([10.0, 13.0]) + monkeypatch.setattr(stages.time, "perf_counter", lambda: next(ticks)) + timings = {} + result = stages._measure( + timings, + "compute", + lambda: calls.append("work") or "result", + lambda: calls.append("sync"), + ) + assert result == "result" + assert calls == ["sync", "work", "sync"] + assert timings == {"compute": 3.0} + + +@pytest.mark.parametrize("stage", stages.STAGES) +def test_setup_imports_only_its_numerical_runtime( + tmp_path, monkeypatch, stage +) -> None: + from cuphoton.xscan.config import PerformanceConfig + + checkpoint = tmp_path / "checkpoint.pt" + schema = tmp_path / "schema.json" + checkpoint.write_bytes(b"checkpoint fixture") + schema.write_text("{}") + config = SimpleNamespace( + checkpoint_dir=str(tmp_path), + checkpoint_sha256=file_sha256(checkpoint), + feature_schema_path=str(schema), + feature_schema_sha256=file_sha256(schema), + stamp_shape=(17, 17), + inference_policy=stages.pipeline._strict_inference_policy_payload(), + device_id=0, + device="cuda:0", + ) + for name in ( + "_validate_feature_schema_source", + "_validate_checkpoint_contract", + "_validate_loaded_model_contract", + ): + monkeypatch.setattr(stages.pipeline, name, lambda *a, **kw: None) + monkeypatch.setattr( + stages.pipeline, "_load_feature_schema_contract", lambda *a, **kw: {} + ) + calls = [] + fake_torch = SimpleNamespace( + device=lambda name: name, + set_num_threads=lambda count: calls.append(("threads", count)), + cuda=SimpleNamespace( + set_device=lambda index: calls.append(("torch", index)), + synchronize=lambda device: calls.append(("sync", device)), + ), + ) + fake_cupy = SimpleNamespace( + cuda=SimpleNamespace( + Device=lambda index: SimpleNamespace( + use=lambda: calls.append(("cupy", index)), + synchronize=lambda: calls.append(("sync", index)), + ) + ) + ) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + monkeypatch.setitem(sys.modules, "cupy", fake_cupy) + monkeypatch.setitem( + sys.modules, + "cuphoton.xscan.training", + SimpleNamespace( + load_model_from_checkpoint=lambda *a, **kw: ( + object(), + {}, + PerformanceConfig(**config.inference_policy), + ) + ), + ) + original_import = builtins.__import__ + forbidden = ( + {"cupy"} if stage == "xscan" else {"torch", "cuphoton.xscan.training"} + ) + + def guarded_import(name, *args, **kwargs): + if any( + name == item or name.startswith(item + ".") for item in forbidden + ): + pytest.fail(f"{stage} imported unused runtime {name}") + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", guarded_import) + runtime = stages._setup(stage, config) + if stage == "xscan": + assert runtime["cp"] is None + assert runtime["torch"] is fake_torch + assert calls == [("threads", 1), ("torch", 0), ("sync", "cuda:0")] + else: + assert runtime["torch"] is None + assert runtime["cp"] is fake_cupy + assert calls == [("cupy", 0), ("sync", 0)] From 8eaba75527343d4c9ffd9e75211595a6bf80f8cd Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 23 Sep 2026 21:49:15 -0700 Subject: [PATCH 2/3] Route pipeline benchmarks through the shared CLI Signed-off-by: Trent Nelson --- docs/components/pipeline-stage-benchmark.md | 9 +- src/cuphoton/xscan/commands.py | 132 ++++++++++++++++ .../{__main__.py => runner.py} | 71 +++------ .../xscan/pipeline_benchmark/stages.py | 25 --- tests/core/test_cli_contract.py | 10 +- tests/xscan/test_pipeline_benchmark.py | 4 +- tests/xscan/test_pipeline_benchmark_cli.py | 149 ++++++++++++++++++ 7 files changed, 321 insertions(+), 79 deletions(-) rename src/cuphoton/xscan/pipeline_benchmark/{__main__.py => runner.py} (90%) create mode 100644 tests/xscan/test_pipeline_benchmark_cli.py diff --git a/docs/components/pipeline-stage-benchmark.md b/docs/components/pipeline-stage-benchmark.md index e1f1f76..d735740 100644 --- a/docs/components/pipeline-stage-benchmark.md +++ b/docs/components/pipeline-stage-benchmark.md @@ -18,7 +18,7 @@ root: ```bash uv sync --locked --extra gpu -uv run --locked --extra gpu python -m cuphoton.xscan.pipeline_benchmark \ +uv run --locked --extra gpu cuphoton xscan benchmark-pipeline \ --output /tmp/cuphoton-pipeline-forward \ --images 4 --image-size 256 --candidates 9 --stamp-size 17 \ --seed 2026 --device cuda:0 --warmup 1 --repeat 3 \ @@ -34,7 +34,7 @@ run serially on the selected GPU. Reverse the treatment order using the *same generated fixture*: ```bash -uv run --locked --extra gpu python -m cuphoton.xscan.pipeline_benchmark \ +uv run --locked --extra gpu cuphoton xscan benchmark-pipeline \ --output /tmp/cuphoton-pipeline-reverse \ --config /tmp/cuphoton-pipeline-forward/input/config.json \ --items /tmp/cuphoton-pipeline-forward/input/items.json \ @@ -68,6 +68,11 @@ scientific outputs shared with the pipeline. Their transfers, hashing and I/O are part of this treatment's cost. Its preliminary rounds prepare caches; subsequent measured rounds still create fresh processes. +To inspect a manual round directly, invoke the command three times with +`--stage xpois`, `--stage xfit` and `--stage xscan`, in that order. Supply the +same `--config`, `--items` and `--output` directory each time. The stage +commands reject changed upstream artifacts and refuse to overwrite a stage. + Both treatments use float64 subtraction and xFit inputs, unweighted xFit stamps, float32 triplets/features, and the same GPU sigmoid. Inference forces AMP, TF32, compilation and cuDNN benchmarking off. An xPOIS variance plane diff --git a/src/cuphoton/xscan/commands.py b/src/cuphoton/xscan/commands.py index 22669eb..6169f9e 100644 --- a/src/cuphoton/xscan/commands.py +++ b/src/cuphoton/xscan/commands.py @@ -7,14 +7,17 @@ from __future__ import annotations import json +import math import os from pathlib import Path +from types import SimpleNamespace from typing import Any, Callable, TypeVar from cuphoton.core.bulk import validate_identifier from cuphoton.core.cli import ( BoolInvariant, CommandError, + FloatInvariant, InvariantAwareCommand, NonNegativeIntegerInvariant, PositiveIntegerInvariant, @@ -2156,3 +2159,132 @@ def run(self) -> None: run_name=self.run_name or None, ) self._emit_json({"run_dir": str(result.run_dir), **result.summary}) + + +class BenchmarkPipelineCommand(XScanCommand): + """Compare the device pipeline with separate file-mediated stages.""" + + _name_ = "benchmark-pipeline" + + output = None + config = None + items = None + images = None + image_size = None + candidates = None + stamp_size = None + seed = None + device = None + warmup = None + repeat = None + timeout = None + order = None + stage = None + + class OutputArg(PathSpecInvariant): + _arg = "--output" + _help = "New output directory for benchmark receipts and artifacts." + _mandatory = True + + class ConfigArg(PathSpecInvariant): + _arg = "--config" + _help = "Existing pipeline config JSON; requires --items." + _default = None + + class ItemsArg(PathSpecInvariant): + _arg = "--items" + _help = "Existing pipeline items JSON; requires --config." + _default = None + + class ImagesArg(PositiveIntegerInvariant): + _arg = "--images" + _help = "Synthetic image-pair count. [default: %default]" + _default = 4 + + class ImageSizeArg(PositiveIntegerInvariant): + _arg = "--image-size" + _help = "Synthetic square image width in pixels. [default: %default]" + _default = 256 + + class CandidatesArg(PositiveIntegerInvariant): + _arg = "--candidates" + _help = "Candidates per synthetic image. [default: %default]" + _default = 9 + + class StampSizeArg(PositiveIntegerInvariant): + _arg = "--stamp-size" + _help = "Square candidate stamp width in pixels. [default: %default]" + _default = 17 + + class SeedArg(NonNegativeIntegerInvariant): + _arg = "--seed" + _help = "Synthetic fixture random seed. [default: %default]" + _default = 2026 + + class DeviceArg(StringInvariant): + _arg = "--device" + _help = "CUDA device for the synthetic fixture. [default: %default]" + _default = "cuda:0" + + class WarmupArg(PositiveIntegerInvariant): + _arg = "--warmup" + _help = "Recorded warmup rounds per treatment. [default: %default]" + _default = 1 + + class RepeatArg(PositiveIntegerInvariant): + _arg = "--repeat" + _help = "Measured rounds per treatment. [default: %default]" + _default = 3 + + class TimeoutArg(FloatInvariant): + _arg = "--timeout" + _help = ( + "Positive child-process timeout in seconds. [default: %default]" + ) + _default = 600.0 + + @classmethod + def validate(cls, value: Any) -> float | None: + converted = super().validate(value) + if converted is not None and ( + not math.isfinite(converted) or converted <= 0 + ): + raise ValueError("must be finite and greater than zero") + return converted + + class OrderArg(SetInvariant): + _arg = "--order" + _help = "Treatment launch order. [default: %default]" + _set = {"pipeline-first", "staged-first"} + _default = "pipeline-first" + + class StageArg(SetInvariant): + _arg = "--stage" + _help = ( + "Comparison or individual subprocess stage. [default: %default]" + ) + _set = {"compare", "pipeline", "xpois", "xfit", "xscan"} + _default = "compare" + + def run(self) -> None: + from .pipeline_benchmark.runner import run_benchmark + + self._call( + run_benchmark, + SimpleNamespace( + output=Path(self.output).expanduser(), + config=self._path(self.config), + items=self._path(self.items), + images=self.images, + image_size=self.image_size, + candidates=self.candidates, + stamp_size=self.stamp_size, + seed=self.seed, + device=self.device, + warmup=self.warmup, + repeat=self.repeat, + timeout=self.timeout, + order=self.order, + stage=self.stage, + ), + ) diff --git a/src/cuphoton/xscan/pipeline_benchmark/__main__.py b/src/cuphoton/xscan/pipeline_benchmark/runner.py similarity index 90% rename from src/cuphoton/xscan/pipeline_benchmark/__main__.py rename to src/cuphoton/xscan/pipeline_benchmark/runner.py index 4942261..1e2e2ee 100644 --- a/src/cuphoton/xscan/pipeline_benchmark/__main__.py +++ b/src/cuphoton/xscan/pipeline_benchmark/runner.py @@ -6,7 +6,6 @@ from __future__ import annotations -import argparse import importlib.metadata import json import os @@ -17,11 +16,12 @@ import sys import time from pathlib import Path +from types import SimpleNamespace from typing import Any import numpy as np -MODULE = "cuphoton.xscan.pipeline_benchmark" +CLI = [sys.executable, "-m", "cuphoton", "xscan", "benchmark-pipeline"] def write_json(path: Path, value: Any) -> None: @@ -104,7 +104,7 @@ def run_child(command: list[str], log: Path, *, timeout: float) -> float: return time.perf_counter() - started -def pipeline_worker(args: argparse.Namespace) -> None: +def pipeline_worker(args: SimpleNamespace) -> None: started = time.perf_counter() import cupy as cp import torch @@ -178,15 +178,14 @@ def pipeline_worker(args: argparse.Namespace) -> None: ) -def measure_pipeline(args: argparse.Namespace) -> dict[str, Any]: +def measure_pipeline(args: SimpleNamespace) -> dict[str, Any]: output = args.output / "pipeline" output.mkdir() elapsed = run_child( [ - sys.executable, - "-m", - MODULE, - "--pipeline-worker", + *CLI, + "--stage", + "pipeline", "--config", str(args.config), "--items", @@ -206,7 +205,7 @@ def measure_pipeline(args: argparse.Namespace) -> dict[str, Any]: return result -def measure_stages(args: argparse.Namespace) -> dict[str, Any]: +def measure_stages(args: SimpleNamespace) -> dict[str, Any]: output = args.output / "staged" output.mkdir() rounds = [] @@ -226,16 +225,14 @@ def measure_stages(args: argparse.Namespace) -> dict[str, Any]: stage_started = time.perf_counter() elapsed = run_child( [ - sys.executable, - "-m", - f"{MODULE}.stages", + *CLI, "--stage", stage, "--config", str(args.config), "--items", str(args.items), - "--output-dir", + "--output", str(root), ], root / f"{stage}.log", @@ -274,7 +271,7 @@ def measure_stages(args: argparse.Namespace) -> dict[str, Any]: return result -def audit(args: argparse.Namespace) -> dict[str, Any]: +def audit(args: SimpleNamespace) -> dict[str, Any]: from cuphoton.core.artifacts import file_sha256 from cuphoton.xscan.device_pipeline import ( _verify_item_hashes, @@ -413,36 +410,24 @@ def provenance() -> dict[str, Any]: } -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--output", type=Path, required=True) - parser.add_argument("--config", type=Path) - parser.add_argument("--items", type=Path) - parser.add_argument("--images", type=int, default=4) - parser.add_argument("--image-size", type=int, default=256) - parser.add_argument("--candidates", type=int, default=9) - parser.add_argument("--stamp-size", type=int, default=17) - parser.add_argument("--seed", type=int, default=2026) - parser.add_argument("--device", default="cuda:0") - parser.add_argument("--warmup", type=int, default=1) - parser.add_argument("--repeat", type=int, default=3) - parser.add_argument("--timeout", type=float, default=600) - parser.add_argument( - "--order", - choices=("pipeline-first", "staged-first"), - default="pipeline-first", - ) - parser.add_argument( - "--pipeline-worker", action="store_true", help=argparse.SUPPRESS - ) - args = parser.parse_args() +def run_benchmark(args: SimpleNamespace) -> None: + """Run a comparison or one child stage through the shared CLI adapter.""" if args.repeat < 1 or args.warmup < 1 or args.timeout <= 0: - parser.error("repeat, warmup and timeout must be positive") + raise ValueError("repeat, warmup and timeout must be positive") if (args.config is None) != (args.items is None): - parser.error("provide both --config and --items, or neither") + raise ValueError("provide both --config and --items, or neither") args.output = args.output.resolve() - if args.pipeline_worker: - pipeline_worker(args) + if args.stage != "compare": + if args.config is None: + raise ValueError("child stages require --config and --items") + if args.stage == "pipeline": + args.output.mkdir(parents=True, exist_ok=True) + pipeline_worker(args) + else: + from .stages import run_stage + + config, items = read_inputs(args.config, args.items) + run_stage(args.stage, config, items, args.output) return args.output.mkdir(parents=True, exist_ok=False) try: @@ -512,7 +497,3 @@ def main() -> None: }, ) raise - - -if __name__ == "__main__": - main() diff --git a/src/cuphoton/xscan/pipeline_benchmark/stages.py b/src/cuphoton/xscan/pipeline_benchmark/stages.py index 6a373c3..ed46236 100644 --- a/src/cuphoton/xscan/pipeline_benchmark/stages.py +++ b/src/cuphoton/xscan/pipeline_benchmark/stages.py @@ -12,7 +12,6 @@ from __future__ import annotations -import argparse import hashlib import json import time @@ -532,27 +531,3 @@ def load_science_arrays(root: Path, index: int) -> dict[str, np.ndarray]: ).astype(np.float64) for name in SCIENCE_KEYS } - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--stage", required=True, choices=STAGES) - parser.add_argument("--config", required=True, type=Path) - parser.add_argument("--items", required=True, type=Path) - parser.add_argument("--output-dir", required=True, type=Path) - args = parser.parse_args() - config = pipeline.DevicePipelineConfig.from_payload( - json.loads(args.config.read_text()) - ) - items = [ - pipeline.DevicePipelineItem.from_payload(item) - for item in json.loads(args.items.read_text()) - ] - summary = run_stage(args.stage, config, items, args.output_dir) - print( - json.dumps({"stage": summary["stage"], "status": summary["status"]}) - ) - - -if __name__ == "__main__": - main() diff --git a/tests/core/test_cli_contract.py b/tests/core/test_cli_contract.py index 74d369d..2d36e66 100644 --- a/tests/core/test_cli_contract.py +++ b/tests/core/test_cli_contract.py @@ -73,15 +73,15 @@ def test_public_command_surface_counts_are_exact() -> None: ("xdr", 1, 0, 14, 0), ("xfit", 3, 3, 25, 1), ("xpois", 7, 7, 146, 1), - ("xscan", 43, 43, 190, 1), + ("xscan", 44, 44, 204, 1), ("xrep", 6, 6, 103, 1), ("xray", 33, 31, 390, 1), ] assert len(per_group) == 6 - assert sum(item[1] for item in per_group) == 93 - assert sum(item[1] + item[4] for item in per_group) == 98 - assert sum(item[2] for item in per_group) == 90 - assert sum(item[3] for item in per_group) == 868 + assert sum(item[1] for item in per_group) == 94 + assert sum(item[1] + item[4] for item in per_group) == 99 + assert sum(item[2] for item in per_group) == 91 + assert sum(item[3] for item in per_group) == 882 def test_public_registry_order_and_component_derivations() -> None: diff --git a/tests/xscan/test_pipeline_benchmark.py b/tests/xscan/test_pipeline_benchmark.py index 1466f4b..5bcd8c7 100644 --- a/tests/xscan/test_pipeline_benchmark.py +++ b/tests/xscan/test_pipeline_benchmark.py @@ -16,8 +16,8 @@ import numpy as np import pytest -from cuphoton.xscan.pipeline_benchmark import __main__ as benchmark -from cuphoton.xscan.pipeline_benchmark.__main__ import ( +from cuphoton.xscan.pipeline_benchmark import runner as benchmark +from cuphoton.xscan.pipeline_benchmark.runner import ( compare_arrays, run_child, ) diff --git a/tests/xscan/test_pipeline_benchmark_cli.py b/tests/xscan/test_pipeline_benchmark_cli.py new file mode 100644 index 0000000..dccc1f7 --- /dev/null +++ b/tests/xscan/test_pipeline_benchmark_cli.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only checks for the benchmark's shared CLI adapter.""" + +from pathlib import Path + +import pytest + +from cuphoton.core.cli import run_component +from cuphoton.xscan.pipeline_benchmark import runner + + +def test_benchmark_help(capsys): + assert run_component("xscan", ["help", "benchmark-pipeline"]) == 0 + output = capsys.readouterr().out + assert "--output" in output + assert "--config" in output + assert "--stage" in output + + +def test_benchmark_defaults_and_path_conversion(monkeypatch): + received = [] + monkeypatch.setattr(runner, "run_benchmark", received.append) + assert ( + run_component("xscan", ["benchmark-pipeline", "--output", "~/bench"]) + == 0 + ) + (args,) = received + assert vars(args) == { + "output": Path("~/bench").expanduser(), + "config": None, + "items": None, + "images": 4, + "image_size": 256, + "candidates": 9, + "stamp_size": 17, + "seed": 2026, + "device": "cuda:0", + "warmup": 1, + "repeat": 3, + "timeout": 600.0, + "order": "pipeline-first", + "stage": "compare", + } + + +def test_benchmark_child_options(monkeypatch): + received = [] + monkeypatch.setattr(runner, "run_benchmark", received.append) + assert ( + run_component( + "xscan", + [ + "benchmark-pipeline", + "--output", + "result", + "--config", + "~/config.json", + "--items", + "~/items.json", + "--stage", + "xfit", + "--order", + "staged-first", + "--timeout", + "12.5", + "--images", + "2", + "--image-size", + "1024", + "--candidates", + "16", + "--stamp-size", + "31", + "--seed", + "0", + "--device", + "cuda:1", + "--warmup", + "2", + "--repeat", + "4", + ], + ) + == 0 + ) + (args,) = received + assert args.output == Path("result") + assert args.config == Path("~/config.json").expanduser() + assert args.items == Path("~/items.json").expanduser() + assert (args.stage, args.order, args.timeout) == ( + "xfit", + "staged-first", + 12.5, + ) + assert ( + args.images, + args.image_size, + args.candidates, + args.stamp_size, + ) == (2, 1024, 16, 31) + assert (args.seed, args.device, args.warmup, args.repeat) == ( + 0, + "cuda:1", + 2, + 4, + ) + + +@pytest.mark.parametrize("timeout", ["0", "-1", "nan", "inf"]) +def test_benchmark_rejects_invalid_timeout(monkeypatch, timeout): + received = [] + monkeypatch.setattr(runner, "run_benchmark", received.append) + assert ( + run_component( + "xscan", + [ + "benchmark-pipeline", + "--output", + "result", + "--timeout", + timeout, + ], + ) + != 0 + ) + assert not received + + +def test_benchmark_reports_unpaired_inputs(tmp_path, capsys): + output = tmp_path / "result" + assert ( + run_component( + "xscan", + [ + "benchmark-pipeline", + "--output", + str(output), + "--config", + str(tmp_path / "config.json"), + ], + ) + == 1 + ) + captured = capsys.readouterr() + assert "provide both --config and --items" in captured.err + captured.out + assert not output.exists() From c8087535ced54e87098af763433dcf55c570e211 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Thu, 24 Sep 2026 17:44:45 -0700 Subject: [PATCH 3/3] Separate benchmark verification cost and record device identity Signed-off-by: Trent Nelson --- docs/components/pipeline-stage-benchmark.md | 32 ++- .../xscan/pipeline_benchmark/runner.py | 63 ++++- .../xscan/pipeline_benchmark/stages.py | 56 +++-- tests/xscan/test_pipeline_benchmark.py | 225 ++++++++++++++++++ tests/xscan/test_pipeline_benchmark_stages.py | 16 ++ 5 files changed, 365 insertions(+), 27 deletions(-) diff --git a/docs/components/pipeline-stage-benchmark.md b/docs/components/pipeline-stage-benchmark.md index d735740..bc03f97 100644 --- a/docs/components/pipeline-stage-benchmark.md +++ b/docs/components/pipeline-stage-benchmark.md @@ -62,11 +62,11 @@ evidence and predictions. The separate-stage treatment starts one xPOIS child, one xFit child and one XScan child per complete round. Each child processes all image items in manifest order, keeping the candidate batch for each image unchanged. It -writes lossless, uncompressed NPY arrays between stages. Those files include -full subtraction images and fit residuals in addition to the compact -scientific outputs shared with the pipeline. Their transfers, hashing and -I/O are part of this treatment's cost. Its preliminary rounds prepare -caches; subsequent measured rounds still create fresh processes. +writes lossless, uncompressed NPY arrays between stages: extracted difference +stamps for xFit, triplets and features for XScan, and the compact scientific +outputs shared with the pipeline. Unused full subtraction images, basis +kernels and fit residuals are not transferred or written. Its preliminary +rounds prepare caches; subsequent measured rounds still create fresh processes. To inspect a manual round directly, invoke the command three times with `--stage xpois`, `--stage xfit` and `--stage xscan`, in that order. Supply the @@ -86,6 +86,8 @@ does not become an xFit variance plane. | Pipeline `batch_seconds` | One ordered image batch through completed device work and per-item result JSON writes. Measured batches reuse the initialized, warmed worker. | | Pipeline `invocation_external_seconds` | Parent-observed process duration, including imports, setup, all warmup and measured rounds, final summary and process exit. | | Staged `batch_seconds` | Parent clock before launching xPOIS through successful XScan process exit, including all three process lifetimes and intermediate files. | +| Staged `extra_hashing_seconds` | Intermediate-artifact SHA-256 reads/writes and original-input rechecks repeated by xFit/XScan, measured inside the raw batch timer. | +| Staged `batch_seconds_without_extra_hashing` | Raw batch time minus that additional verification time; retains process startup, numerical work, transfers and file I/O. | | Staged per-process `external_seconds` | Parent-observed duration of that individual child, including imports and shutdown. | Stage summaries also record internal setup, reads, uploads, computation, @@ -96,8 +98,13 @@ timers as isolated GPU kernel durations would be misleading. Use the completed batch timers for the main workflow comparison. The report's `fresh_stages_over_warm_pipeline_ratio` divides the median -staged batch time by the median warm pipeline batch time. It includes -repeated startup and file costs in the staged treatment. Keep setup and +staged batch time **without extra hashing** by the median warm pipeline batch +time. Both arms still include their common original-input hashes on read and +once after execution; checkpoint/schema verification remains in setup. All +verification still executes. The adjustment subtracts measured hash time; it +is not a separate hash-disabled run and cannot undo cache effects caused by +verification. `raw_fresh_stages_over_warm_pipeline_ratio` preserves the ratio +of unadjusted medians. Both include repeated startup and file costs. Keep setup and whole-invocation measurements alongside that ratio when discussing a service that may process only a few batches. @@ -114,14 +121,17 @@ arrays, including shapes, dtypes and NaN locations, plus matching candidate identity/order, fit metadata, subtraction diagnostics and predictions. It verifies retained file hashes and rechecks the original inputs. This checks equivalence between two compositions of the same algorithms; it does not -independently establish their astronomical accuracy. Full intermediate -images are retained by the staged treatment but are outside the pipeline's -compact parity contract. +independently establish their astronomical accuracy. Only arrays needed by +the next stage or by this compact parity contract are retained. `manifest.json` records the configuration, ordered input descriptors and runtime settings. Treatment subdirectories retain each attempted round and child log. `parity.json` records numerical comparisons; a successful -`report.json` contains all round times and the measured medians. A failure +`report.json` contains all round times and the measured medians. Report-level +provenance includes the selected GPU name/UUID, NVIDIA driver version, CUDA +driver API/runtime versions and the imported CuPy version. Package discovery +also records the installed distribution providing `cupy`. Driver-query +failures are retained explicitly rather than silently omitting the field. A failure stops the run and leaves its logs and `failure.json` for inspection. A successful process exit alone does not satisfy the parity check. diff --git a/src/cuphoton/xscan/pipeline_benchmark/runner.py b/src/cuphoton/xscan/pipeline_benchmark/runner.py index 1e2e2ee..242c5d8 100644 --- a/src/cuphoton/xscan/pipeline_benchmark/runner.py +++ b/src/cuphoton/xscan/pipeline_benchmark/runner.py @@ -165,6 +165,9 @@ def pipeline_worker(args: SimpleNamespace) -> None: rounds.append(record) write_json(output / "timing.json", record) print(json.dumps(record), flush=True) + gpu_uuid = str(torch.cuda.get_device_properties(ordinal).uuid) + if not gpu_uuid.startswith(("GPU-", "MIG-")): + gpu_uuid = "GPU-" + gpu_uuid write_json( args.output / "summary.json", { @@ -174,6 +177,14 @@ def pipeline_worker(args: SimpleNamespace) -> None: "configuration_sha256": config.configuration_sha256, "rounds": rounds, "gpu_name": torch.cuda.get_device_name(ordinal), + "device_provenance": { + "device": config.device, + "gpu_name": torch.cuda.get_device_name(ordinal), + "gpu_uuid": gpu_uuid, + "cuda_driver_api_version": cp.cuda.runtime.driverGetVersion(), + "cuda_runtime_version": cp.cuda.runtime.runtimeGetVersion(), + "cupy_version": cp.__version__, + }, }, ) @@ -263,6 +274,15 @@ def measure_stages(args: SimpleNamespace) -> dict[str, Any]: raise record["batch_seconds"] = time.perf_counter() - started record["status"] = "success" + record["extra_hashing_seconds"] = sum( + json.loads((root / stage / "summary.json").read_text())[ + "extra_hashing_seconds" + ] + for stage in ("xpois", "xfit", "xscan") + ) + record["batch_seconds_without_extra_hashing"] = ( + record["batch_seconds"] - record["extra_hashing_seconds"] + ) rounds.append(record) write_json(root / "timing.json", record) print(json.dumps({"treatment": "staged", **record}), flush=True) @@ -380,7 +400,10 @@ def provenance() -> dict[str, Any]: "xscan/model.py", ) packages = {} - for name in ("cuphoton", "numpy", "cupy-cuda13x", "torch"): + cupy_packages = importlib.metadata.packages_distributions().get( + "cupy", [] + ) + for name in ("cuphoton", "numpy", "torch", *cupy_packages): try: packages[name] = importlib.metadata.version(name) except importlib.metadata.PackageNotFoundError: @@ -478,10 +501,46 @@ def run_benchmark(args: SimpleNamespace) -> None: report[name]["measured_batch_median_seconds"] = statistics.median( values ) - report["fresh_stages_over_warm_pipeline_ratio"] = ( + report["staged"][ + "measured_batch_median_seconds_without_extra_hashing" + ] = statistics.median( + row["batch_seconds_without_extra_hashing"] + for row in report["staged"]["rounds"] + if row["phase"] == "measured" + ) + report["provenance"]["device"] = report["pipeline"][ + "device_provenance" + ] + try: + driver = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=driver_version", + "--format=csv,noheader", + "--id=" + report["provenance"]["device"]["gpu_uuid"], + ], + check=True, + capture_output=True, + text=True, + timeout=10, + ) + except (OSError, subprocess.SubprocessError) as exc: + report["provenance"]["device"]["driver_version"] = None + report["provenance"]["device"]["driver_query_error"] = str(exc) + else: + report["provenance"]["device"]["driver_version"] = ( + driver.stdout.strip() + ) + report["raw_fresh_stages_over_warm_pipeline_ratio"] = ( report["staged"]["measured_batch_median_seconds"] / report["pipeline"]["measured_batch_median_seconds"] ) + report["fresh_stages_over_warm_pipeline_ratio"] = ( + report["staged"][ + "measured_batch_median_seconds_without_extra_hashing" + ] + / report["pipeline"]["measured_batch_median_seconds"] + ) write_json(args.output / "report.json", report) print( json.dumps( diff --git a/src/cuphoton/xscan/pipeline_benchmark/stages.py b/src/cuphoton/xscan/pipeline_benchmark/stages.py index ed46236..eb4f361 100644 --- a/src/cuphoton/xscan/pipeline_benchmark/stages.py +++ b/src/cuphoton/xscan/pipeline_benchmark/stages.py @@ -91,11 +91,29 @@ def _read_summary(root: Path, stage: str) -> dict[str, Any]: return value -def _read_array(root: Path, descriptor: dict[str, Any]) -> np.ndarray: +def _artifact_hash( + path: Path, timings: dict[str, float] | None = None +) -> str: + started = time.perf_counter() + digest = file_sha256(path) + if timings is not None: + timings["artifact_hash_seconds"] = ( + timings.get("artifact_hash_seconds", 0.0) + + time.perf_counter() + - started + ) + return digest + + +def _read_array( + root: Path, + descriptor: dict[str, Any], + timings: dict[str, float] | None = None, +) -> np.ndarray: path = (root / descriptor["path"]).resolve() if not path.is_relative_to(root.resolve()): raise ValueError("intermediate artifact escapes its round directory") - if file_sha256(path) != descriptor["sha256"]: + if _artifact_hash(path, timings) != descriptor["sha256"]: raise ValueError(f"intermediate artifact hash changed: {path.name}") array = np.load(path, allow_pickle=False) if ( @@ -108,7 +126,7 @@ def _read_array(root: Path, descriptor: dict[str, Any]) -> np.ndarray: raise ValueError( f"intermediate artifact contract changed: {path.name}" ) - if file_sha256(path) != descriptor["sha256"]: + if _artifact_hash(path, timings) != descriptor["sha256"]: raise ValueError( f"intermediate artifact changed while reading: {path.name}" ) @@ -116,7 +134,10 @@ def _read_array(root: Path, descriptor: dict[str, Any]) -> np.ndarray: def _write_arrays( - root: Path, directory: Path, arrays: dict[str, np.ndarray] + root: Path, + directory: Path, + arrays: dict[str, np.ndarray], + timings: dict[str, float] | None = None, ) -> dict[str, Any]: directory.mkdir() artifacts = {} @@ -129,7 +150,7 @@ def _write_arrays( np.save(stream, array, allow_pickle=False) artifacts[name] = { "path": str(path.relative_to(root)), - "sha256": file_sha256(path), + "sha256": _artifact_hash(path, timings), "dtype": array.dtype.str, "shape": list(array.shape), "nbytes": array.nbytes, @@ -219,7 +240,7 @@ def sync() -> None: def previous(which: str, name: str) -> np.ndarray: return _read_array( - root, priors[which]["items"][index]["artifacts"][name] + root, priors[which]["items"][index]["artifacts"][name], timings ) if stage == "xpois": @@ -268,14 +289,8 @@ def compute() -> tuple[dict[str, Any], dict[str, Any]]: arrays = { f"xpois.{name}": getattr(result, name) for name in ( - "kernel", - "matched", - "residual", - "fit_mask", - "background", "kernel_coefficients", "background_coefficients", - "basis_kernels", ) } arrays.update(stamps=stamps, difference=difference) @@ -318,7 +333,7 @@ def compute() -> tuple[dict[str, Any], dict[str, Any]]: return { **{ f"xfit.{name}": getattr(result, name) - for name in (*XFIT_FIELDS, "residuals") + for name in XFIT_FIELDS }, "xfit.features": features.values, }, { @@ -460,6 +475,7 @@ def run_stage( root, stage_dir / f"item-{index:04d}", arrays, + timings, ), ) _measure( @@ -470,6 +486,11 @@ def run_stage( when=f"during {stage} execution", ), ) + # The pipeline and xPOIS both hash original inputs on read and once + # after execution. Only the later stages' rechecks are additional. + timings["extra_hashing_seconds"] = timings.get( + "artifact_hash_seconds", 0.0 + ) + (timings["input_recheck_seconds"] if stage != "xpois" else 0.0) timings["wall_seconds"] = time.perf_counter() - item_started records.append( { @@ -492,11 +513,18 @@ def run_stage( name: file_sha256(root / name / "summary.json") for name in priors }, "items": records, + "extra_hashing_seconds": sum( + record["timings_seconds"]["extra_hashing_seconds"] + for record in records + ), "timing_note": ( "Setup starts inside run_stage; parent process wall includes " "imports and shutdown. Compute/upload/download synchronize " "explicitly; these are host elapsed, not GPU-only. Wall excludes " - "this final summary write. NPY writes close without fsync." + "this final summary write. NPY writes close without fsync. " + "Extra hashing counts intermediate digests and original-input " + "rechecks in xFit/xScan; input hashes shared with the pipeline " + "remain in the adjusted comparison." ), } _write_json(stage_dir / "summary.json", summary) diff --git a/tests/xscan/test_pipeline_benchmark.py b/tests/xscan/test_pipeline_benchmark.py index 5bcd8c7..431bd2d 100644 --- a/tests/xscan/test_pipeline_benchmark.py +++ b/tests/xscan/test_pipeline_benchmark.py @@ -88,6 +88,13 @@ def child(*args: object, **kwargs: object) -> float: result = next(outcomes) if isinstance(result, BaseException): raise result + command = args[0] + stage = command[command.index("--stage") + 1] + root = Path(command[command.index("--output") + 1]) / stage + root.mkdir() + (root / "summary.json").write_text( + json.dumps({"extra_hashing_seconds": 0.0}) + ) return result monkeypatch.setattr(benchmark, "run_child", child) @@ -182,3 +189,221 @@ def run_item(*_: object) -> object: assert not (tmp_path / "round-001/item-0001.json").exists() assert not (tmp_path / "round-002").exists() assert not (tmp_path / "summary.json").exists() + + +@pytest.fixture +def benchmark_fixture(tmp_path: Path): + from cuphoton.xscan.pipeline_benchmark.fixture import prepare_fixture + + return prepare_fixture( + tmp_path / "inputs", + images=1, + image_size=48, + candidates=2, + stamp_size=9, + seed=2026, + device="cuda:0", + ) + + +def test_prepare_fixture_has_repeatable_inputs_and_cpu_model( + tmp_path: Path, benchmark_fixture +) -> None: + import torch + + from cuphoton.xscan.pipeline_benchmark.fixture import prepare_fixture + + config, items = benchmark.read_inputs(*benchmark_fixture) + other = prepare_fixture( + tmp_path / "repeat", + images=1, + image_size=48, + candidates=2, + stamp_size=9, + seed=2026, + device="cuda:0", + ) + other_config, other_items = benchmark.read_inputs(*other) + assert items[0].candidates == other_items[0].candidates + for name in ("reference", "target", "variance", "fit_mask"): + left, right = getattr(items[0], name), getattr(other_items[0], name) + assert left.sha256 == right.sha256 + np.testing.assert_array_equal(np.load(left.path), np.load(right.path)) + models = [ + torch.load( + Path(value.checkpoint_dir) / "checkpoint.pt", weights_only=False + ) + for value in (config, other_config) + ] + assert models[0]["benchmark_fixture"]["trained"] is False + for key, value in models[0]["model_state"].items(): + assert value.device.type == "cpu" + torch.testing.assert_close( + value, models[1]["model_state"][key], rtol=0, atol=0 + ) + + +@pytest.mark.parametrize("order", ["pipeline-first", "staged-first"]) +def test_compare_cpu_orchestration_and_audit( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + benchmark_fixture, + order: str, +) -> None: + from cuphoton.xscan import device_pipeline + from cuphoton.xscan.pipeline_benchmark import stages + + config, items = benchmark.read_inputs(*benchmark_fixture) + science = {name: np.array([0.25, 0.75]) for name in stages.SCIENCE_KEYS} + predictions = [ + { + **candidate.to_payload(), + "logit": 0.25, + "probability": 0.75, + "decision": True, + } + for candidate in items[0].candidates + ] + metadata = { + "xpois": {"xpois_chi2": 1.0, "basis_terms": []}, + "xfit": {"parameter_names": ["amplitude"], "feature_names": ["flux"]}, + "xscan": {"predictions": predictions}, + } + # Exercise the real file contracts, stage summaries, round bookkeeping and + # audit. Only CUDA computation and process launching are replaced on CPU. + monkeypatch.setattr(stages, "_setup", lambda *_: {}) + monkeypatch.setattr( + stages, + "_run_item", + lambda stage, *_: ( + { + key: value + for key, value in science.items() + if key.startswith(stage + ".") + }, + metadata[stage], + {}, + ), + ) + monkeypatch.setattr( + device_pipeline, + "decode_device_pipeline_evidence", + lambda *_, **__: science, + ) + launched = [] + + def child(command, log, *, timeout): + stage = command[command.index("--stage") + 1] + output = Path(command[command.index("--output") + 1]) + launched.append(stage) + if stage != "pipeline": + stages.run_stage(stage, config, items, output) + else: + for index, phase in enumerate(("warmup", "measured")): + root = output / f"round-{index:03d}" + root.mkdir() + benchmark.write_json( + root / "item-0000.json", + { + "item_id": items[0].item_id, + "configuration_sha256": config.configuration_sha256, + "predictions": predictions, + "xpois": {"chi2": 1.0}, + "scientific_evidence": { + "xpois": {"basis_terms": []}, + "parameter_names": ["amplitude"], + "feature_names": ["flux"], + }, + }, + ) + benchmark.write_json( + output / "summary.json", + { + "device_provenance": { + "gpu_uuid": "GPU-test", + "gpu_name": "CPU test double", + }, + "rounds": [ + {"phase": "warmup", "batch_seconds": 2.0}, + {"phase": "measured", "batch_seconds": 1.0}, + ], + }, + ) + return 0.1 + + monkeypatch.setattr(benchmark, "run_child", child) + monkeypatch.setattr( + benchmark.subprocess, + "run", + lambda *_, **__: SimpleNamespace(stdout="610.57.04\n"), + ) + args = argparse.Namespace( + stage="compare", + output=tmp_path / "compare", + config=benchmark_fixture[0], + items=benchmark_fixture[1], + warmup=1, + repeat=1, + timeout=10, + order=order, + ) + benchmark.run_benchmark(args) + report = json.loads((args.output / "report.json").read_text()) + assert report["parity_passed"] + assert ( + len(json.loads((args.output / "parity.json").read_text())["checks"]) + == 2 + ) + assert launched.count("pipeline") == 1 + assert ( + launched.count("xpois") + == launched.count("xfit") + == launched.count("xscan") + == 2 + ) + assert (launched[0] == "pipeline") == (order == "pipeline-first") + assert report["provenance"]["device"]["driver_version"] == "610.57.04" + assert ( + report["fresh_stages_over_warm_pipeline_ratio"] + < report["raw_fresh_stages_over_warm_pipeline_ratio"] + ) + for row in report["staged"]["rounds"]: + assert row["extra_hashing_seconds"] > 0 + assert row["batch_seconds_without_extra_hashing"] == pytest.approx( + row["batch_seconds"] - row["extra_hashing_seconds"] + ) + # Candidate identity and exact scientific values both gate acceptance. + path = args.output / "pipeline/round-001/item-0000.json" + result = json.loads(path.read_text()) + result["predictions"][0]["candidate_id"] = "changed" + benchmark.write_json(path, result) + with pytest.raises(ValueError, match="candidate identity"): + benchmark.audit(args) + result["predictions"] = predictions + benchmark.write_json(path, result) + monkeypatch.setattr( + device_pipeline, + "decode_device_pipeline_evidence", + lambda *_, **__: { + name: values + 1 for name, values in science.items() + }, + ) + assert not benchmark.audit(args)["passed"] + with pytest.raises(FileExistsError): + stages.run_stage( + "xpois", config, items, args.output / "staged/round-001" + ) + + +def test_provenance_discovers_the_installed_cupy_distribution( + monkeypatch, +) -> None: + monkeypatch.setattr( + benchmark.importlib.metadata, + "packages_distributions", + lambda: {"cupy": ["cupy"]}, + ) + monkeypatch.setattr( + benchmark.importlib.metadata, "version", lambda name: "test-" + name + ) + assert benchmark.provenance()["packages"]["cupy"] == "test-cupy" diff --git a/tests/xscan/test_pipeline_benchmark_stages.py b/tests/xscan/test_pipeline_benchmark_stages.py index e6144f9..bff6910 100644 --- a/tests/xscan/test_pipeline_benchmark_stages.py +++ b/tests/xscan/test_pipeline_benchmark_stages.py @@ -258,3 +258,19 @@ def guarded_import(name, *args, **kwargs): assert runtime["torch"] is None assert runtime["cp"] is fake_cupy assert calls == [("cupy", 0), ("sync", 0)] + + +def test_hash_accounting_counts_both_reads_and_the_write( + tmp_path: Path, monkeypatch +) -> None: + ticks = iter(float(i) for i in range(6)) + monkeypatch.setattr(stages.time, "perf_counter", lambda: next(ticks)) + timings = {} + artifact = stages._write_arrays( + tmp_path, tmp_path / "arrays", {"values": np.arange(3)}, timings + )["values"] + assert timings["artifact_hash_seconds"] == 1.0 + np.testing.assert_array_equal( + stages._read_array(tmp_path, artifact, timings), np.arange(3) + ) + assert timings["artifact_hash_seconds"] == 3.0