From abbb5501096c5b000de6e2eafdb6adbec205e998 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Mon, 21 Sep 2026 16:07:56 -0700 Subject: [PATCH 1/2] Run the device pipeline in persistent Dragon workers Reuse one placed device context across complete image-pair jobs. Bind compact launch and result descriptors to input, configuration, and scientific evidence identities, and retain Torch-before-CuPy startup. Document direct execution and Dragon worker reuse with the required checkpoint and feature-schema contracts. Signed-off-by: Trent Nelson --- docs/components/xscan.md | 127 +++ src/cuphoton/xscan/dragon_pipeline.py | 1124 ++++++++++++++++++++++ tests/xscan/test_dragon_pipeline.py | 1243 +++++++++++++++++++++++++ 3 files changed, 2494 insertions(+) create mode 100644 src/cuphoton/xscan/dragon_pipeline.py create mode 100644 tests/xscan/test_dragon_pipeline.py diff --git a/docs/components/xscan.md b/docs/components/xscan.md index a11fe7f1..a5412e86 100644 --- a/docs/components/xscan.md +++ b/docs/components/xscan.md @@ -240,6 +240,133 @@ The `*.blackwell.example.yaml` files demonstrate throughput-oriented settings for recent NVIDIA GPUs. They are starting points, not universal performance recommendations. +## Persistent XPOIS, xFit and XScan pipeline + +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` +loads the model once and accepts serial jobs on one CUDA device. This path +requires CUDA 13, CuPy and Torch (`uv sync --locked --extra gpu`). + +Use a triplet fusion checkpoint and the exact `schema.json` from its training +xFit feature bundle. The schema must describe unmasked, unweighted Gaussian +difference fits with the configured stamp shape. Pipeline inference disables +AMP, TF32, compilation and cuDNN benchmarking through an explicit checkpoint +policy. This policy does not enable PyTorch's deterministic-algorithm mode. + +Prepare descriptors from caller-owned NPY images; this example uses an +existing 63-pixel checkpoint and an interior candidate in images of at least +95 by 95 pixels. Change the candidate coordinates and kernel settings for +your data. Optional item `variance` and `fit_mask` descriptors apply to XPOIS; +xFit consumes unweighted difference stamps. + +```python +from pathlib import Path + +import numpy as np + +from cuphoton.core.artifacts import file_sha256 +from cuphoton.xscan.device_pipeline import ( + DevicePipelineCandidate, + DevicePipelineConfig, + DevicePipelineItem, + DeviceWorkerContext, + DeviceXPOISPipelineConfig, + NpyArrayDescriptor, + decode_device_pipeline_evidence, + run_device_pipeline_item, +) + + +def describe(path): + path = Path(path).resolve() + array = np.load(path, allow_pickle=False, mmap_mode="r") + return NpyArrayDescriptor( + path=str(path), sha256=file_sha256(path), + shape=tuple(array.shape), dtype=array.dtype.str, + ) + + +checkpoint = Path("/path/to/training-run").resolve() +schema = Path("/path/to/training-features/schema.json").resolve() +config = DevicePipelineConfig( + device="cuda:0", + checkpoint_dir=str(checkpoint), + checkpoint_sha256=file_sha256(checkpoint / "checkpoint.pt"), + feature_schema_path=str(schema), + feature_schema_sha256=file_sha256(schema), + stamp_shape=(63, 63), + decision_threshold=0.5, + xpois=DeviceXPOISPipelineConfig( + kernel_shape=(15, 15), basis_sigmas=(1.5,), basis_degrees=(0,), + ), +) +items = tuple( + DevicePipelineItem( + item_id=f"pair-{index}", + reference=describe(reference), target=describe(target), + candidates=(DevicePipelineCandidate( + candidate_id=f"candidate-{index}", + center_x=47, center_y=47, source_index=0, + ),), + ) + for index, (reference, target) in enumerate([ + ("/path/to/reference-0.npy", "/path/to/target-0.npy"), + ("/path/to/reference-1.npy", "/path/to/target-1.npy"), + ]) +) +``` + +For direct execution, import Torch before CuPy calls CUDA, bind both libraries +to the configured device, and retain the context between items: + +```python +import torch +import cupy as cp + +torch.cuda.set_device(config.device_id) +cp.cuda.Device(config.device_id).use() +context = DeviceWorkerContext.initialize(config) +results = [run_device_pipeline_item(item, context) for item in items] +arrays = decode_device_pipeline_evidence( + results[0].scientific_evidence, config=config, +) +``` + +For Dragon, use the descriptor preparation in a fresh coordinator process +without importing Torch or CuPy there. Run the following entry point under +your site's installed Dragon launcher, with all input, checkpoint, schema +and output paths accessible on every worker: + +```python +from cuphoton.xscan.dragon_pipeline import run_dragon_device_pipeline + +if __name__ == "__main__": + batch = run_dragon_device_pipeline( + items=items, config=config, + output_root=Path("/path/to/pipeline-runs"), max_workers=1, + ) + if batch.status != "success": + raise RuntimeError(f"Pipeline failed: {batch.run_dir}") + print(batch.run_dir) +``` + +Dragon places one worker on each selected GPU and maps it to local `cuda:0`. +Each worker reuses its context for complete image-pair jobs; increase +`max_workers` to use more GPUs. Candidate order is preserved within each +item. Workers validate placement before loading Torch and probing CuPy. +Results include predictions, configuration/input hashes and compact +scientific evidence; Dragon saves each result in +`items//summary.json` and checks it again in the coordinator. +The evidence decoder returns the 22 named arrays for comparison. + +The pipeline retains device owners through the blocking terminal copy and +synchronizes failed work before reuse. Failed cleanup makes the context +unusable. Transfer receipts count pipeline-owned uploads and the packed +terminal download; internal XPOIS/xFit control transfers are outside that +count. Timings are host elapsed times, so distinguish model initialization, +first-item compilation and warmed context reuse when comparing runs. + ## Review workflow ```bash diff --git a/src/cuphoton/xscan/dragon_pipeline.py b/src/cuphoton/xscan/dragon_pipeline.py new file mode 100644 index 00000000..0244c345 --- /dev/null +++ b/src/cuphoton/xscan/dragon_pipeline.py @@ -0,0 +1,1124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Dragon orchestration for the persistent XPOIS-to-XScan device seam.""" + +from __future__ import annotations + +import hashlib +import importlib +import json +import math +import os +import time +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any, cast + +import numpy as np + +from cuphoton.core.artifacts import file_sha256 +from cuphoton.core.bulk import ( + Placement, + WorkItem, + atomic_write_json, + error_payload, + item_ids_sha256, + json_mapping, + read_json_mapping, + timestamp_utc, +) +from cuphoton.xpois import dragon as _dragon + +from .device_pipeline import ( + DEVICE_PIPELINE_RESULT_SCHEMA, + DevicePipelineConfig, + DevicePipelineItem, + DeviceWorkerContext, + NpyArrayDescriptor, + decode_device_pipeline_evidence, + run_device_pipeline_item, +) + +DRAGON_DEVICE_PIPELINE_MANIFEST_SCHEMA = ( + "cuphoton.xscan.device-pipeline.dragon-manifest/v1" +) +DRAGON_DEVICE_PIPELINE_INPUT_IDENTITY_SCHEMA = ( + "cuphoton.xscan.device-pipeline.dragon-input-identity/v1" +) +DRAGON_DEVICE_PIPELINE_OPTIONS_SCHEMA = ( + "cuphoton.xscan.device-pipeline.dragon-options/v1" +) +DRAGON_DEVICE_PIPELINE_RUN_SCHEMA = ( + "cuphoton.xscan.device-pipeline.dragon-run/v1" +) +DRAGON_DEVICE_PIPELINE_SUMMARY_SCHEMA = ( + "cuphoton.xscan.device-pipeline.dragon-summary/v1" +) +DRAGON_DEVICE_PIPELINE_ITEM_SCHEMA = ( + "cuphoton.xscan.device-pipeline.dragon-item/v1" +) +DRAGON_DEVICE_PIPELINE_SHARD_SCHEMA = ( + "cuphoton.xscan.device-pipeline.dragon-shard/v1" +) + +_BACKEND = "cupy" +_RUNTIME_SCHEMA = "cuphoton.xscan.device-pipeline.worker-runtime/v1" +_RESULT_METADATA_FIELDS = frozenset( + { + "schema", + "summary_sha256", + "result_sha256", + "configuration_sha256", + "checkpoint_sha256", + "feature_schema_sha256", + "input_sha256", + "candidate_count", + "positive_count", + "scientific_evidence_sha256", + "scientific_evidence_bytes", + "transfers", + } +) +_RESULT_IDENTITY_FIELDS = ( + "schema", + "item_id", + "device", + "checkpoint_sha256", + "feature_schema_sha256", + "configuration_sha256", + "input_sha256", + "xpois", + "predictions", + "scientific_evidence", + "transfers", +) +_RESULT_PAYLOAD_FIELDS = frozenset( + (*_RESULT_IDENTITY_FIELDS, "timings", "result_sha256") +) +_PREDICTION_FIELDS = frozenset( + { + "candidate_id", + "center_x", + "center_y", + "source_index", + "logit", + "probability", + "decision", + } +) +_XPOIS_RESULT_FIELDS = frozenset({"chi2", "dof", "fit_pixel_count"}) +_TIMING_FIELDS = frozenset({"stages_seconds", "total_seconds", "semantics"}) +_TIMING_STAGES = frozenset( + { + "input_read", + "input_h2d", + "xpois", + "stamp_extraction", + "xfit", + "feature_transform", + "evidence_pack", + "dlpack_handoff", + "xscan", + "terminal_pack", + "terminal_d2h", + } +) + + +def _collect_device_pipeline_gpu_identity(backend: str) -> Mapping[str, Any]: + """Load Torch before CuPy touches CUDA in the placed worker. + + Torch establishes its CUDA library load order during import. A prior CuPy + runtime call can make the later cuBLASLt model initialization fail. + """ + + importlib.import_module("torch") + return _dragon._collect_gpu_identity(backend) + + +_TRANSFER_FIELDS = frozenset( + { + "input_h2d_bytes", + "terminal_d2h_bytes", + "dlpack_shared_bytes", + "pipeline_full_array_d2h_bytes", + "pipeline_compact_h2d_bytes", + "terminal_d2h_calls", + "accounting_scope", + "device_stage_internal_transfers", + } +) +_XFIT_SCALAR_EVIDENCE_SEGMENTS = ( + "xfit.status_codes", + "xfit.converged", + "xfit.evaluations", + "xfit.residual_norm", + "xfit.chi_square", + "xfit.valid_pixel_count", + "xfit.valid_pixel_fraction", + "xfit.null_chi_square", + "xfit.delta_chi_square", + "xfit.fractional_null_improvement", + "xfit.degrees_of_freedom", + "xfit.reduced_chi_square", + "xfit.uncertainty_valid", + "xfit.uncertainty_reason_codes", +) +_EVIDENCE_SEGMENTS = ( + "xpois.kernel_coefficients", + "xpois.background_coefficients", + *_XFIT_SCALAR_EVIDENCE_SEGMENTS, + "xfit.parameters", + "xfit.standard_errors", + "xfit.covariance", + "xfit.features", + "xscan.logits", + "xscan.probabilities", +) + + +def _json_sha256(value: Any) -> str: + encoded = json.dumps( + value, + allow_nan=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def _is_sha256(value: Any) -> bool: + return ( + isinstance(value, str) + and len(value) == 64 + and all(character in "0123456789abcdef" for character in value) + ) + + +def _exact_mapping( + value: Any, *, expected: frozenset[str], field: str +) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError(f"{field} must be a mapping") + result = dict(value) + if set(result) != expected: + raise ValueError(f"{field} has invalid fields") + return result + + +def _finite_float( + value: Any, *, field: str, minimum: float | None = None +) -> float: + if isinstance(value, bool) or not isinstance(value, float): + raise TypeError(f"{field} must be a float") + if not math.isfinite(value): + raise ValueError(f"{field} must be finite") + if minimum is not None and value < minimum: + raise ValueError(f"{field} must be at least {minimum}") + return value + + +def _strict_integer( + value: Any, *, field: str, minimum: int | None = None +) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{field} must be an integer") + if minimum is not None and value < minimum: + raise ValueError(f"{field} must be at least {minimum}") + return value + + +def _strict_name_list(value: Any, *, field: str) -> list[str]: + if ( + not isinstance(value, list) + or not value + or any(not isinstance(name, str) or not name for name in value) + or len(set(value)) != len(value) + ): + raise ValueError(f"{field} must contain unique non-empty strings") + return value + + +def _expected_input_h2d_bytes(item: DevicePipelineItem) -> int: + return sum( + math.prod(descriptor.shape) + * ( + np.dtype(np.bool_).itemsize + if name == "fit_mask" + else np.dtype(np.float64).itemsize + ) + for name, descriptor in _item_descriptors(item) + ) + + +def _expected_dlpack_shared_bytes( + decoded: Mapping[str, np.ndarray], + *, + config: DevicePipelineConfig, + candidate_count: int, +) -> int: + stamp_bytes = ( + candidate_count + * 3 + * math.prod(config.stamp_shape) + * np.dtype(np.float32).itemsize + ) + feature_bytes = ( + decoded["xfit.features"].size * np.dtype(np.float32).itemsize + ) + evidence_bytes = sum( + values.size * np.dtype(np.float64).itemsize + for name, values in decoded.items() + if name not in {"xscan.logits", "xscan.probabilities"} + ) + return int(stamp_bytes + feature_bytes + evidence_bytes) + + +def _strict_device_pipeline_result_payload( + payload: Mapping[str, Any], + *, + item: DevicePipelineItem, + config: DevicePipelineConfig, +) -> dict[str, Any]: + """Validate one complete result/v2 and its canonical evidence receipt.""" + + values = json_mapping( + payload, field=f"device pipeline result {item.item_id!r}" + ) + if set(values) != _RESULT_PAYLOAD_FIELDS: + raise ValueError("device pipeline result has invalid fields") + if ( + values["schema"] != DEVICE_PIPELINE_RESULT_SCHEMA + or values["item_id"] != item.item_id + or values["device"] != config.device + or values["configuration_sha256"] != config.configuration_sha256 + or values["checkpoint_sha256"] != config.checkpoint_sha256 + or values["feature_schema_sha256"] != config.feature_schema_sha256 + or values["input_sha256"] != _input_sha256(item) + ): + raise ValueError("device pipeline result identity differs") + + xpois = _exact_mapping( + values["xpois"], + expected=_XPOIS_RESULT_FIELDS, + field="device pipeline XPOIS result", + ) + _finite_float( + xpois["chi2"], field="device pipeline XPOIS chi2", minimum=0.0 + ) + dof = _strict_integer(xpois["dof"], field="device pipeline XPOIS dof") + fit_pixel_count = _strict_integer( + xpois["fit_pixel_count"], + field="device pipeline XPOIS fit_pixel_count", + minimum=1, + ) + + raw_predictions = values["predictions"] + if not isinstance(raw_predictions, list) or len(raw_predictions) != len( + item.candidates + ): + raise ValueError("device pipeline predictions and candidates differ") + predictions: list[dict[str, Any]] = [] + for index, (raw_prediction, candidate) in enumerate( + zip(raw_predictions, item.candidates, strict=True) + ): + prediction = _exact_mapping( + raw_prediction, + expected=_PREDICTION_FIELDS, + field=f"device pipeline prediction {index}", + ) + if ( + type(prediction["candidate_id"]) + is not type(candidate.candidate_id) + or prediction["candidate_id"] != candidate.candidate_id + ): + raise ValueError( + "device pipeline prediction candidate ID differs" + ) + for field, expected in ( + ("center_x", candidate.center_x), + ("center_y", candidate.center_y), + ("source_index", candidate.source_index), + ): + if ( + _strict_integer( + prediction[field], + field=f"device pipeline prediction {index} {field}", + ) + != expected + ): + raise ValueError( + f"device pipeline prediction {index} {field} differs" + ) + _finite_float( + prediction["logit"], + field=f"device pipeline prediction {index} logit", + ) + probability = _finite_float( + prediction["probability"], + field=f"device pipeline prediction {index} probability", + minimum=0.0, + ) + if probability > 1.0: + raise ValueError( + "device pipeline prediction probability exceeds one" + ) + if not isinstance(prediction["decision"], bool): + raise TypeError( + "device pipeline prediction decision must be boolean" + ) + if prediction["decision"] is not ( + probability >= config.decision_threshold + ): + raise ValueError("device pipeline prediction decision differs") + predictions.append(prediction) + + transfers = _exact_mapping( + values["transfers"], + expected=_TRANSFER_FIELDS, + field="device pipeline transfer receipt", + ) + for name in ( + "input_h2d_bytes", + "terminal_d2h_bytes", + "dlpack_shared_bytes", + "pipeline_full_array_d2h_bytes", + "pipeline_compact_h2d_bytes", + "terminal_d2h_calls", + ): + _strict_integer( + transfers[name], + field=f"device pipeline transfer receipt {name}", + minimum=0, + ) + if ( + transfers["pipeline_full_array_d2h_bytes"] != 0 + or transfers["pipeline_compact_h2d_bytes"] != 0 + or transfers["terminal_d2h_calls"] != 1 + or transfers["accounting_scope"] != "pipeline-owned" + or transfers["device_stage_internal_transfers"] != "not-instrumented" + ): + raise ValueError("device pipeline transfer receipt differs") + + timings = _exact_mapping( + values["timings"], + expected=_TIMING_FIELDS, + field="device pipeline timing receipt", + ) + stages = _exact_mapping( + timings["stages_seconds"], + expected=_TIMING_STAGES, + field="device pipeline stage timings", + ) + for name, value in stages.items(): + _finite_float( + value, + field=f"device pipeline stage timing {name}", + minimum=0.0, + ) + _finite_float( + timings["total_seconds"], + field="device pipeline total timing", + minimum=0.0, + ) + if timings["semantics"] != "host-elapsed; terminal-d2h-is-blocking": + raise ValueError("device pipeline timing semantics differ") + + evidence = values["scientific_evidence"] + if not isinstance(evidence, Mapping): + raise TypeError( + "device pipeline scientific evidence must be a mapping" + ) + decoded = decode_device_pipeline_evidence(evidence, config=config) + if tuple(decoded) != _EVIDENCE_SEGMENTS: + raise ValueError( + "device pipeline scientific evidence segments differ" + ) + candidate_count = _strict_integer( + evidence.get("candidate_count"), + field="device pipeline scientific evidence candidate_count", + minimum=1, + ) + if candidate_count != len(predictions): + raise ValueError( + "device pipeline scientific evidence candidate count differs" + ) + parameter_names = _strict_name_list( + evidence.get("parameter_names"), + field="device pipeline scientific evidence parameter_names", + ) + feature_names = _strict_name_list( + evidence.get("feature_names"), + field="device pipeline scientific evidence feature_names", + ) + expected_shapes = { + **{ + name: (candidate_count,) + for name in _XFIT_SCALAR_EVIDENCE_SEGMENTS + }, + "xfit.parameters": (candidate_count, len(parameter_names)), + "xfit.standard_errors": (candidate_count, len(parameter_names)), + "xfit.covariance": ( + candidate_count, + len(parameter_names), + len(parameter_names), + ), + "xfit.features": (candidate_count, len(feature_names)), + "xscan.logits": (candidate_count,), + "xscan.probabilities": (candidate_count,), + } + if any( + tuple(decoded[name].shape) != shape + for name, shape in expected_shapes.items() + ): + raise ValueError("device pipeline scientific evidence shape differs") + kernel_coefficients = decoded["xpois.kernel_coefficients"] + background_coefficients = decoded["xpois.background_coefficients"] + if ( + kernel_coefficients.ndim != 1 + or kernel_coefficients.size < 1 + or background_coefficients.ndim != 1 + or background_coefficients.size < 1 + ): + raise ValueError("device pipeline XPOIS evidence shape differs") + if dof != fit_pixel_count - int( + kernel_coefficients.size + background_coefficients.size + ): + raise ValueError("device pipeline XPOIS degrees of freedom differ") + for index, prediction in enumerate(predictions): + if ( + float(decoded["xscan.logits"][index]) != prediction["logit"] + or float(decoded["xscan.probabilities"][index]) + != prediction["probability"] + ): + raise ValueError( + "device pipeline predictions and scientific evidence differ" + ) + if transfers["terminal_d2h_bytes"] != evidence.get("packed_byte_count"): + raise ValueError( + "device pipeline evidence and terminal transfer bytes differ" + ) + if transfers["input_h2d_bytes"] != _expected_input_h2d_bytes(item): + raise ValueError("device pipeline input H2D byte count differs") + if transfers["dlpack_shared_bytes"] != _expected_dlpack_shared_bytes( + decoded, + config=config, + candidate_count=candidate_count, + ): + raise ValueError("device pipeline DLPack shared byte count differs") + + result_sha256 = values["result_sha256"] + if not _is_sha256(result_sha256): + raise ValueError("device pipeline result SHA-256 is invalid") + identity = {name: values[name] for name in _RESULT_IDENTITY_FIELDS} + if _json_sha256(identity) != result_sha256: + raise ValueError("device pipeline result SHA-256 differs") + return values + + +def _item_descriptors( + item: DevicePipelineItem, +) -> tuple[tuple[str, NpyArrayDescriptor], ...]: + values: list[tuple[str, NpyArrayDescriptor]] = [ + ("reference", item.reference), + ("target", item.target), + ] + if item.variance is not None: + values.append(("variance", item.variance)) + if item.fit_mask is not None: + values.append(("fit_mask", item.fit_mask)) + return tuple(values) + + +def _input_sha256(item: DevicePipelineItem) -> dict[str, str]: + return { + name: descriptor.sha256 + for name, descriptor in _item_descriptors(item) + } + + +def _observe_content_file( + path: Path, + *, + expected_sha256: str, + description: str, + cache: dict[Path, tuple[str, int]], +) -> int: + try: + if not path.is_file(): + raise FileNotFoundError(path) + size_bytes = path.stat().st_size + except OSError as exc: + raise FileNotFoundError( + f"{description} is not readable: {path}" + ) from exc + observed = cache.get(path) + if observed is None: + observed = (file_sha256(path), size_bytes) + cache[path] = observed + if observed != (expected_sha256, size_bytes): + raise RuntimeError(f"{description} changed before Dragon launch") + return size_bytes + + +def _preflight_items( + items: Sequence[DevicePipelineItem], +) -> tuple[tuple[WorkItem, ...], list[dict[str, Any]]]: + if not items: + raise ValueError("Dragon device pipeline items must not be empty") + if any(not isinstance(item, DevicePipelineItem) for item in items): + raise TypeError("items must contain only DevicePipelineItem values") + item_ids = [item.item_id for item in items] + if len(set(item_ids)) != len(item_ids): + raise ValueError("Dragon device pipeline item IDs must be unique") + + observed_files: dict[Path, tuple[str, int]] = {} + work_items: list[WorkItem] = [] + identities: list[dict[str, Any]] = [] + for item in items: + by_path: dict[str, dict[str, Any]] = {} + for role, descriptor in _item_descriptors(item): + path = Path(descriptor.path) + size_bytes = _observe_content_file( + path, + expected_sha256=descriptor.sha256, + description=f"{item.item_id} {role} input", + cache=observed_files, + ) + descriptor_payload = descriptor.to_payload() + existing = by_path.get(descriptor.path) + if existing is None: + by_path[descriptor.path] = { + **descriptor_payload, + "size_bytes": size_bytes, + "roles": [role], + } + else: + observed_descriptor = { + name: existing[name] + for name in ("path", "sha256", "shape", "dtype") + } + if observed_descriptor != descriptor_payload: + raise ValueError( + f"item {item.item_id!r} describes one input path " + "with conflicting identities" + ) + existing["roles"].append(role) + files = [by_path[path] for path in sorted(by_path)] + weight_bytes = sum(value["size_bytes"] for value in files) + work_items.append( + WorkItem( + item_id=item.item_id, + payload=item.to_payload(), + weight_bytes=weight_bytes, + ) + ) + identities.append({"item_id": item.item_id, "files": files}) + return tuple(work_items), identities + + +def _preflight( + items: Sequence[DevicePipelineItem], + *, + config: DevicePipelineConfig, +) -> tuple[ + tuple[WorkItem, ...], + dict[str, Any], + dict[str, Any], + str, +]: + if not isinstance(config, DevicePipelineConfig): + raise TypeError("config must be a DevicePipelineConfig") + if config.device != "cuda:0": + raise ValueError( + "Dragon device workers require config.device='cuda:0'; " + "Dragon maps each physical GPU affinity to local CUDA ordinal 0" + ) + + observed_files: dict[Path, tuple[str, int]] = {} + checkpoint_path = Path(config.checkpoint_dir) / "checkpoint.pt" + checkpoint_size = _observe_content_file( + checkpoint_path, + expected_sha256=config.checkpoint_sha256, + description="XScan checkpoint", + cache=observed_files, + ) + feature_schema_path = Path(config.feature_schema_path) + feature_schema_size = _observe_content_file( + feature_schema_path, + expected_sha256=config.feature_schema_sha256, + description="xFit feature schema", + cache=observed_files, + ) + work_items, item_identities = _preflight_items(items) + manifest_payload = { + "schema": DRAGON_DEVICE_PIPELINE_MANIFEST_SCHEMA, + "backend": _BACKEND, + "configuration": config.to_payload(), + "configuration_sha256": config.configuration_sha256, + "items": [item.to_dict() for item in work_items], + } + manifest_sha256 = _json_sha256(manifest_payload) + input_identity_payload = { + "schema": DRAGON_DEVICE_PIPELINE_INPUT_IDENTITY_SCHEMA, + "content_sha256": { + "checkpoint": config.checkpoint_sha256, + "feature_schema": config.feature_schema_sha256, + }, + "configuration_sha256": config.configuration_sha256, + "configuration_files": { + "checkpoint": { + "path": str(checkpoint_path), + "size_bytes": checkpoint_size, + }, + "feature_schema": { + "path": str(feature_schema_path), + "size_bytes": feature_schema_size, + }, + }, + "items": item_identities, + } + return ( + work_items, + manifest_payload, + input_identity_payload, + manifest_sha256, + ) + + +def _strict_worker_options( + payload: Mapping[str, Any], +) -> DevicePipelineConfig: + options = json_mapping(payload, field="Dragon device pipeline options") + expected = { + "schema", + "backend", + "configuration", + "configuration_sha256", + } + if set(options) != expected: + raise ValueError("Dragon device pipeline options have invalid fields") + if options["schema"] != DRAGON_DEVICE_PIPELINE_OPTIONS_SCHEMA: + raise ValueError("unsupported Dragon device pipeline options schema") + if options["backend"] != _BACKEND: + raise ValueError("Dragon device pipeline backend must be 'cupy'") + configuration = options["configuration"] + if not isinstance(configuration, Mapping): + raise TypeError( + "Dragon device pipeline configuration must be a mapping" + ) + config = DevicePipelineConfig.from_payload(configuration) + if config.device != "cuda:0": + raise ValueError( + "Dragon device pipeline worker device must be cuda:0" + ) + if options["configuration_sha256"] != config.configuration_sha256: + raise ValueError( + "Dragon device pipeline configuration SHA-256 differs" + ) + return config + + +def _strict_work_item(payload: Mapping[str, Any]) -> WorkItem: + values = json_mapping(payload, field="Dragon device pipeline work item") + if set(values) != {"item_id", "payload", "weight_bytes"}: + raise ValueError( + "Dragon device pipeline work item has invalid fields" + ) + if not isinstance(values["item_id"], str): + raise TypeError( + "Dragon device pipeline work item ID must be a string" + ) + return WorkItem.from_dict(values) + + +def _publish_item_result( + *, + work_item: WorkItem, + item_dir: Path, + config: DevicePipelineConfig, + context: DeviceWorkerContext, +) -> Mapping[str, Any]: + started = time.perf_counter() + item = DevicePipelineItem.from_payload(work_item.payload) + if item.item_id != work_item.item_id: + raise ValueError("Dragon work item and device item IDs differ") + result = run_device_pipeline_item(item, context) + if not callable(getattr(result, "to_payload", None)): + raise TypeError("device pipeline runner returned an invalid result") + payload = _strict_device_pipeline_result_payload( + result.to_payload(), item=item, config=config + ) + + item_dir.mkdir(parents=True, exist_ok=False) + summary_path = item_dir / "summary.json" + try: + atomic_write_json(summary_path, payload) + except BaseException as exc: + try: + summary_path.unlink(missing_ok=True) + item_dir.rmdir() + except OSError as cleanup_error: + exc.add_note( + "device pipeline item cleanup failed: " + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + raise + + predictions = cast(list[Mapping[str, Any]], payload["predictions"]) + timing = cast(Mapping[str, Any], payload["timings"]) + evidence = cast(Mapping[str, Any], payload["scientific_evidence"]) + return { + "run_dir": str(item_dir), + "summary_path": str(summary_path), + "requested_backend": _BACKEND, + "backend": _BACKEND, + "device": config.device, + "runtime": { + "schema": _RUNTIME_SCHEMA, + "context_load_seconds": context.load_seconds, + "configuration_sha256": config.configuration_sha256, + "checkpoint_sha256": config.checkpoint_sha256, + "feature_schema_sha256": config.feature_schema_sha256, + }, + "timings_sec": timing["stages_seconds"], + "wall_sec": {"item_runner": time.perf_counter() - started}, + "device_pipeline": { + "schema": DEVICE_PIPELINE_RESULT_SCHEMA, + "summary_sha256": file_sha256(summary_path), + "result_sha256": payload["result_sha256"], + "configuration_sha256": config.configuration_sha256, + "checkpoint_sha256": config.checkpoint_sha256, + "feature_schema_sha256": config.feature_schema_sha256, + "input_sha256": payload["input_sha256"], + "candidate_count": len(predictions), + "positive_count": sum( + prediction.get("decision") is True + for prediction in predictions + ), + "scientific_evidence_sha256": evidence["packed_sha256"], + "scientific_evidence_bytes": evidence["packed_byte_count"], + "transfers": payload["transfers"], + }, + } + + +def _dragon_device_pipeline_worker( + run_id: str, + run_dir_raw: str, + placement_payload: Mapping[str, Any], + item_payloads: Sequence[Mapping[str, Any]], + options_payload: Mapping[str, Any], + results_queue: Any, + allow_loopback_alias: bool, +) -> None: + """Dragon target that owns one persistent context for one whole shard.""" + + worker_started = time.perf_counter() + started_at = timestamp_utc() + run_dir: Path | None = None + placement: Placement | None = None + items: tuple[WorkItem, ...] = () + try: + run_dir = Path(run_dir_raw) + placement = Placement(**dict(placement_payload)) + items = tuple(_strict_work_item(payload) for payload in item_payloads) + config = _strict_worker_options(options_payload) + context_attempted = False + context: DeviceWorkerContext | None = None + context_error: Exception | None = None + + def item_runner( + work_item: WorkItem, + item_dir: Path, + _options: Mapping[str, Any], + ) -> Mapping[str, Any]: + nonlocal context_attempted, context, context_error + if not context_attempted: + context_attempted = True + try: + context = DeviceWorkerContext.initialize(config) + except Exception as exc: + context_error = exc + if context_error is not None: + raise context_error + if context is None: + raise RuntimeError( + "Dragon device worker context is unavailable" + ) + return _publish_item_result( + work_item=work_item, + item_dir=item_dir, + config=config, + context=context, + ) + + result = _dragon._execute_shard( + run_id=run_id, + run_dir=run_dir, + placement=placement, + items=items, + options=json_mapping( + options_payload, field="Dragon device pipeline options" + ), + item_runner=item_runner, + gpu_identity_loader=_collect_device_pipeline_gpu_identity, + backend=_BACKEND, + record_schema=DRAGON_DEVICE_PIPELINE_ITEM_SCHEMA, + shard_schema=DRAGON_DEVICE_PIPELINE_SHARD_SCHEMA, + allow_loopback_alias=allow_loopback_alias, + ) + except Exception as exc: + result = { + "schema": DRAGON_DEVICE_PIPELINE_SHARD_SCHEMA, + "worker_id": ( + placement.worker_id if placement is not None else None + ), + "status": "failed", + "item_count": len(items), + "success_count": 0, + "failed_count": len(items), + "weight_bytes": sum(item.weight_bytes for item in items), + "item_ids_sha256": item_ids_sha256(items), + "started_at_utc": started_at, + "completed_at_utc": timestamp_utc(), + "worker_wall_sec": time.perf_counter() - worker_started, + "timings_sec": {}, + "provenance": None, + "record_write_errors": [], + "error": error_payload(exc), + } + if run_dir is not None and placement is not None: + worker_path = ( + run_dir / "workers" / f"worker-{placement.worker_id:04d}.json" + ) + try: + if not os.path.lexists(worker_path): + atomic_write_json(worker_path, result) + except Exception as artifact_error: + result["artifact_error"] = error_payload(artifact_error) + results_queue.put(result) + + +def _finite_nonnegative(value: Any) -> bool: + return ( + not isinstance(value, bool) + and isinstance(value, (int, float)) + and math.isfinite(value) + and value >= 0 + ) + + +def _success_record_problems( + record: Mapping[str, Any], + *, + output_root: Path, + config: DevicePipelineConfig, + items: Mapping[str, DevicePipelineItem], +) -> tuple[str, ...]: + problems: set[str] = set() + item_id = record.get("item_id") + item = items.get(item_id) if isinstance(item_id, str) else None + record_device = record.get("device") + if record_device != config.device: + problems.add("device") + runtime = record.get("runtime") + if ( + not isinstance(runtime, Mapping) + or set(runtime) + != { + "schema", + "context_load_seconds", + "configuration_sha256", + "checkpoint_sha256", + "feature_schema_sha256", + } + or runtime.get("schema") != _RUNTIME_SCHEMA + or not _finite_nonnegative(runtime.get("context_load_seconds")) + or runtime.get("configuration_sha256") != config.configuration_sha256 + or runtime.get("checkpoint_sha256") != config.checkpoint_sha256 + or runtime.get("feature_schema_sha256") + != config.feature_schema_sha256 + ): + problems.add("runtime") + + metadata = record.get("device_pipeline") + if ( + not isinstance(metadata, Mapping) + or set(metadata) != _RESULT_METADATA_FIELDS + or metadata.get("schema") != DEVICE_PIPELINE_RESULT_SCHEMA + or not _is_sha256(metadata.get("summary_sha256")) + or not _is_sha256(metadata.get("result_sha256")) + or metadata.get("configuration_sha256") != config.configuration_sha256 + or metadata.get("checkpoint_sha256") != config.checkpoint_sha256 + or metadata.get("feature_schema_sha256") + != config.feature_schema_sha256 + or not _is_sha256(metadata.get("scientific_evidence_sha256")) + or isinstance(metadata.get("scientific_evidence_bytes"), bool) + or not isinstance(metadata.get("scientific_evidence_bytes"), int) + or metadata["scientific_evidence_bytes"] <= 0 + or item is None + or metadata.get("input_sha256") != _input_sha256(item) + ): + problems.add("device_pipeline") + return tuple(sorted(problems)) + + assert item is not None + candidate_count = metadata.get("candidate_count") + positive_count = metadata.get("positive_count") + if ( + isinstance(candidate_count, bool) + or not isinstance(candidate_count, int) + or candidate_count != len(item.candidates) + or isinstance(positive_count, bool) + or not isinstance(positive_count, int) + or not 0 <= positive_count <= candidate_count + ): + problems.add("device_pipeline") + transfers = metadata.get("transfers") + if ( + not isinstance(transfers, Mapping) + or transfers.get("pipeline_full_array_d2h_bytes") != 0 + or transfers.get("pipeline_compact_h2d_bytes") != 0 + or transfers.get("terminal_d2h_calls") != 1 + ): + problems.add("device_pipeline") + + run_id = record.get("run_id") + summary_relative = record.get("summary_path") + expected_summary = f"items/{item.item_id}/summary.json" + if not isinstance(run_id, str) or summary_relative != expected_summary: + problems.add("summary_path") + return tuple(sorted(problems)) + assert isinstance(summary_relative, str) + summary_path = output_root / run_id / summary_relative + try: + summary = read_json_mapping(summary_path) + if file_sha256(summary_path) != metadata["summary_sha256"]: + problems.add("device_pipeline") + except Exception: + problems.add("summary_path") + return tuple(sorted(problems)) + + try: + summary = _strict_device_pipeline_result_payload( + summary, item=item, config=config + ) + except (TypeError, ValueError): + problems.add("device_pipeline") + return tuple(sorted(problems)) + evidence = cast(Mapping[str, Any], summary["scientific_evidence"]) + timing = cast(Mapping[str, Any], summary["timings"]) + predictions = cast(list[Mapping[str, Any]], summary["predictions"]) + if record_device != summary["device"]: + problems.add("device") + if ( + summary["result_sha256"] != metadata["result_sha256"] + or evidence["packed_sha256"] != metadata["scientific_evidence_sha256"] + or evidence["packed_byte_count"] + != metadata["scientific_evidence_bytes"] + or summary.get("transfers") != transfers + or sum(prediction["decision"] is True for prediction in predictions) + != positive_count + or record.get("timings_sec") != timing["stages_seconds"] + ): + problems.add("device_pipeline") + return tuple(sorted(problems)) + + +def _failed_record_problems( + record: Mapping[str, Any], + *, + output_root: Path, + items: Mapping[str, DevicePipelineItem], +) -> tuple[str, ...]: + item_id = record.get("item_id") + run_id = record.get("run_id") + if ( + not isinstance(item_id, str) + or item_id not in items + or not isinstance(run_id, str) + ): + return () + item_dir = output_root / run_id / "items" / item_id + return ("failed_item_output",) if os.path.lexists(item_dir) else () + + +def run_dragon_device_pipeline( + *, + items: Sequence[DevicePipelineItem], + config: DevicePipelineConfig, + output_root: Path, + run_id: str | None = None, + max_workers: int | None = None, + result_timeout_sec: float = 120.0, + worker_timeout_sec: float = 3600.0, +) -> _dragon.DragonBatchResult: + """Run compact file-backed pipeline items on persistent Dragon workers.""" + + invocation_start = time.perf_counter() + started_at = timestamp_utc() + preflight_started = time.perf_counter() + normalized_items = tuple(items) + ( + work_items, + manifest_payload, + input_identity_payload, + manifest_sha256, + ) = _preflight(normalized_items, config=config) + preflight_seconds = time.perf_counter() - preflight_started + options_payload = { + "schema": DRAGON_DEVICE_PIPELINE_OPTIONS_SCHEMA, + "backend": _BACKEND, + "configuration": config.to_payload(), + "configuration_sha256": config.configuration_sha256, + } + by_id = {item.item_id: item for item in normalized_items} + resolved_output_root = output_root.expanduser().resolve() + return _dragon.run_dragon_work_items( + items=work_items, + output_root=resolved_output_root, + run_id=run_id, + max_workers=max_workers, + result_timeout_sec=result_timeout_sec, + worker_timeout_sec=worker_timeout_sec, + backend=_BACKEND, + options_payload=options_payload, + manifest_payload=manifest_payload, + input_identity_payload=input_identity_payload, + manifest_sha256=manifest_sha256, + worker_target=_dragon_device_pipeline_worker, + run_prefix="dragon-device-pipeline", + run_schema=DRAGON_DEVICE_PIPELINE_RUN_SCHEMA, + summary_schema=DRAGON_DEVICE_PIPELINE_SUMMARY_SCHEMA, + record_schema=DRAGON_DEVICE_PIPELINE_ITEM_SCHEMA, + shard_schema=DRAGON_DEVICE_PIPELINE_SHARD_SCHEMA, + coordinator_timings={"manifest_preflight_sec": preflight_seconds}, + invocation_start=invocation_start, + started_at=started_at, + success_record_validator=lambda record: _success_record_problems( + record, + output_root=resolved_output_root, + config=config, + items=by_id, + ), + failed_record_validator=lambda record: _failed_record_problems( + record, + output_root=resolved_output_root, + items=by_id, + ), + ) + + +__all__ = [ + "DRAGON_DEVICE_PIPELINE_INPUT_IDENTITY_SCHEMA", + "DRAGON_DEVICE_PIPELINE_ITEM_SCHEMA", + "DRAGON_DEVICE_PIPELINE_MANIFEST_SCHEMA", + "DRAGON_DEVICE_PIPELINE_OPTIONS_SCHEMA", + "DRAGON_DEVICE_PIPELINE_RUN_SCHEMA", + "DRAGON_DEVICE_PIPELINE_SHARD_SCHEMA", + "DRAGON_DEVICE_PIPELINE_SUMMARY_SCHEMA", + "run_dragon_device_pipeline", +] diff --git a/tests/xscan/test_dragon_pipeline.py b/tests/xscan/test_dragon_pipeline.py new file mode 100644 index 00000000..e03f9059 --- /dev/null +++ b/tests/xscan/test_dragon_pipeline.py @@ -0,0 +1,1243 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Dragon adapter tests for the persistent device pipeline.""" + +from __future__ import annotations + +import base64 +import hashlib +import json +import math +import os +import queue +import socket +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest + +from cuphoton.core.artifacts import file_sha256 +from cuphoton.xfit.api import DipoleFitUncertaintyReason +from cuphoton.xfit.solver import LMStatus +from cuphoton.xpois import ( + GaussianBasisComponent, + build_gaussian_polynomial_basis, + triangular_degree_pairs, +) +from cuphoton.xpois import dragon as dragon_module +from cuphoton.xscan import dragon_pipeline +from cuphoton.xscan.device_pipeline import ( + DEVICE_PIPELINE_EVIDENCE_SCHEMA, + DEVICE_PIPELINE_RESULT_SCHEMA, + DevicePipelineCandidate, + DevicePipelineConfig, + DevicePipelineItem, + DevicePipelineTransferReceipt, + DeviceXFitPipelineConfig, + DeviceXPOISPipelineConfig, + NpyArrayDescriptor, +) +from cuphoton.xscan.xfit_features import FEATURE_NAMES, _canonical_feature_row + +_XFIT_SCALAR_EVIDENCE_SEGMENTS = ( + "xfit.status_codes", + "xfit.converged", + "xfit.evaluations", + "xfit.residual_norm", + "xfit.chi_square", + "xfit.valid_pixel_count", + "xfit.valid_pixel_fraction", + "xfit.null_chi_square", + "xfit.delta_chi_square", + "xfit.fractional_null_improvement", + "xfit.degrees_of_freedom", + "xfit.reduced_chi_square", + "xfit.uncertainty_valid", + "xfit.uncertainty_reason_codes", +) +_TIMING_STAGES = ( + "input_read", + "input_h2d", + "xpois", + "stamp_extraction", + "xfit", + "feature_transform", + "evidence_pack", + "dlpack_handoff", + "xscan", + "terminal_pack", + "terminal_d2h", +) + + +def _descriptor(path: Path) -> NpyArrayDescriptor: + values = np.load(path, allow_pickle=False) + return NpyArrayDescriptor( + path=str(path.resolve()), + sha256=file_sha256(path), + shape=tuple(values.shape), + dtype=values.dtype.str, + ) + + +def _config( + tmp_path: Path, + *, + device: str = "cuda:0", + xpois: DeviceXPOISPipelineConfig | None = None, +) -> DevicePipelineConfig: + checkpoint_dir = tmp_path / "checkpoint" + checkpoint_dir.mkdir() + checkpoint_path = checkpoint_dir / "checkpoint.pt" + checkpoint_path.write_bytes(b"checkpoint") + feature_schema_path = tmp_path / "feature-schema.json" + feature_schema_path.write_text("{}\n", encoding="utf-8") + return DevicePipelineConfig( + device=device, + checkpoint_dir=str(checkpoint_dir.resolve()), + checkpoint_sha256=file_sha256(checkpoint_path), + feature_schema_path=str(feature_schema_path.resolve()), + feature_schema_sha256=file_sha256(feature_schema_path), + stamp_shape=(11, 11), + decision_threshold=0.5, + xpois=xpois + or DeviceXPOISPipelineConfig( + kernel_shape=(3, 3), + basis_sigmas=(0.8,), + basis_degrees=(0,), + ), + xfit=DeviceXFitPipelineConfig(max_evaluations=4), + ) + + +def _item(tmp_path: Path, item_id: str) -> DevicePipelineItem: + reference_path = tmp_path / f"{item_id}-reference.npy" + target_path = tmp_path / f"{item_id}-target.npy" + np.save(reference_path, np.zeros((31, 31), dtype=np.float32)) + np.save(target_path, np.ones((31, 31), dtype=np.float64)) + return DevicePipelineItem( + item_id=item_id, + reference=_descriptor(reference_path), + target=_descriptor(target_path), + candidates=( + DevicePipelineCandidate( + candidate_id=f"{item_id}-candidate", + center_x=15, + center_y=15, + source_index=0, + ), + ), + ) + + +def _payload_sha256(payload: dict[str, Any]) -> str: + encoded = json.dumps( + payload, + allow_nan=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def _fake_scientific_evidence( + *, + config: DevicePipelineConfig, + candidate_count: int, + logit: float, + probability: float, + xpois: DeviceXPOISPipelineConfig | None = None, + variance_present: bool = False, +) -> dict[str, Any]: + xpois = xpois or config.xpois + parameter_names = [ + "amplitude", + "sigma_x", + "sigma_y", + "theta", + "x_pos", + "y_pos", + "x_neg", + "y_neg", + ] + feature_names = list(FEATURE_NAMES) + parameters = np.tile( + np.asarray( + [[2.0, 1.5, 1.25, 0.0, 1.0, 0.0, -1.0, 0.0]], + dtype=" dict[str, Any]: + identity = { + name: value + for name, value in payload.items() + if name not in {"timings", "result_sha256"} + } + payload["result_sha256"] = _payload_sha256(identity) + return payload + + +def _fake_result_payload( + item: DevicePipelineItem, + config: DevicePipelineConfig, + *, + evidence_xpois: DeviceXPOISPipelineConfig | None = None, + evidence_variance_present: bool = False, +) -> dict[str, Any]: + logit = float(np.float32(0.25)) + probability = float(np.float32(1.0 / (1.0 + np.exp(-np.float32(logit))))) + evidence = _fake_scientific_evidence( + config=config, + candidate_count=len(item.candidates), + logit=logit, + probability=probability, + xpois=evidence_xpois, + variance_present=evidence_variance_present, + ) + background_count = next( + segment + for segment in evidence["layout"] + if segment["name"] == "xpois.background_coefficients" + )["shape"][0] + coefficient_count = ( + len(evidence["xpois"]["basis_terms"]) + background_count + ) + fit_pixel_count = coefficient_count + 9 + input_descriptors = [ + ("reference", item.reference), + ("target", item.target), + ] + if item.variance is not None: + input_descriptors.append(("variance", item.variance)) + if item.fit_mask is not None: + input_descriptors.append(("fit_mask", item.fit_mask)) + input_h2d_bytes = sum( + math.prod(descriptor.shape) + * (1 if name == "fit_mask" else np.dtype(np.float64).itemsize) + for name, descriptor in input_descriptors + ) + evidence_device_elements = evidence["packed_element_count"] - 2 * len( + item.candidates + ) + dlpack_shared_bytes = ( + len(item.candidates) + * 3 + * math.prod(config.stamp_shape) + * np.dtype(np.float32).itemsize + + len(item.candidates) + * len(FEATURE_NAMES) + * np.dtype(np.float32).itemsize + + evidence_device_elements * np.dtype(np.float64).itemsize + ) + payload = { + "schema": DEVICE_PIPELINE_RESULT_SCHEMA, + "item_id": item.item_id, + "device": config.device, + "checkpoint_sha256": config.checkpoint_sha256, + "feature_schema_sha256": config.feature_schema_sha256, + "configuration_sha256": config.configuration_sha256, + "input_sha256": { + name: descriptor.sha256 for name, descriptor in input_descriptors + }, + "xpois": { + "chi2": 1.0, + "dof": 9, + "fit_pixel_count": fit_pixel_count, + }, + "predictions": [ + { + **candidate.to_payload(), + "logit": logit, + "probability": probability, + "decision": True, + } + for candidate in item.candidates + ], + "transfers": DevicePipelineTransferReceipt( + input_h2d_bytes=input_h2d_bytes, + terminal_d2h_bytes=8, + dlpack_shared_bytes=dlpack_shared_bytes, + ).to_payload(), + "timings": { + "stages_seconds": { + name: 0.01 if name == "input_read" else 0.0 + for name in _TIMING_STAGES + }, + "total_seconds": 0.02, + "semantics": "host-elapsed; terminal-d2h-is-blocking", + }, + "scientific_evidence": evidence, + } + payload["transfers"]["terminal_d2h_bytes"] = payload[ + "scientific_evidence" + ]["packed_byte_count"] + return _seal_fake_result(payload) + + +def _fake_result( + item: DevicePipelineItem, config: DevicePipelineConfig +) -> Any: + payload = _fake_result_payload(item, config) + return SimpleNamespace(to_payload=lambda: payload) + + +class _FakeQueue(queue.Queue): + def close(self) -> None: + return None + + +class _FakePolicy: + Placement = SimpleNamespace(HOST_NAME="host-name") + + def __init__( + self, + *, + placement: Any, + host_name: str, + gpu_affinity: list[int], + ): + self.placement = placement + self.host_name = host_name + self.gpu_affinity = gpu_affinity + + +class _FakeTemplate: + created: list[_FakeTemplate] = [] + + def __init__(self, *, target: Any, args: tuple[Any, ...], policy: Any): + self.target = target + self.args = args + self.policy = policy + self.argdata = b"x" * 512 + type(self).created.append(self) + + +class _FakeGroup: + init_args: tuple[Any, ...] | None = None + init_kwargs: dict[str, Any] | None = None + + def __init__( + self, *, restart: bool, ignore_error_on_exit: bool, walltime: float + ) -> None: + assert restart is False + assert ignore_error_on_exit is False + assert walltime > 0 + self.templates: list[_FakeTemplate] = [] + self.exit_status: list[tuple[int, int]] = [] + + def add_process(self, *, nproc: int, template: _FakeTemplate) -> None: + assert nproc == 1 + self.templates.append(template) + + def init(self, *args: Any, **kwargs: Any) -> None: + type(self).init_args = args + type(self).init_kwargs = kwargs + + def start(self) -> None: + original = os.environ.get("CUDA_VISIBLE_DEVICES") + try: + for index, template in enumerate(self.templates): + os.environ["CUDA_VISIBLE_DEVICES"] = str( + template.policy.gpu_affinity[0] + ) + template.target(*template.args) + self.exit_status.append((1000 + index, 0)) + finally: + if original is None: + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + else: + os.environ["CUDA_VISIBLE_DEVICES"] = original + + def join(self, timeout: float | None = None) -> None: + assert timeout is not None and timeout > 0 + + def stop(self, patience: float = 5.0) -> None: + assert patience == 5.0 + + def close(self, patience: float = 5.0) -> None: + assert patience == 5.0 + + @property + def inactive_puids(self) -> list[tuple[int, int]]: + return self.exit_status + + +class _CorruptingSummaryGroup(_FakeGroup): + def start(self) -> None: + super().start() + descriptor_path = Path(self.templates[0].args[1]) + run_dir = descriptor_path.parent.parent + descriptor = json.loads(descriptor_path.read_text(encoding="utf-8")) + item_id = descriptor["items"][0]["item_id"] + summary_path = run_dir / "items" / item_id / "summary.json" + summary = json.loads(summary_path.read_text(encoding="utf-8")) + summary["scientific_evidence"]["candidate_count"] += 1 + _seal_fake_result(summary) + summary_path.write_text( + json.dumps(summary, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + record_path = run_dir / "records" / f"{item_id}.json" + record = json.loads(record_path.read_text(encoding="utf-8")) + record["device_pipeline"]["summary_sha256"] = file_sha256( + summary_path + ) + record["device_pipeline"]["result_sha256"] = summary["result_sha256"] + record_path.write_text( + json.dumps(record, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +class _CorruptingRecordDeviceGroup(_FakeGroup): + def start(self) -> None: + super().start() + descriptor_path = Path(self.templates[0].args[1]) + run_dir = descriptor_path.parent.parent + descriptor = json.loads(descriptor_path.read_text(encoding="utf-8")) + item_id = descriptor["items"][0]["item_id"] + summary_path = run_dir / "items" / item_id / "summary.json" + summary = json.loads(summary_path.read_text(encoding="utf-8")) + summary["device"] = "cuda:7" + _seal_fake_result(summary) + summary_path.write_text( + json.dumps(summary, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + record_path = run_dir / "records" / f"{item_id}.json" + record = json.loads(record_path.read_text(encoding="utf-8")) + record["device"] = "cuda:7" + record["device_pipeline"]["summary_sha256"] = file_sha256( + summary_path + ) + record["device_pipeline"]["result_sha256"] = summary["result_sha256"] + record_path.write_text( + json.dumps(record, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +class _FakeSystem: + nodes = (1,) + + +class _FakeNode: + hostname = socket.gethostname() + gpus = [3] + + def __init__(self, node_id: int): + assert node_id == 1 + + +def _install_fake_dragon( + monkeypatch: pytest.MonkeyPatch, + *, + context_type: type[Any], + group_type: type[Any] = _FakeGroup, +) -> None: + api = dragon_module._DragonAPI( + System=_FakeSystem, + Node=_FakeNode, + Policy=_FakePolicy, + ProcessGroup=group_type, + ProcessTemplate=_FakeTemplate, + Queue=_FakeQueue, + ) + monkeypatch.setattr(dragon_module, "_load_dragon_api", lambda: api) + monkeypatch.setattr( + dragon_pipeline, + "_collect_device_pipeline_gpu_identity", + lambda backend: { + "backend": backend, + "device_index": 0, + "name": "fake GPU", + "uuid": "GPU-fake", + "pci_bus_id": "00000000:03:00.0", + "identity_error": None, + "identity_warnings": [], + }, + ) + execute_shard = dragon_module._execute_shard + + def execute_without_import_check(**kwargs: Any) -> dict[str, Any]: + return execute_shard(require_clean_cuda_imports=False, **kwargs) + + monkeypatch.setattr( + dragon_module, "_execute_shard", execute_without_import_check + ) + monkeypatch.setattr(dragon_pipeline, "DeviceWorkerContext", context_type) + + +def test_device_pipeline_gpu_identity_loads_torch_before_cupy_probe( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[tuple[str, str]] = [] + identity = { + "backend": "cupy", + "device_index": 0, + "name": "fake GPU", + } + + monkeypatch.setattr( + dragon_pipeline.importlib, + "import_module", + lambda name: calls.append(("import", name)), + ) + + def collect(backend: str) -> dict[str, Any]: + calls.append(("identity", backend)) + return identity + + monkeypatch.setattr(dragon_module, "_collect_gpu_identity", collect) + + assert ( + dragon_pipeline._collect_device_pipeline_gpu_identity("cupy") + is identity + ) + assert calls == [("import", "torch"), ("identity", "cupy")] + + +def test_device_pipeline_gpu_identity_preserves_probe_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + dragon_pipeline.importlib, + "import_module", + lambda _name: object(), + ) + + def fail(_backend: str) -> dict[str, Any]: + raise RuntimeError("identity failed") + + monkeypatch.setattr(dragon_module, "_collect_gpu_identity", fail) + + with pytest.raises(RuntimeError, match="identity failed"): + dragon_pipeline._collect_device_pipeline_gpu_identity("cupy") + + +def test_dragon_pipeline_reuses_one_context_with_released_dragon_init( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = _config(tmp_path) + items = (_item(tmp_path, "pair-0"), _item(tmp_path, "pair-1")) + contexts: list[Any] = [] + calls: list[tuple[str, Any]] = [] + + class FakeContext: + load_seconds = 0.125 + + def __init__(self, initialized_config: DevicePipelineConfig): + self.config = initialized_config + + @classmethod + def initialize(cls, initialized_config: DevicePipelineConfig) -> Any: + context = cls(initialized_config) + contexts.append(context) + return context + + _FakeTemplate.created = [] + _FakeGroup.init_args = None + _FakeGroup.init_kwargs = None + _install_fake_dragon(monkeypatch, context_type=FakeContext) + + def run_item(item: DevicePipelineItem, context: Any) -> Any: + calls.append((item.item_id, context)) + return _fake_result(item, context.config) + + monkeypatch.setattr(dragon_pipeline, "run_device_pipeline_item", run_item) + result = dragon_pipeline.run_dragon_device_pipeline( + items=items, + config=config, + output_root=tmp_path / "runs", + run_id="adapter-test", + max_workers=1, + result_timeout_sec=1.0, + worker_timeout_sec=10.0, + ) + + assert result.status == "success" + assert len(contexts) == 1 + assert [item_id for item_id, _ in calls] == ["pair-0", "pair-1"] + assert all(context is contexts[0] for _, context in calls) + assert _FakeGroup.init_args == () + assert _FakeGroup.init_kwargs == {} + assert len(_FakeTemplate.created) == 1 + assert len(_FakeTemplate.created[0].args) == 5 + assert "candidates" not in repr(_FakeTemplate.created[0].args[:4]) + launch = json.loads( + (result.run_dir / "launch" / "worker-0000.json").read_text() + ) + assert launch["worker_target"] == { + "module": "cuphoton.xscan.dragon_pipeline", + "qualname": "_dragon_device_pipeline_worker", + } + assert launch["options"]["configuration"] == config.to_payload() + assert len(launch["items"]) == 2 + record = json.loads( + (result.run_dir / "records" / "pair-0.json").read_text() + ) + assert record["device_pipeline"]["candidate_count"] == 1 + assert record["device_pipeline"]["positive_count"] == 1 + item_summary = json.loads( + (result.run_dir / "items" / "pair-0" / "summary.json").read_text() + ) + assert item_summary["schema"] == DEVICE_PIPELINE_RESULT_SCHEMA + evidence = item_summary["scientific_evidence"] + assert evidence["schema"] == DEVICE_PIPELINE_EVIDENCE_SCHEMA + assert ( + record["device_pipeline"]["scientific_evidence_sha256"] + == (evidence["packed_sha256"]) + ) + assert ( + record["device_pipeline"]["scientific_evidence_bytes"] + == evidence["packed_byte_count"] + ) + queued = json.dumps(result.summary["shard_results"], sort_keys=True) + assert "scientific_evidence" not in queued + assert "packed_base64" not in queued + assert result.summary["terminal_record_errors"] == [] + + +def test_dragon_pipeline_caches_context_initialization_failure( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + config = _config(tmp_path) + items = (_item(tmp_path, "pair-0"), _item(tmp_path, "pair-1")) + attempts = 0 + + class FailedContext: + @classmethod + def initialize(cls, initialized_config: DevicePipelineConfig) -> Any: + nonlocal attempts + del initialized_config + attempts += 1 + raise RuntimeError("context load failed") + + _install_fake_dragon(monkeypatch, context_type=FailedContext) + monkeypatch.setattr( + dragon_pipeline, + "run_device_pipeline_item", + lambda *_args: pytest.fail( + "items must not run after context failure" + ), + ) + result = dragon_pipeline.run_dragon_device_pipeline( + items=items, + config=config, + output_root=tmp_path / "runs", + run_id="context-failure", + max_workers=1, + result_timeout_sec=1.0, + worker_timeout_sec=10.0, + ) + + assert result.status == "failed" + assert attempts == 1 + assert result.summary["terminal_record_audit"]["failed_item_ids"] == [ + "pair-0", + "pair-1", + ] + for item in items: + record = json.loads( + (result.run_dir / "records" / f"{item.item_id}.json").read_text() + ) + assert record["error"]["message"] == "context load failed" + assert not (result.run_dir / "items" / item.item_id).exists() + + +def test_dragon_pipeline_rejects_nonlocal_device_before_dragon( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + config = _config(tmp_path, device="cuda:1") + item = _item(tmp_path, "pair-0") + monkeypatch.setattr( + dragon_module, + "run_dragon_work_items", + lambda **_kwargs: pytest.fail("Dragon must not load"), + ) + + with pytest.raises(ValueError, match="local CUDA ordinal 0"): + dragon_pipeline.run_dragon_device_pipeline( + items=(item,), + config=config, + output_root=tmp_path / "runs", + ) + + +def test_dragon_pipeline_rejects_changed_input_before_dragon( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + config = _config(tmp_path) + item = _item(tmp_path, "pair-0") + np.save(item.reference.path, np.ones((31, 31), dtype=np.float32)) + monkeypatch.setattr( + dragon_module, + "run_dragon_work_items", + lambda **_kwargs: pytest.fail("Dragon must not load"), + ) + + with pytest.raises(RuntimeError, match="changed before Dragon launch"): + dragon_pipeline.run_dragon_device_pipeline( + items=(item,), + config=config, + output_root=tmp_path / "runs", + ) + + +def test_worker_options_and_nested_item_identity_are_strict( + tmp_path: Path, +) -> None: + config = _config(tmp_path) + options = { + "schema": dragon_pipeline.DRAGON_DEVICE_PIPELINE_OPTIONS_SCHEMA, + "backend": "cupy", + "configuration": config.to_payload(), + "configuration_sha256": config.configuration_sha256, + } + assert dragon_pipeline._strict_worker_options(options) == config + with pytest.raises(ValueError, match="invalid fields"): + dragon_pipeline._strict_worker_options({**options, "extra": True}) + + item = _item(tmp_path, "pair-0") + work_item = dragon_module.WorkItem( + item_id="different", payload=item.to_payload(), weight_bytes=1 + ) + context = SimpleNamespace(load_seconds=0.0) + with pytest.raises(ValueError, match="IDs differ"): + dragon_pipeline._publish_item_result( + work_item=work_item, + item_dir=tmp_path / "item-output", + config=config, + context=context, + ) + + +def test_publish_accepts_exact_transfer_bytes_for_all_input_roles( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + config = _config(tmp_path) + base_item = _item(tmp_path, "pair-0") + variance_path = tmp_path / "variance.npy" + fit_mask_path = tmp_path / "fit-mask.npy" + np.save(variance_path, np.ones((31, 31), dtype=np.float32)) + np.save(fit_mask_path, np.ones((31, 31), dtype=bool)) + item = DevicePipelineItem( + item_id=base_item.item_id, + reference=base_item.reference, + target=base_item.target, + variance=_descriptor(variance_path), + fit_mask=_descriptor(fit_mask_path), + candidates=( + *base_item.candidates, + DevicePipelineCandidate( + candidate_id="pair-0-candidate-2", + center_x=16, + center_y=15, + source_index=1, + ), + ), + ) + payload = _fake_result_payload(item, config) + monkeypatch.setattr( + dragon_pipeline, + "run_device_pipeline_item", + lambda *_args: SimpleNamespace(to_payload=lambda: payload), + ) + + metadata = dragon_pipeline._publish_item_result( + work_item=dragon_module.WorkItem( + item_id=item.item_id, + payload=item.to_payload(), + weight_bytes=1, + ), + item_dir=tmp_path / "item-output", + config=config, + context=SimpleNamespace(load_seconds=0.0), + ) + + transfers = metadata["device_pipeline"]["transfers"] + assert transfers["input_h2d_bytes"] == 3 * 31 * 31 * 8 + 31 * 31 + evidence_elements = payload["scientific_evidence"]["packed_element_count"] + assert transfers["dlpack_shared_bytes"] == ( + 2 * 3 * 11 * 11 * 4 + + 2 * len(FEATURE_NAMES) * 4 + + (evidence_elements - 4) * 8 + ) + + +@pytest.mark.parametrize( + ("corruption", "error_type", "message"), + [ + ("top-level", ValueError, "result has invalid fields"), + ("prediction", ValueError, "prediction 0 has invalid fields"), + ("prediction-type", TypeError, "prediction 0 logit must be a float"), + ("prediction-value", ValueError, "prediction decision differs"), + ( + "prediction-alignment", + ValueError, + "predictions and scientific evidence differ", + ), + ("xpois", ValueError, "XPOIS result has invalid fields"), + ("xpois-type", TypeError, "XPOIS dof must be an integer"), + ("timing", ValueError, "timing receipt has invalid fields"), + ("timing-value", ValueError, "total timing must be at least 0.0"), + ("transfer", ValueError, "transfer receipt has invalid fields"), + ("transfer-value", ValueError, "transfer receipt differs"), + ("input-h2d-bytes", ValueError, "input H2D byte count differs"), + ("dlpack-bytes", ValueError, "DLPack shared byte count differs"), + ("evidence", ValueError, "packed SHA-256 changed"), + ("result", ValueError, "result SHA-256 differs"), + ], +) +def test_publish_rejects_corrupt_v2_result_before_creating_output( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + corruption: str, + error_type: type[Exception], + message: str, +) -> None: + config = _config(tmp_path) + item = _item(tmp_path, "pair-0") + payload = _fake_result_payload(item, config) + if corruption == "top-level": + payload["unexpected"] = True + _seal_fake_result(payload) + elif corruption == "prediction": + payload["predictions"][0]["unexpected"] = True + _seal_fake_result(payload) + elif corruption == "prediction-type": + payload["predictions"][0]["logit"] = 1 + _seal_fake_result(payload) + elif corruption == "prediction-value": + payload["predictions"][0]["decision"] = False + _seal_fake_result(payload) + elif corruption == "prediction-alignment": + payload["predictions"][0]["logit"] = 0.5 + _seal_fake_result(payload) + elif corruption == "xpois": + payload["xpois"]["unexpected"] = True + _seal_fake_result(payload) + elif corruption == "xpois-type": + payload["xpois"]["dof"] = True + _seal_fake_result(payload) + elif corruption == "timing": + payload["timings"]["unexpected"] = True + elif corruption == "timing-value": + payload["timings"]["total_seconds"] = -1.0 + elif corruption == "transfer": + payload["transfers"]["unexpected"] = True + _seal_fake_result(payload) + elif corruption == "transfer-value": + payload["transfers"]["terminal_d2h_calls"] = 0 + _seal_fake_result(payload) + elif corruption == "input-h2d-bytes": + payload["transfers"]["input_h2d_bytes"] += 1 + _seal_fake_result(payload) + elif corruption == "dlpack-bytes": + payload["transfers"]["dlpack_shared_bytes"] += 1 + _seal_fake_result(payload) + elif corruption == "evidence": + payload["scientific_evidence"]["packed_sha256"] = "0" * 64 + _seal_fake_result(payload) + else: + payload["result_sha256"] = "0" * 64 + monkeypatch.setattr( + dragon_pipeline, + "run_device_pipeline_item", + lambda *_args: SimpleNamespace(to_payload=lambda: payload), + ) + item_dir = tmp_path / "item-output" + + with pytest.raises(error_type, match=message): + dragon_pipeline._publish_item_result( + work_item=dragon_module.WorkItem( + item_id=item.item_id, + payload=item.to_payload(), + weight_bytes=1, + ), + item_dir=item_dir, + config=config, + context=SimpleNamespace(load_seconds=0.0), + ) + + assert not item_dir.exists() + + +@pytest.mark.parametrize( + "corruption", + [ + "kernel-shape", + "flux-policy", + "flux-reference-order", + "background-count", + "variance-feature", + ], +) +def test_publish_rejects_rehashed_evidence_not_bound_to_config( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + corruption: str, +) -> None: + if corruption == "flux-reference-order": + xpois = DeviceXPOISPipelineConfig( + kernel_shape=(3, 3), + basis_sigmas=(0.8, 1.2), + basis_degrees=(0, 0), + flux_conserve=True, + flux_reference_index=1, + ) + elif corruption == "background-count": + xpois = DeviceXPOISPipelineConfig( + kernel_shape=(3, 3), + basis_sigmas=(0.8,), + basis_degrees=(0,), + background_degree=1, + ) + else: + xpois = DeviceXPOISPipelineConfig( + kernel_shape=(3, 3), + basis_sigmas=(0.8,), + basis_degrees=(0,), + ) + config = _config(tmp_path, xpois=xpois) + item = _item(tmp_path, "pair-0") + evidence_xpois = None + if corruption == "background-count": + evidence_xpois = DeviceXPOISPipelineConfig( + kernel_shape=(3, 3), + basis_sigmas=(0.8,), + basis_degrees=(0,), + background_degree=0, + ) + payload = _fake_result_payload( + item, + config, + evidence_xpois=evidence_xpois, + evidence_variance_present=corruption == "variance-feature", + ) + evidence = payload["scientific_evidence"] + if corruption == "kernel-shape": + evidence["xpois"]["kernel_shape"] = [5, 5] + elif corruption == "flux-policy": + evidence["xpois"]["flux_conserve"] = True + elif corruption == "flux-reference-order": + terms = evidence["xpois"]["basis_terms"] + assert [term["component_index"] for term in terms] == [1, 0] + evidence["xpois"]["basis_terms"] = [ + {**terms[1], "zero_sum": False}, + {**terms[0], "zero_sum": True}, + ] + _seal_fake_result(payload) + monkeypatch.setattr( + dragon_pipeline, + "run_device_pipeline_item", + lambda *_args: SimpleNamespace(to_payload=lambda: payload), + ) + item_dir = tmp_path / "item-output" + + with pytest.raises(ValueError, match="scientific evidence"): + dragon_pipeline._publish_item_result( + work_item=dragon_module.WorkItem( + item_id=item.item_id, + payload=item.to_payload(), + weight_bytes=1, + ), + item_dir=item_dir, + config=config, + context=SimpleNamespace(load_seconds=0.0), + ) + + assert not item_dir.exists() + + +def test_coordinator_revalidates_file_backed_v2_result( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + config = _config(tmp_path) + item = _item(tmp_path, "pair-0") + + class FakeContext: + load_seconds = 0.125 + + def __init__(self, initialized_config: DevicePipelineConfig): + self.config = initialized_config + + @classmethod + def initialize(cls, initialized_config: DevicePipelineConfig) -> Any: + return cls(initialized_config) + + _install_fake_dragon( + monkeypatch, + context_type=FakeContext, + group_type=_CorruptingSummaryGroup, + ) + monkeypatch.setattr( + dragon_pipeline, + "run_device_pipeline_item", + lambda value, context: _fake_result(value, context.config), + ) + + result = dragon_pipeline.run_dragon_device_pipeline( + items=(item,), + config=config, + output_root=tmp_path / "runs", + run_id="corrupt-file-backed-result", + max_workers=1, + result_timeout_sec=1.0, + worker_timeout_sec=10.0, + ) + + assert result.status == "failed" + assert result.summary["shard_result_audit"]["ok"] is True + assert result.summary["process_exit_audit"]["ok"] is True + assert result.summary["terminal_record_errors"] == [ + { + "item_id": item.item_id, + "type": "InvalidTerminalRecord", + "message": "record 0 has invalid field(s): device_pipeline", + } + ] + + +def test_coordinator_rejects_record_device_mismatch( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + config = _config(tmp_path) + item = _item(tmp_path, "pair-0") + + class FakeContext: + load_seconds = 0.125 + + def __init__(self, initialized_config: DevicePipelineConfig): + self.config = initialized_config + + @classmethod + def initialize(cls, initialized_config: DevicePipelineConfig) -> Any: + return cls(initialized_config) + + _install_fake_dragon( + monkeypatch, + context_type=FakeContext, + group_type=_CorruptingRecordDeviceGroup, + ) + monkeypatch.setattr( + dragon_pipeline, + "run_device_pipeline_item", + lambda value, context: _fake_result(value, context.config), + ) + + result = dragon_pipeline.run_dragon_device_pipeline( + items=(item,), + config=config, + output_root=tmp_path / "runs", + run_id="corrupt-record-device", + max_workers=1, + result_timeout_sec=1.0, + worker_timeout_sec=10.0, + ) + + assert result.status == "failed" + assert result.summary["shard_result_audit"]["ok"] is True + assert result.summary["process_exit_audit"]["ok"] is True + assert result.summary["terminal_record_errors"] == [ + { + "item_id": item.item_id, + "type": "InvalidTerminalRecord", + "message": ( + "record 0 has invalid field(s): device, device_pipeline" + ), + } + ] From 6348319ee5258a831beaa145f1ec5c8adac49348 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 22 Sep 2026 21:30:08 -0700 Subject: [PATCH 2/2] Reject overflowing Dragon runtime metadata Treat JSON integers outside floating range as invalid runtime metadata. This lets the coordinator persist a failed terminal summary when a worker's context-load duration cannot be represented as a finite number. Signed-off-by: Trent Nelson --- src/cuphoton/xscan/dragon_pipeline.py | 13 ++++----- tests/xscan/test_dragon_pipeline.py | 39 ++++++++++++++++++++++++--- 2 files changed, 43 insertions(+), 9 deletions(-) diff --git a/src/cuphoton/xscan/dragon_pipeline.py b/src/cuphoton/xscan/dragon_pipeline.py index 0244c345..f83bf077 100644 --- a/src/cuphoton/xscan/dragon_pipeline.py +++ b/src/cuphoton/xscan/dragon_pipeline.py @@ -903,12 +903,13 @@ def item_runner( def _finite_nonnegative(value: Any) -> bool: - return ( - not isinstance(value, bool) - and isinstance(value, (int, float)) - and math.isfinite(value) - and value >= 0 - ) + if isinstance(value, bool) or not isinstance(value, (int, float)): + return False + try: + return math.isfinite(value) and value >= 0 + except OverflowError: + # JSON integers need not fit the float conversion used by isfinite. + return False def _success_record_problems( diff --git a/tests/xscan/test_dragon_pipeline.py b/tests/xscan/test_dragon_pipeline.py index e03f9059..2db9e2dd 100644 --- a/tests/xscan/test_dragon_pipeline.py +++ b/tests/xscan/test_dragon_pipeline.py @@ -581,6 +581,22 @@ def start(self) -> None: ) +class _OverflowingRuntimeRecordGroup(_FakeGroup): + def start(self) -> None: + super().start() + descriptor_path = Path(self.templates[0].args[1]) + run_dir = descriptor_path.parent.parent + descriptor = json.loads(descriptor_path.read_text(encoding="utf-8")) + item_id = descriptor["items"][0]["item_id"] + record_path = run_dir / "records" / f"{item_id}.json" + record = json.loads(record_path.read_text(encoding="utf-8")) + record["runtime"]["context_load_seconds"] = 10**1000 + record_path.write_text( + json.dumps(record, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + class _CorruptingRecordDeviceGroup(_FakeGroup): def start(self) -> None: super().start() @@ -1143,8 +1159,19 @@ def test_publish_rejects_rehashed_evidence_not_bound_to_config( assert not item_dir.exists() +@pytest.mark.parametrize( + ("group_type", "invalid_field"), + [ + (_CorruptingSummaryGroup, "device_pipeline"), + (_OverflowingRuntimeRecordGroup, "runtime"), + ], + ids=["corrupt-summary", "overflowing-runtime"], +) def test_coordinator_revalidates_file_backed_v2_result( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + group_type: type[_FakeGroup], + invalid_field: str, ) -> None: config = _config(tmp_path) item = _item(tmp_path, "pair-0") @@ -1162,7 +1189,7 @@ def initialize(cls, initialized_config: DevicePipelineConfig) -> Any: _install_fake_dragon( monkeypatch, context_type=FakeContext, - group_type=_CorruptingSummaryGroup, + group_type=group_type, ) monkeypatch.setattr( dragon_pipeline, @@ -1187,9 +1214,15 @@ def initialize(cls, initialized_config: DevicePipelineConfig) -> Any: { "item_id": item.item_id, "type": "InvalidTerminalRecord", - "message": "record 0 has invalid field(s): device_pipeline", + "message": f"record 0 has invalid field(s): {invalid_field}", } ] + persisted = json.loads(result.summary_path.read_text(encoding="utf-8")) + assert persisted["status"] == "failed" + assert ( + persisted["terminal_record_errors"] + == result.summary["terminal_record_errors"] + ) def test_coordinator_rejects_record_device_mismatch(