diff --git a/docs/components/pipeline-stage-benchmark.md b/docs/components/pipeline-stage-benchmark.md new file mode 100644 index 0000000..bc03f97 --- /dev/null +++ b/docs/components/pipeline-stage-benchmark.md @@ -0,0 +1,148 @@ +# 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 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 \ + --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 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 \ + --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: 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 +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 +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 `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, +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 **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. + +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. 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. 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. + +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/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/__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/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/runner.py b/src/cuphoton/xscan/pipeline_benchmark/runner.py new file mode 100644 index 0000000..242c5d8 --- /dev/null +++ b/src/cuphoton/xscan/pipeline_benchmark/runner.py @@ -0,0 +1,558 @@ +# 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 importlib.metadata +import json +import os +import platform +import signal +import statistics +import subprocess +import sys +import time +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import numpy as np + +CLI = [sys.executable, "-m", "cuphoton", "xscan", "benchmark-pipeline"] + + +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: SimpleNamespace) -> 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) + 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", + { + "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), + "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__, + }, + }, + ) + + +def measure_pipeline(args: SimpleNamespace) -> dict[str, Any]: + output = args.output / "pipeline" + output.mkdir() + elapsed = run_child( + [ + *CLI, + "--stage", + "pipeline", + "--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: SimpleNamespace) -> 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( + [ + *CLI, + "--stage", + stage, + "--config", + str(args.config), + "--items", + str(args.items), + "--output", + 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" + 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) + result = {"rounds": rounds} + write_json(output / "summary.json", result) + return result + + +def audit(args: SimpleNamespace) -> 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 = {} + 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: + 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 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: + raise ValueError("repeat, warmup and timeout must be positive") + if (args.config is None) != (args.items is None): + raise ValueError("provide both --config and --items, or neither") + args.output = args.output.resolve() + 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: + 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["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( + {"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 diff --git a/src/cuphoton/xscan/pipeline_benchmark/stages.py b/src/cuphoton/xscan/pipeline_benchmark/stages.py new file mode 100644 index 0000000..eb4f361 --- /dev/null +++ b/src/cuphoton/xscan/pipeline_benchmark/stages.py @@ -0,0 +1,561 @@ +# 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 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 _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 _artifact_hash(path, timings) != 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 _artifact_hash(path, timings) != 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], + timings: dict[str, float] | None = None, +) -> 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": _artifact_hash(path, timings), + "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], timings + ) + + 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_coefficients", + "background_coefficients", + ) + } + 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 + }, + "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, + timings, + ), + ) + _measure( + timings, + "input_recheck_seconds", + lambda: pipeline._verify_item_hashes( + item, + 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( + { + "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, + "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. " + "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) + 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 + } 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 new file mode 100644 index 0000000..431bd2d --- /dev/null +++ b/tests/xscan/test_pipeline_benchmark.py @@ -0,0 +1,409 @@ +# 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 runner as benchmark +from cuphoton.xscan.pipeline_benchmark.runner 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 + 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) + 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() + + +@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_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() diff --git a/tests/xscan/test_pipeline_benchmark_stages.py b/tests/xscan/test_pipeline_benchmark_stages.py new file mode 100644 index 0000000..bff6910 --- /dev/null +++ b/tests/xscan/test_pipeline_benchmark_stages.py @@ -0,0 +1,276 @@ +# 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)] + + +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