From ffdf63cb5a9746c6c981ccf33039d89e762e78d3 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Fri, 21 Aug 2026 19:02:06 -0700 Subject: [PATCH 1/4] Add Dragon-backed XPOIS image-pair batches Run complete image-pair fits on explicitly placed GPU workers from a validated manifest. Balance shards by input size and record complete per-item results in immutable run directories. Signed-off-by: Trent Nelson --- THIRD_PARTY_NOTICES.md | 11 + docs/components/xpois.md | 101 ++ examples/xpois/dragon_batch.py | 16 + src/cuphoton/core/bulk.py | 270 +++++ src/cuphoton/xpois/batch.py | 868 ++++++++++++++++ src/cuphoton/xpois/commands.py | 255 +++-- src/cuphoton/xpois/dragon.py | 1327 ++++++++++++++++++++++++ src/cuphoton/xpois/workflows.py | 113 +- tests/core/test_bulk.py | 140 +++ tests/core/test_cli_contract.py | 10 +- tests/xpois/test_batch.py | 321 ++++++ tests/xpois/test_cli.py | 134 +++ tests/xpois/test_dragon.py | 1712 +++++++++++++++++++++++++++++++ 13 files changed, 5164 insertions(+), 114 deletions(-) create mode 100644 examples/xpois/dragon_batch.py create mode 100644 src/cuphoton/core/bulk.py create mode 100644 src/cuphoton/xpois/batch.py create mode 100644 src/cuphoton/xpois/dragon.py create mode 100644 tests/core/test_bulk.py create mode 100644 tests/xpois/test_batch.py create mode 100644 tests/xpois/test_dragon.py diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 5aaad4c..c501cd5 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -54,6 +54,17 @@ more than one license. | `dev` | `pytest>=8.3` | `9.1.1` | `MIT` | [pytest](https://github.com/pytest-dev/pytest) | `uv / PyPI` | | `dev` | `ruff>=0.15.12` | `0.15.20` | `MIT` | [Ruff](https://github.com/astral-sh/ruff) | `uv / PyPI` | +## Optional distributed runtime inventory + +XPOIS distributed execution requires the runtime selected by the caller. +These runtimes are installed separately from cuPhoton and are not included +in its wheel or dependency lock. See the +[XPOIS guide](docs/components/xpois.md) for installation and launch details. + +| Runtime | License and upstream notices | Use and installation | +| --- | --- | --- | +| DragonHPC (`dragonhpc`; import `dragon`) | [MIT](https://github.com/DragonHPC/dragon/blob/0.14.2/LICENSE) | Required for Dragon workers, placement and communication. Install separately in cuPhoton's Python environment; no Dragon source or binaries are bundled. | + ## Native system dependency inventory | Package | Version or version range | License identifier | Upstream | Use in cuPhoton | Distribution | diff --git a/docs/components/xpois.md b/docs/components/xpois.md index 5f8863b..69c625d 100644 --- a/docs/components/xpois.md +++ b/docs/components/xpois.md @@ -134,3 +134,104 @@ represent neighboring-pixel covariance, reference-target covariance, resampling covariance, or fitted-kernel uncertainty. The standardized residual is a descriptive diagnostic, not a whitened residual or calibrated significance image. Prefer held-out pixels when assessing fit quality. + +## Dragon image-pair batches + +`fit-batch-dragon` distributes complete reference/target image-pair fits across +GPU workers. The coordinator reads a manifest, assigns one Dragon +ProcessGroup worker to each selected GPU, and balances work by input size. +Arrays and output files stay on shared storage; workers send compact result +records through Dragon queues. Each image-pair fit runs on one GPU. + +### Runtime dependency + +The Dragon executor requires [DragonHPC](https://dragonhpc.github.io/dragon/doc/_build/html/index.html) +(Python distribution `dragonhpc`, import `dragon`) in the same Python +environment as cuPhoton on every participating node. Installing cuPhoton, +including its `gpu` extra, does not install DragonHPC. DragonHPC is installed +separately and is not included in cuPhoton's dependency lock or wheel. + +For example, install the released DragonHPC 0.14.2 package into a CUDA 13 +cuPhoton environment: + +```bash +uv sync --locked --python 3.12 --extra gpu +uv pip install --python .venv/bin/python "dragonhpc==0.14.2" +``` + +The [DragonHPC 0.14.2 wheels](https://pypi.org/project/dragonhpc/0.14.2/#files) +support CPython 3.11 through 3.13 on Linux x86-64 and AArch64 with glibc 2.28 +or newer. Use one of those Python versions for this installation; a CPython +3.14 wheel is not published for this DragonHPC release. + +Use the installed `.venv/bin/dragon` launcher with the examples below. +`uv sync` removes packages outside the project lock, so repeat the DragonHPC +installation after resynchronizing the environment. Use the same cuPhoton +environment and DragonHPC version on every node. Each run records the +DragonHPC version it discovers. See the +[runtime notices](../../THIRD_PARTY_NOTICES.md#optional-distributed-runtime-inventory) +for licensing and installation details. + +### Manifest and launch + +Inputs and output directories must be accessible at the same paths on every +node. The command accepts a strict JSON or YAML manifest: + +```yaml +schema: cuphoton.xpois.image-pairs/v1 +pairs: + - id: detector-0001 + reference: /shared/input/reference-0001.fits + target: /shared/input/target-0001.fits + reference_hdu: 1 + target_hdu: 1 + variance: /shared/input/variance-0001.fits + variance_hdu: 1 +``` + +Loading the manifest records each unique input's resolved path, byte size, +and nanosecond modification time. These values are checked after coordinator +preflight and before and after each item. A change that preserves both size +and modification time is not detected; use immutable input data or external +checksums when content identity matters. + +Launch through Dragon within an existing scheduler allocation: + +```bash +.venv/bin/dragon examples/xpois/dragon_batch.py \ + --manifest /shared/manifests/image-pairs.yaml \ + --output-dir /shared/results/xpois-dragon \ + --name image-pairs-gpu4 \ + --max-workers 4 \ + --worker-timeout-sec 3600 \ + --backend cupy +``` + +The wrapper invokes `cuphoton xpois fit-batch-dragon`. This command requires +an explicit GPU backend: `cupy`, `numba-cuda`, or `cutile`. It rejects `auto` +instead of falling back to CPU when a GPU package is unavailable. + +The coordinator enumerates actual `Node.gpus` IDs and selects workers +round-robin across hosts. Each worker checks its host and singleton +`CUDA_VISIBLE_DEVICES` assignment before importing the numerical backend. +Loopback hostname aliases are accepted only for a single-node allocation. + +### Results and limits + +Every attempt uses a new, immutable run directory. Each item has an atomic +terminal record under `records/`; an ordinary item error does not prevent +the remaining items in its shard from running. The final `summary.json` +checks for missing, duplicate, unexpected, malformed, or assignment-inconsistent +item and worker results, as well as nonzero worker exits. Worker wall time +defaults to one hour. The coordinator attempts bounded stop and close cleanup +after failed starts or joins. + +`coordinator_wall_sec` includes setup, worker cleanup, and terminal-result +checks. Writing the final summary falls outside that interval. Per-item +`timings_sec` separates input reads, preprocessing, solving, artifact writes, +review work, and other postprocessing; `wall_sec` retains the enclosing +workflow and item-runner measurements. + +The executor does not retry failed items, recover dead workers, or split one +image-pair solve across GPUs. Inspect the terminal status before consuming a +run's outputs, and start a new run after a failed attempt. diff --git a/examples/xpois/dragon_batch.py b/examples/xpois/dragon_batch.py new file mode 100644 index 0000000..58e4d57 --- /dev/null +++ b/examples/xpois/dragon_batch.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Run XPOIS image-pair batches with Dragon workers.""" + +from __future__ import annotations + +import sys + +from cuphoton.core.cli import run_component + +if __name__ == "__main__": + raise SystemExit( + run_component("xpois", ["fit-batch-dragon", *sys.argv[1:]]) + ) diff --git a/src/cuphoton/core/bulk.py b/src/cuphoton/core/bulk.py new file mode 100644 index 0000000..ca9e6a3 --- /dev/null +++ b/src/cuphoton/core/bulk.py @@ -0,0 +1,270 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Durable contracts for independent whole-item bulk work.""" + +from __future__ import annotations + +import json +import os +import re +import uuid +from collections import Counter +from collections.abc import Mapping, Sequence +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +_IDENTIFIER = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,127}") + + +@dataclass(frozen=True) +class WorkItem: + """One independent, JSON-described unit of work.""" + + item_id: str + payload: Mapping[str, Any] + weight_bytes: int = 0 + + def __post_init__(self) -> None: + validate_identifier(self.item_id, field="item_id") + if ( + isinstance(self.weight_bytes, bool) + or not isinstance(self.weight_bytes, int) + or self.weight_bytes < 0 + ): + raise ValueError("weight_bytes must be a non-negative integer") + object.__setattr__( + self, + "payload", + json_mapping(self.payload, field="work item payload"), + ) + + def to_dict(self) -> dict[str, Any]: + """Return the JSON-compatible wire representation.""" + + return { + "item_id": self.item_id, + "payload": dict(self.payload), + "weight_bytes": self.weight_bytes, + } + + @classmethod + def from_dict(cls, payload: Mapping[str, Any]) -> WorkItem: + """Restore an item from its wire representation.""" + + return cls( + item_id=str(payload["item_id"]), + payload=payload["payload"], + weight_bytes=payload.get("weight_bytes", 0), + ) + + +@dataclass(frozen=True) +class Placement: + """One explicitly selected Dragon host/GPU worker placement.""" + + worker_id: int + host: str + gpu_id: int + + def __post_init__(self) -> None: + if isinstance(self.worker_id, bool) or not isinstance( + self.worker_id, int + ): + raise TypeError("worker_id must be an integer") + if self.worker_id < 0: + raise ValueError("worker_id must be non-negative") + if not self.host: + raise ValueError("placement host must not be empty") + if isinstance(self.gpu_id, bool) or not isinstance(self.gpu_id, int): + raise TypeError("Dragon GPU IDs must be integers") + if self.gpu_id < 0: + raise ValueError("Dragon GPU IDs must be non-negative") + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-compatible placement record.""" + + return asdict(self) + + +def partition_byte_balanced( + items: Sequence[WorkItem], worker_count: int +) -> tuple[tuple[WorkItem, ...], ...]: + """Deterministically assign whole items by descending input bytes.""" + + if isinstance(worker_count, bool) or not isinstance(worker_count, int): + raise TypeError("worker_count must be an integer") + if worker_count <= 0: + raise ValueError("worker_count must be positive") + _require_unique_item_ids(items) + shards: list[list[WorkItem]] = [[] for _ in range(worker_count)] + weights = [0] * worker_count + for item in sorted( + items, + key=lambda value: (-value.weight_bytes, value.item_id), + ): + worker_id = min( + range(worker_count), + key=lambda value: ( + weights[value], + len(shards[value]), + value, + ), + ) + shards[worker_id].append(item) + weights[worker_id] += item.weight_bytes + return tuple(tuple(shard) for shard in shards) + + +def audit_terminal_records( + expected_item_ids: Sequence[str], + records: Sequence[Mapping[str, Any]], +) -> dict[str, Any]: + """Audit exactly one valid terminal record for every expected item.""" + + expected = list(expected_item_ids) + expected_counts = Counter(expected) + if any(count != 1 for count in expected_counts.values()): + raise ValueError("expected item IDs must be unique") + observed = [str(record.get("item_id", "")) for record in records] + observed_counts = Counter(observed) + expected_set = set(expected) + missing = sorted(expected_set - set(observed_counts)) + duplicates = sorted( + item_id for item_id, count in observed_counts.items() if count > 1 + ) + unexpected = sorted(set(observed_counts) - expected_set) + invalid_status = sorted( + str(record.get("item_id", "")) + for record in records + if record.get("status") not in {"success", "failed"} + ) + failed = sorted( + str(record.get("item_id", "")) + for record in records + if record.get("status") == "failed" + and str(record.get("item_id", "")) in expected_set + ) + ok = ( + not missing + and not duplicates + and not unexpected + and not invalid_status + ) + return { + "ok": ok, + "expected_count": len(expected), + "terminal_record_count": len(records), + "missing_item_ids": missing, + "duplicate_item_ids": duplicates, + "unexpected_item_ids": unexpected, + "invalid_status_item_ids": invalid_status, + "failed_item_ids": failed, + } + + +def atomic_write_json(path: Path, payload: Mapping[str, Any]) -> None: + """Atomically replace one JSON mapping with a trailing newline.""" + + normalized = json_mapping(payload, field=f"payload for {path}") + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name( + f".{path.name}.tmp-{os.getpid()}-{uuid.uuid4().hex}" + ) + try: + with temporary.open("w", encoding="utf-8") as handle: + handle.write( + json.dumps(normalized, indent=2, sort_keys=True) + "\n" + ) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + directory_flags = os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) + directory_fd = os.open(path.parent, directory_flags) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + finally: + temporary.unlink(missing_ok=True) + + +def read_json_mapping(path: Path) -> dict[str, Any]: + """Read one JSON object and reject other JSON value types.""" + + payload = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(payload, dict): + raise ValueError(f"expected a JSON object in {path}") + return payload + + +def json_mapping(payload: Mapping[str, Any], *, field: str) -> dict[str, Any]: + """Round-trip one mapping through strict JSON.""" + + if not isinstance(payload, Mapping): + raise TypeError(f"{field} must be a mapping") + try: + encoded = json.dumps(dict(payload), sort_keys=True, allow_nan=False) + decoded = json.loads(encoded) + except (TypeError, ValueError) as exc: + raise ValueError(f"{field} must be JSON-compatible: {exc}") from exc + if not isinstance(decoded, dict): + raise TypeError(f"{field} must encode a JSON object") + return decoded + + +def validate_identifier(value: str, *, field: str) -> str: + """Validate a filesystem-safe item or run identity.""" + + if not isinstance(value, str) or not _IDENTIFIER.fullmatch(value): + raise ValueError(f"{field} must match {_IDENTIFIER.pattern!r}") + if value in {".", ".."}: + raise ValueError(f"{field} cannot be {value!r}") + return value + + +def new_run_id(prefix: str) -> str: + """Create a collision-resistant UTC run identity.""" + + validate_identifier(prefix, field="run prefix") + stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + return f"{prefix}-{stamp}-{uuid.uuid4().hex[:8]}" + + +def timestamp_utc() -> str: + """Return the current timezone-aware UTC timestamp.""" + + return datetime.now(timezone.utc).isoformat() + + +def error_payload(exc: BaseException) -> dict[str, str]: + """Return a compact structured exception description.""" + + return {"type": type(exc).__name__, "message": str(exc)} + + +def _require_unique_item_ids(items: Sequence[WorkItem]) -> None: + counts = Counter(item.item_id for item in items) + duplicates = sorted( + item_id for item_id, count in counts.items() if count > 1 + ) + if duplicates: + raise ValueError("duplicate work item IDs: " + ", ".join(duplicates)) + + +__all__ = [ + "Placement", + "WorkItem", + "atomic_write_json", + "audit_terminal_records", + "error_payload", + "json_mapping", + "new_run_id", + "partition_byte_balanced", + "read_json_mapping", + "timestamp_utc", + "validate_identifier", +] diff --git a/src/cuphoton/xpois/batch.py b/src/cuphoton/xpois/batch.py new file mode 100644 index 0000000..b23c5c5 --- /dev/null +++ b/src/cuphoton/xpois/batch.py @@ -0,0 +1,868 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Whole-image batch contracts for XPOIS.""" + +from __future__ import annotations + +import hashlib +import json +import math +import time +from collections.abc import Mapping +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import yaml +from astropy.io import fits + +from cuphoton.core.bulk import WorkItem, validate_identifier + +from .data import ( + ERROR_EXTENSION_NAMES, + FITS_SUFFIXES, + MASK_EXTENSION_NAMES, + VARIANCE_EXTENSION_NAMES, +) + +IMAGE_PAIR_MANIFEST_SCHEMA = "cuphoton.xpois.image-pairs/v1" + +_ROOT_FIELDS = frozenset({"schema", "pairs"}) +_PAIR_REQUIRED_FIELDS = frozenset({"id", "reference", "target"}) +_PATH_FIELDS = ( + "reference", + "target", + "variance", + "reference_mask", + "target_mask", + "fit_mask", +) +_HDU_FIELDS = ( + "reference_hdu", + "target_hdu", + "variance_hdu", + "reference_mask_hdu", + "target_mask_hdu", +) +_PAIR_FIELDS = ( + _PAIR_REQUIRED_FIELDS | frozenset(_PATH_FIELDS) | frozenset(_HDU_FIELDS) +) +_SUPPORTED_BACKENDS = frozenset( + {"auto", "cpu", "cupy", "numba-cuda", "cutile"} +) +_MASK_POLICIES = frozenset({"none", "strict", "hsc-masklite", "masklite"}) +_INPUT_IDENTITY_KEY = "_input_identity" +_INPUT_IDENTITY_POLICY = "size-mtime-ns" +_INPUT_IDENTITY_LIMITATION = ( + "Content hashes are not computed by default; size and nanosecond mtime " + "cannot detect content changes that preserve both values." +) + + +class _UniqueKeySafeLoader(yaml.SafeLoader): + """YAML loader that rejects duplicate mapping keys.""" + + +def _construct_unique_yaml_mapping( + loader: yaml.SafeLoader, + node: yaml.MappingNode, + deep: bool = False, +) -> dict[Any, Any]: + loader.flatten_mapping(node) + result: dict[Any, Any] = {} + for key_node, value_node in node.value: + key = loader.construct_object(key_node, deep=deep) + try: + duplicate = key in result + except TypeError as exc: + raise yaml.constructor.ConstructorError( + "while constructing a mapping", + node.start_mark, + "found an unhashable key", + key_node.start_mark, + ) from exc + if duplicate: + raise yaml.constructor.ConstructorError( + "while constructing a mapping", + node.start_mark, + f"found duplicate key {key!r}", + key_node.start_mark, + ) + result[key] = loader.construct_object(value_node, deep=deep) + return result + + +_UniqueKeySafeLoader.add_constructor( + yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, + _construct_unique_yaml_mapping, +) + + +@dataclass(frozen=True) +class InputFileIdentity: + """Cheap, reproducible identity for one resolved input file.""" + + path: Path + size_bytes: int + mtime_ns: int + + @classmethod + def capture(cls, path: Path) -> InputFileIdentity: + """Capture size and nanosecond mtime without reading file content.""" + + resolved = path.expanduser().resolve() + stat = resolved.stat() + return cls( + path=resolved, + size_bytes=stat.st_size, + mtime_ns=stat.st_mtime_ns, + ) + + @classmethod + def from_payload(cls, payload: Mapping[str, Any]) -> InputFileIdentity: + """Restore a strict worker-side identity record.""" + + values = _require_mapping(payload, "input identity") + _reject_unknown_fields( + values, + frozenset({"path", "size_bytes", "mtime_ns"}), + "input identity", + ) + path_value = values.get("path") + if not isinstance(path_value, str) or not path_value: + raise ValueError("input identity path must be a non-empty string") + path = Path(path_value) + if not path.is_absolute(): + raise ValueError("input identity path must be absolute") + size_bytes = _require_non_negative_integer( + values.get("size_bytes"), field="input identity size_bytes" + ) + mtime_ns = _require_non_negative_integer( + values.get("mtime_ns"), field="input identity mtime_ns" + ) + return cls(path, size_bytes, mtime_ns) + + def to_payload(self) -> dict[str, Any]: + """Return a JSON-compatible identity record.""" + + return { + "path": str(self.path), + "size_bytes": self.size_bytes, + "mtime_ns": self.mtime_ns, + } + + def verify(self) -> None: + """Reject an input that changed since identity capture.""" + + try: + actual = self.capture(self.path) + except OSError as exc: + raise RuntimeError( + f"input changed since manifest load: {self.path}: {exc}" + ) from exc + if actual != self: + raise RuntimeError( + "input changed since manifest load: " + f"{self.path} expected size={self.size_bytes}, " + f"mtime_ns={self.mtime_ns}; got size={actual.size_bytes}, " + f"mtime_ns={actual.mtime_ns}" + ) + + +@dataclass(frozen=True) +class ImagePairSpec: + """One independent reference/target pair from a batch manifest.""" + + item_id: str + reference: Path + target: Path + reference_hdu: int | None = None + target_hdu: int | None = None + variance: Path | None = None + variance_hdu: int | None = None + reference_mask: Path | None = None + target_mask: Path | None = None + reference_mask_hdu: int | None = None + target_mask_hdu: int | None = None + fit_mask: Path | None = None + + def to_payload(self) -> dict[str, Any]: + """Return the canonical JSON-compatible representation.""" + + payload: dict[str, Any] = {"id": self.item_id} + for field in _PATH_FIELDS: + value = getattr(self, field) + payload[field] = str(value) if value is not None else None + for field in _HDU_FIELDS: + payload[field] = getattr(self, field) + return payload + + @classmethod + def from_payload(cls, payload: Mapping[str, Any]) -> ImagePairSpec: + """Restore a validated pair passed to a worker.""" + + values = dict(payload) + item_id = str(values.pop("id")) + for field in _PATH_FIELDS: + value = values.get(field) + values[field] = Path(value) if value is not None else None + return cls(item_id=item_id, **values) + + +@dataclass(frozen=True) +class ImagePairManifest: + """Validated image-pair manifest and its canonical identity.""" + + source_path: Path + pairs: tuple[ImagePairSpec, ...] + input_identities: tuple[InputFileIdentity, ...] + sha256: str + + def canonical_payload(self) -> dict[str, Any]: + """Return the resolved manifest used for execution.""" + + return { + "schema": IMAGE_PAIR_MANIFEST_SCHEMA, + "pairs": [pair.to_payload() for pair in self.pairs], + } + + def input_identity_payload(self) -> dict[str, Any]: + """Return the stat-based identity contract used for this execution.""" + + return { + "schema": "cuphoton.xpois.input-identity/v1", + "policy": _INPUT_IDENTITY_POLICY, + "content_hashes": False, + "limitation": _INPUT_IDENTITY_LIMITATION, + "files": [item.to_payload() for item in self.input_identities], + } + + def verify_input_identity(self) -> None: + """Reject files changed since the manifest was loaded.""" + + for identity in self.input_identities: + identity.verify() + + def work_items(self) -> tuple[WorkItem, ...]: + """Convert pairs to scheduler-neutral work items.""" + + by_path = { + identity.path: identity for identity in self.input_identities + } + items = [] + for pair in self.pairs: + paths = [ + path + for field in _PATH_FIELDS + if (path := getattr(pair, field)) is not None + ] + unique_paths = sorted(set(paths), key=str) + unique_identities = [ + by_path[path].to_payload() for path in unique_paths + ] + payload = pair.to_payload() + payload[_INPUT_IDENTITY_KEY] = unique_identities + items.append( + WorkItem( + item_id=pair.item_id, + payload=payload, + weight_bytes=sum( + by_path[path].size_bytes for path in unique_paths + ), + ) + ) + return tuple(items) + + +@dataclass(frozen=True) +class BatchFitOptions: + """Global fit options shared by all image pairs.""" + + kernel_shape: tuple[int, int] + basis_sigmas: tuple[float, ...] + basis_degrees: tuple[int, ...] + mask_policy: str = "none" + crop_y0: int | None = None + crop_x0: int | None = None + crop_height: int | None = None + crop_width: int | None = None + auto_stamp_mask: bool = False + auto_stamp_size: int = 31 + auto_stamp_count: int = 5 + auto_peak_percentile: float = 99.5 + background_degree: int = 0 + flux_conserve: bool = False + backend: str = "cupy" + + def __post_init__(self) -> None: + if len(self.kernel_shape) != 2 or any( + isinstance(value, bool) or value <= 0 or value % 2 == 0 + for value in self.kernel_shape + ): + raise ValueError( + "kernel_shape must contain two positive odd values" + ) + if not self.basis_sigmas or len(self.basis_sigmas) != len( + self.basis_degrees + ): + raise ValueError( + "basis_sigmas and basis_degrees must be non-empty and aligned" + ) + if any( + not math.isfinite(value) or value <= 0 + for value in self.basis_sigmas + ): + raise ValueError("basis sigmas must be finite and positive") + if any(value < 0 for value in self.basis_degrees): + raise ValueError("basis degrees must be non-negative") + if ( + isinstance(self.auto_stamp_size, bool) + or not isinstance(self.auto_stamp_size, int) + or self.auto_stamp_size <= 0 + or self.auto_stamp_size % 2 == 0 + ): + raise ValueError("auto_stamp_size must be a positive odd integer") + if ( + isinstance(self.auto_stamp_count, bool) + or not isinstance(self.auto_stamp_count, int) + or self.auto_stamp_count <= 0 + ): + raise ValueError("auto_stamp_count must be a positive integer") + if ( + isinstance(self.auto_peak_percentile, bool) + or not isinstance(self.auto_peak_percentile, (int, float)) + or not math.isfinite(self.auto_peak_percentile) + or not 0 <= self.auto_peak_percentile <= 100 + ): + raise ValueError( + "auto_peak_percentile must be a finite number from 0 to 100" + ) + if self.mask_policy not in _MASK_POLICIES: + raise ValueError(f"unsupported mask policy: {self.mask_policy}") + if self.backend not in _SUPPORTED_BACKENDS: + raise ValueError(f"unsupported fit backend: {self.backend}") + crop = ( + self.crop_y0, + self.crop_x0, + self.crop_height, + self.crop_width, + ) + if any(value is None for value in crop) and not all( + value is None for value in crop + ): + raise ValueError("specify either all crop parameters or none") + if self.crop_y0 is not None: + assert self.crop_x0 is not None + assert self.crop_height is not None + assert self.crop_width is not None + if self.crop_y0 < 0 or self.crop_x0 < 0: + raise ValueError("crop origins must be non-negative") + if self.crop_height <= 0 or self.crop_width <= 0: + raise ValueError("crop dimensions must be positive") + + def to_payload(self) -> dict[str, Any]: + """Return JSON-compatible worker options.""" + + payload = asdict(self) + payload["kernel_shape"] = list(self.kernel_shape) + payload["basis_sigmas"] = list(self.basis_sigmas) + payload["basis_degrees"] = list(self.basis_degrees) + return payload + + @classmethod + def from_payload(cls, payload: Mapping[str, Any]) -> BatchFitOptions: + """Restore options passed through Dragon.""" + + values = dict(payload) + values["kernel_shape"] = tuple(values["kernel_shape"]) + values["basis_sigmas"] = tuple(values["basis_sigmas"]) + values["basis_degrees"] = tuple(values["basis_degrees"]) + return cls(**values) + + +def load_image_pair_manifest(path: Path) -> ImagePairManifest: + """Load and strictly validate a JSON or YAML pair manifest.""" + + resolved = path.expanduser().resolve() + if not resolved.is_file(): + raise FileNotFoundError(f"manifest does not exist: {resolved}") + if resolved.suffix.lower() not in {".json", ".yaml", ".yml"}: + raise ValueError("image-pair manifest must be JSON or YAML") + try: + if resolved.suffix.lower() == ".json": + raw = json.loads( + resolved.read_text(encoding="utf-8"), + object_pairs_hook=_construct_unique_json_mapping, + ) + else: + raw = yaml.load( + resolved.read_text(encoding="utf-8"), + Loader=_UniqueKeySafeLoader, + ) + except (json.JSONDecodeError, yaml.YAMLError, ValueError) as exc: + raise ValueError(f"invalid image-pair manifest: {exc}") from exc + root = _require_mapping(raw, "image-pair manifest") + _reject_unknown_fields(root, _ROOT_FIELDS, "image-pair manifest") + if root.get("schema") != IMAGE_PAIR_MANIFEST_SCHEMA: + raise ValueError( + f"manifest schema must be {IMAGE_PAIR_MANIFEST_SCHEMA!r}" + ) + raw_pairs = root.get("pairs") + if not isinstance(raw_pairs, list) or not raw_pairs: + raise ValueError("manifest pairs must be a non-empty list") + + pairs: list[ImagePairSpec] = [] + seen_ids: set[str] = set() + for index, raw_pair in enumerate(raw_pairs): + entry = _require_mapping(raw_pair, f"pairs[{index}]") + _reject_unknown_fields(entry, _PAIR_FIELDS, f"pairs[{index}]") + missing = sorted(_PAIR_REQUIRED_FIELDS - set(entry)) + if missing: + raise ValueError( + f"pairs[{index}] is missing required field(s): " + + ", ".join(missing) + ) + item_id = validate_identifier(entry["id"], field=f"pairs[{index}].id") + if item_id in seen_ids: + raise ValueError(f"duplicate image-pair ID: {item_id!r}") + seen_ids.add(item_id) + values: dict[str, Any] = {"item_id": item_id} + for field in _PATH_FIELDS: + values[field] = _resolve_input_path( + entry.get(field), + base_dir=resolved.parent, + field=f"pairs[{index}].{field}", + required=field in {"reference", "target"}, + ) + for field in _HDU_FIELDS: + values[field] = _validate_hdu( + entry.get(field), field=f"pairs[{index}].{field}" + ) + _validate_hdu_dependencies(values, index=index) + pairs.append(ImagePairSpec(**values)) + + canonical = { + "schema": IMAGE_PAIR_MANIFEST_SCHEMA, + "pairs": [pair.to_payload() for pair in pairs], + } + input_paths = sorted( + { + path + for pair in pairs + for field in _PATH_FIELDS + if (path := getattr(pair, field)) is not None + }, + key=str, + ) + input_identities = tuple( + InputFileIdentity.capture(path) for path in input_paths + ) + identity_payload = { + "policy": _INPUT_IDENTITY_POLICY, + "files": [item.to_payload() for item in input_identities], + } + encoded = json.dumps( + {"manifest": canonical, "input_identity": identity_payload}, + sort_keys=True, + separators=(",", ":"), + ).encode() + return ImagePairManifest( + source_path=resolved, + pairs=tuple(pairs), + input_identities=input_identities, + sha256=hashlib.sha256(encoded).hexdigest(), + ) + + +def preflight_image_pair_manifest( + manifest: ImagePairManifest, options: BatchFitOptions +) -> None: + """Check shapes/selectors without importing or initializing CUDA.""" + + manifest.verify_input_identity() + for pair in manifest.pairs: + reference_shape = _probe_array_shape( + pair.reference, pair.reference_hdu, kind="image" + ) + target_shape = _probe_array_shape( + pair.target, pair.target_hdu, kind="image" + ) + _require_matching_shape( + pair.item_id, "target", target_shape, reference_shape + ) + if pair.variance is not None: + variance_shape = _probe_array_shape( + pair.variance, pair.variance_hdu, kind="variance" + ) + _require_matching_shape( + pair.item_id, "variance", variance_shape, reference_shape + ) + if pair.fit_mask is not None: + if pair.fit_mask.suffix.lower() != ".npy": + raise ValueError( + f"pair {pair.item_id!r} fit_mask must be NPY" + ) + fit_shape = _probe_array_shape(pair.fit_mask, None, kind="image") + expected = reference_shape + if options.crop_height is not None: + assert options.crop_width is not None + expected = (options.crop_height, options.crop_width) + _require_matching_shape( + pair.item_id, "fit_mask", fit_shape, expected + ) + if options.auto_stamp_mask: + raise ValueError( + f"pair {pair.item_id!r} supplies fit_mask while " + "auto_stamp_mask is enabled" + ) + supplied_masks = any( + value is not None + for value in ( + pair.reference_mask, + pair.target_mask, + pair.reference_mask_hdu, + pair.target_mask_hdu, + ) + ) + if options.mask_policy == "none" and supplied_masks: + raise ValueError( + f"pair {pair.item_id!r} supplies masks with " + "mask_policy='none'" + ) + if options.mask_policy != "none": + for label, path, hdu in ( + ( + "reference_mask", + pair.reference_mask or pair.reference, + pair.reference_mask_hdu, + ), + ( + "target_mask", + pair.target_mask or pair.target, + pair.target_mask_hdu, + ), + ): + shape = _probe_array_shape(path, hdu, kind="mask") + _require_matching_shape( + pair.item_id, label, shape, reference_shape + ) + if options.crop_y0 is not None: + assert options.crop_x0 is not None + assert options.crop_height is not None + assert options.crop_width is not None + if ( + options.crop_y0 + options.crop_height > reference_shape[0] + or options.crop_x0 + options.crop_width > reference_shape[1] + ): + raise ValueError( + f"pair {pair.item_id!r} crop exceeds image shape " + f"{reference_shape}" + ) + manifest.verify_input_identity() + + +def run_image_pair_item( + item: WorkItem, + item_output_dir: Path, + options: BatchFitOptions, +) -> dict[str, Any]: + """Fit one pair after the caller has established GPU placement.""" + + item_start = time.perf_counter() + from .ois import GaussianBasisComponent + from .workflows import run_constant_kernel_fit + + payload = dict(item.payload) + raw_identities = payload.pop(_INPUT_IDENTITY_KEY, None) + identities = _load_input_identities(raw_identities) + pair = ImagePairSpec.from_payload(payload) + expected_paths = { + path + for field in _PATH_FIELDS + if (path := getattr(pair, field)) is not None + } + if {identity.path for identity in identities} != expected_paths: + raise ValueError( + "work item input identity does not match pair inputs" + ) + _verify_input_identities(identities) + components = [ + GaussianBasisComponent(sigma=sigma, degree=degree) + for sigma, degree in zip(options.basis_sigmas, options.basis_degrees) + ] + workflow_error: Exception | None = None + try: + result = run_constant_kernel_fit( + reference_path=pair.reference, + target_path=pair.target, + output_root=item_output_dir.parent, + name=item_output_dir.name, + reference_hdu=pair.reference_hdu, + target_hdu=pair.target_hdu, + kernel_shape=options.kernel_shape, + components=components, + variance_path=pair.variance, + variance_hdu=pair.variance_hdu, + reference_mask_path=pair.reference_mask, + target_mask_path=pair.target_mask, + reference_mask_hdu=pair.reference_mask_hdu, + target_mask_hdu=pair.target_mask_hdu, + mask_policy=options.mask_policy, + crop_y0=options.crop_y0, + crop_x0=options.crop_x0, + crop_height=options.crop_height, + crop_width=options.crop_width, + fit_mask_path=pair.fit_mask, + auto_stamp_mask=options.auto_stamp_mask, + auto_stamp_size=options.auto_stamp_size, + auto_stamp_count=options.auto_stamp_count, + auto_peak_percentile=options.auto_peak_percentile, + background_degree=options.background_degree, + flux_conserve=options.flux_conserve, + backend=options.backend, + workflow_name="fit_batch_item", + run_prefix="fit-batch-item", + ) + except Exception as exc: + workflow_error = exc + raise + finally: + try: + _verify_input_identities(identities) + except Exception as identity_exc: + if workflow_error is None: + raise + workflow_error.add_note( + "post-item input identity verification failed: " + f"{type(identity_exc).__name__}: {identity_exc}" + ) + timings = dict(result.summary.get("timings_sec", {})) + wall = dict(result.summary.get("wall_sec", {})) + wall["item_runner"] = time.perf_counter() - item_start + return { + "summary_path": str(result.run_dir / "summary.json"), + "run_dir": str(result.run_dir), + "requested_backend": result.summary["requested_backend"], + "backend": result.summary["backend"], + "device": result.summary["device"], + "runtime": result.summary.get("runtime", {}), + "timings_sec": timings, + "wall_sec": wall, + } + + +def _load_input_identities(value: Any) -> tuple[InputFileIdentity, ...]: + if not isinstance(value, list) or not value: + raise ValueError("work item input identity must be a non-empty list") + identities = tuple(InputFileIdentity.from_payload(item) for item in value) + paths = [identity.path for identity in identities] + if len(set(paths)) != len(paths): + raise ValueError("work item input identity paths must be unique") + return identities + + +def _verify_input_identities( + identities: tuple[InputFileIdentity, ...], +) -> None: + for identity in identities: + identity.verify() + + +def _construct_unique_json_mapping( + pairs: list[tuple[str, Any]], +) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, value in pairs: + if key in result: + raise ValueError(f"found duplicate key {key!r}") + result[key] = value + return result + + +def _require_mapping(value: Any, field: str) -> Mapping[str, Any]: + if not isinstance(value, Mapping): + raise ValueError(f"{field} must be a mapping") + if any(not isinstance(key, str) for key in value): + raise ValueError(f"{field} keys must be strings") + return value + + +def _reject_unknown_fields( + payload: Mapping[str, Any], allowed: frozenset[str], field: str +) -> None: + unknown = sorted(set(payload) - allowed) + if unknown: + raise ValueError( + f"{field} has unknown field(s): " + ", ".join(unknown) + ) + + +def _resolve_input_path( + value: Any, + *, + base_dir: Path, + field: str, + required: bool, +) -> Path | None: + if value is None and not required: + return None + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{field} must be a non-empty path") + path = Path(value.strip()).expanduser() + if not path.is_absolute(): + path = base_dir / path + resolved = path.resolve() + if not resolved.is_file(): + raise FileNotFoundError(f"{field} does not exist: {resolved}") + if resolved.suffix.lower() not in FITS_SUFFIXES | {".npy"}: + raise ValueError(f"{field} must name a FITS or NPY image: {resolved}") + return resolved + + +def _validate_hdu(value: Any, *, field: str) -> int | None: + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError(f"{field} must be a non-negative integer or null") + return value + + +def _require_non_negative_integer(value: Any, *, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError(f"{field} must be a non-negative integer") + return value + + +def _validate_hdu_dependencies( + values: Mapping[str, Any], *, index: int +) -> None: + if values["variance_hdu"] is not None and values["variance"] is None: + raise ValueError(f"pairs[{index}].variance_hdu requires variance") + for hdu_field, path_field in ( + ("reference_hdu", "reference"), + ("target_hdu", "target"), + ("variance_hdu", "variance"), + ): + path = values[path_field] + if ( + values[hdu_field] is not None + and path is not None + and path.suffix.lower() == ".npy" + ): + raise ValueError( + f"pairs[{index}].{hdu_field} requires a FITS input" + ) + for hdu_field, path_field, fallback_field in ( + ("reference_mask_hdu", "reference_mask", "reference"), + ("target_mask_hdu", "target_mask", "target"), + ): + path = values[path_field] or values[fallback_field] + if values[hdu_field] is not None and path.suffix.lower() == ".npy": + raise ValueError( + f"pairs[{index}].{hdu_field} requires a FITS input" + ) + + +def _probe_array_shape( + path: Path, hdu: int | None, *, kind: str +) -> tuple[int, int]: + if path.suffix.lower() == ".npy": + if hdu is not None: + raise ValueError(f"HDU selectors are not supported for {path}") + array = np.load(path, mmap_mode="r", allow_pickle=False) + if array.ndim != 2: + raise ValueError(f"{path} is not a 2D image") + return (int(array.shape[0]), int(array.shape[1])) + + with fits.open(path, memmap=True, lazy_load_hdus=False) as hdul: + candidates = [ + (index, _fits_image_shape(item)) + for index, item in enumerate(hdul) + ] + candidates = [ + (index, shape) for index, shape in candidates if shape is not None + ] + if hdu is not None: + if hdu >= len(hdul): + raise ValueError(f"HDU {hdu} is out of range for {path}") + shape = _fits_image_shape(hdul[hdu]) + if shape is None: + raise ValueError(f"HDU {hdu} in {path} is not a 2D image") + return shape + if kind == "image": + if candidates: + return candidates[0][1] + raise ValueError(f"could not find a 2D image HDU in {path}") + + expected_names = ( + VARIANCE_EXTENSION_NAMES + if kind == "variance" + else MASK_EXTENSION_NAMES + ) + named = [] + errors = [] + for index, shape in candidates: + extname = ( + str(hdul[index].header.get("EXTNAME", "")).strip().upper() + ) + if extname in expected_names: + named.append((index, shape)) + if extname in ERROR_EXTENSION_NAMES: + errors.append(index) + if len(named) == 1: + return named[0][1] + if len(named) > 1: + raise ValueError( + f"multiple {kind}-like HDUs in {path}; specify one" + ) + if kind == "variance" and len(candidates) == 1: + if candidates[0][0] in errors: + raise ValueError( + f"variance FITS {path} uses an error/sigma extension" + ) + return candidates[0][1] + raise ValueError(f"{kind} FITS {path} is ambiguous; specify an HDU") + + +def _fits_image_shape(hdu: Any) -> tuple[int, int] | None: + if not isinstance( + hdu, + (fits.PrimaryHDU, fits.ImageHDU, fits.CompImageHDU), + ): + return None + header = hdu.header + if int(header.get("NAXIS", 0)) != 2: + return None + height = int(header.get("NAXIS2", 0)) + width = int(header.get("NAXIS1", 0)) + if height <= 0 or width <= 0: + return None + return (height, width) + + +def _require_matching_shape( + item_id: str, + label: str, + actual: tuple[int, int], + expected: tuple[int, int], +) -> None: + if actual != expected: + raise ValueError( + f"pair {item_id!r} {label} shape {actual} does not match " + f"expected shape {expected}" + ) + + +__all__ = [ + "BatchFitOptions", + "IMAGE_PAIR_MANIFEST_SCHEMA", + "ImagePairManifest", + "ImagePairSpec", + "InputFileIdentity", + "load_image_pair_manifest", + "preflight_image_pair_manifest", + "run_image_pair_item", +] diff --git a/src/cuphoton/xpois/commands.py b/src/cuphoton/xpois/commands.py index 6a18412..8309eea 100644 --- a/src/cuphoton/xpois/commands.py +++ b/src/cuphoton/xpois/commands.py @@ -23,7 +23,9 @@ StringInvariant, ) +from .batch import BatchFitOptions from .data import inspect_hsc_data_tree +from .dragon import run_dragon_image_pair_batch from .ois import ( EXPLICIT_BACKENDS, SUPPORTED_BACKENDS, @@ -121,19 +123,8 @@ def run(self) -> None: self._emit_json(summary) -class _KernelSolveCommand(XPOISCommand): - reference = None - target = None - reference_hdu = None - target_hdu = None - variance = None - variance_hdu = None - reference_mask = None - target_mask = None - reference_mask_hdu = None - target_mask_hdu = None +class _KernelSolveOptionsCommand(XPOISCommand): mask_policy = None - fit_mask = None auto_stamp_mask = None auto_stamp_size = None auto_stamp_count = None @@ -151,71 +142,6 @@ class _KernelSolveCommand(XPOISCommand): background_degree = None flux_conserve = None - class ReferenceArg(ImagePathInvariant): - _arg = "--reference" - _help = "Reference image path (.fits or .npy)." - _mandatory = True - - class TargetArg(ImagePathInvariant): - _arg = "--target" - _help = "Target image path (.fits or .npy)." - _mandatory = True - - class ReferenceHduArg(NonNegativeIntegerInvariant): - _arg = "--reference-hdu" - _help = "Optional explicit HDU index for the reference FITS file." - _mandatory = False - _default = None - - class TargetHduArg(NonNegativeIntegerInvariant): - _arg = "--target-hdu" - _help = "Optional explicit HDU index for the target FITS file." - _mandatory = False - _default = None - - class VarianceArg(PathSpecInvariant): - _arg = "--variance" - _help = "Optional target variance image path (.fits or .npy)." - _mandatory = False - _default = None - - class VarianceHduArg(NonNegativeIntegerInvariant): - _arg = "--variance-hdu" - _help = "Optional explicit HDU index for the variance FITS file." - _mandatory = False - _default = None - - class ReferenceMaskArg(PathSpecInvariant): - _arg = "--reference-mask" - _help = ( - "Optional reference mask path (.fits or .npy). " - "Defaults to --reference when mask-policy is enabled on FITS " - "input." - ) - _mandatory = False - _default = None - - class TargetMaskArg(PathSpecInvariant): - _arg = "--target-mask" - _help = ( - "Optional target mask path (.fits or .npy). " - "Defaults to --target when mask-policy is enabled on FITS input." - ) - _mandatory = False - _default = None - - class ReferenceMaskHduArg(NonNegativeIntegerInvariant): - _arg = "--reference-mask-hdu" - _help = "Optional explicit HDU index for the reference FITS mask." - _mandatory = False - _default = None - - class TargetMaskHduArg(NonNegativeIntegerInvariant): - _arg = "--target-mask-hdu" - _help = "Optional explicit HDU index for the target FITS mask." - _mandatory = False - _default = None - class MaskPolicyArg(SetInvariant): _arg = "--mask-policy" _help = ( @@ -226,12 +152,6 @@ class MaskPolicyArg(SetInvariant): _default = "none" _set = {"none", "strict", "hsc-masklite", "masklite"} - class FitMaskArg(PathSpecInvariant): - _arg = "--fit-mask" - _help = "Optional boolean .npy mask selecting fit pixels." - _mandatory = False - _default = None - class AutoStampMaskArg(BoolInvariant): _arg = "--auto-stamp-mask" _help = "Auto-build a compact-source stamp fit mask." @@ -380,6 +300,91 @@ def _output_root(self) -> Path: return self.context.runs_dir +class _KernelSolveCommand(_KernelSolveOptionsCommand): + reference = None + target = None + reference_hdu = None + target_hdu = None + variance = None + variance_hdu = None + reference_mask = None + target_mask = None + reference_mask_hdu = None + target_mask_hdu = None + fit_mask = None + + class ReferenceArg(ImagePathInvariant): + _arg = "--reference" + _help = "Reference image path (.fits or .npy)." + _mandatory = True + + class TargetArg(ImagePathInvariant): + _arg = "--target" + _help = "Target image path (.fits or .npy)." + _mandatory = True + + class ReferenceHduArg(NonNegativeIntegerInvariant): + _arg = "--reference-hdu" + _help = "Optional explicit HDU index for the reference FITS file." + _mandatory = False + _default = None + + class TargetHduArg(NonNegativeIntegerInvariant): + _arg = "--target-hdu" + _help = "Optional explicit HDU index for the target FITS file." + _mandatory = False + _default = None + + class VarianceArg(PathSpecInvariant): + _arg = "--variance" + _help = "Optional target variance image path (.fits or .npy)." + _mandatory = False + _default = None + + class VarianceHduArg(NonNegativeIntegerInvariant): + _arg = "--variance-hdu" + _help = "Optional explicit HDU index for the variance FITS file." + _mandatory = False + _default = None + + class ReferenceMaskArg(PathSpecInvariant): + _arg = "--reference-mask" + _help = ( + "Optional reference mask path (.fits or .npy). " + "Defaults to --reference when mask-policy is enabled on FITS " + "input." + ) + _mandatory = False + _default = None + + class TargetMaskArg(PathSpecInvariant): + _arg = "--target-mask" + _help = ( + "Optional target mask path (.fits or .npy). " + "Defaults to --target when mask-policy is enabled on FITS input." + ) + _mandatory = False + _default = None + + class ReferenceMaskHduArg(NonNegativeIntegerInvariant): + _arg = "--reference-mask-hdu" + _help = "Optional explicit HDU index for the reference FITS mask." + _mandatory = False + _default = None + + class TargetMaskHduArg(NonNegativeIntegerInvariant): + _arg = "--target-mask-hdu" + _help = "Optional explicit HDU index for the target FITS mask." + _mandatory = False + _default = None + + class FitMaskArg(PathSpecInvariant): + _arg = "--fit-mask" + _help = "Optional boolean .npy mask selecting fit pixels." + _mandatory = False + _default = None + + class _FitCommand(_KernelSolveCommand): backend = None @@ -490,6 +495,90 @@ def run(self) -> None: self._emit_json(result.summary) +class FitBatchDragonCommand(_KernelSolveOptionsCommand): + """Fit an image-pair manifest with explicitly placed Dragon workers.""" + + manifest = None + max_workers = None + result_timeout_sec = None + worker_timeout_sec = None + backend = None + + class ManifestArg(PathSpecInvariant): + _arg = "--manifest" + _help = "Strict JSON or YAML image-pair manifest." + _mandatory = True + + class MaxWorkersArg(PositiveIntegerInvariant): + _arg = "--max-workers" + _help = "Maximum explicitly placed GPU workers." + _mandatory = False + _default = None + + class ResultTimeoutSecArg(FloatInvariant): + _arg = "--result-timeout-sec" + _help = ( + "Seconds of grace for ProcessGroup join and compact worker " + "results. [default: %default]" + ) + _mandatory = False + _default = 60.0 + _min = 0.001 + + class WorkerTimeoutSecArg(FloatInvariant): + _arg = "--worker-timeout-sec" + _help = ( + "Dragon ProcessGroup worker wall time in seconds. " + "[default: %default]" + ) + _mandatory = False + _default = 3600.0 + _min = 0.001 + + class BackendArg(SetInvariant): + _arg = "--backend" + _help = "GPU fit backend. [default: %default]" + _mandatory = False + _default = "cupy" + _set = {"cupy", "numba-cuda", "cutile"} + + def run(self) -> None: + components = self._components() + options = self._call( + BatchFitOptions, + kernel_shape=(self.kernel_height, self.kernel_width), + basis_sigmas=tuple(item.sigma for item in components), + basis_degrees=tuple(item.degree for item in components), + mask_policy=self.mask_policy, + crop_y0=self.crop_y0, + crop_x0=self.crop_x0, + crop_height=self.crop_height, + crop_width=self.crop_width, + auto_stamp_mask=bool(self.auto_stamp_mask), + auto_stamp_size=self.auto_stamp_size, + auto_stamp_count=self.auto_stamp_count, + auto_peak_percentile=self.auto_peak_percentile, + background_degree=self.background_degree, + flux_conserve=bool(self.flux_conserve), + backend=self.backend, + ) + result = self._call( + run_dragon_image_pair_batch, + manifest_path=Path(self.manifest).expanduser(), + output_root=self._output_root(), + run_id=self.run_name or None, + max_workers=self.max_workers, + result_timeout_sec=self.result_timeout_sec, + worker_timeout_sec=self.worker_timeout_sec, + options=options, + ) + self._emit_json(result.to_dict()) + if result.status != "success": + raise CommandError( + "Dragon batch failed; inspect " + str(result.summary_path) + ) + + class BenchmarkBackendsCommand(_KernelSolveCommand): """Benchmark constant-kernel CPU/CuPy backends and parity. diff --git a/src/cuphoton/xpois/dragon.py b/src/cuphoton/xpois/dragon.py new file mode 100644 index 0000000..6c70156 --- /dev/null +++ b/src/cuphoton/xpois/dragon.py @@ -0,0 +1,1327 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Optional Dragon ProcessGroup adapter for XPOIS image pairs.""" + +from __future__ import annotations + +import hashlib +import math +import os +import queue +import socket +import sys +import time +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path +from typing import Any + +from cuphoton.core.bulk import ( + Placement, + WorkItem, + atomic_write_json, + audit_terminal_records, + error_payload, + json_mapping, + new_run_id, + partition_byte_balanced, + read_json_mapping, + timestamp_utc, + validate_identifier, +) + +from .batch import ( + BatchFitOptions, + load_image_pair_manifest, + preflight_image_pair_manifest, + run_image_pair_item, +) + +_DRAGON_GPU_BACKENDS = frozenset({"cupy", "numba-cuda", "cutile"}) + + +@dataclass(frozen=True) +class DragonBatchResult: + """Handle for a terminal XPOIS Dragon batch.""" + + run_id: str + run_dir: Path + summary_path: Path + status: str + summary: Mapping[str, Any] + + def to_dict(self) -> dict[str, Any]: + """Return a compact command result.""" + + return { + "run_id": self.run_id, + "run_dir": str(self.run_dir), + "summary_path": str(self.summary_path), + "status": self.status, + } + + +@dataclass(frozen=True) +class _DragonAPI: + System: Any + Node: Any + Policy: Any + ProcessGroup: Any + ProcessTemplate: Any + Queue: Any + + +def run_dragon_image_pair_batch( + *, + manifest_path: Path, + output_root: Path, + run_id: str | None, + max_workers: int | None, + result_timeout_sec: float, + options: BatchFitOptions, + worker_timeout_sec: float = 3600.0, +) -> DragonBatchResult: + """Run deterministic XPOIS shards with one Dragon worker per GPU.""" + + invocation_start = time.perf_counter() + started_at = timestamp_utc() + coordinator_timings: dict[str, float] = {} + if ( + isinstance(result_timeout_sec, bool) + or not isinstance(result_timeout_sec, (int, float)) + or not math.isfinite(result_timeout_sec) + or result_timeout_sec <= 0 + ): + raise ValueError("result_timeout_sec must be positive") + if max_workers is not None and ( + isinstance(max_workers, bool) + or not isinstance(max_workers, int) + or max_workers <= 0 + ): + raise ValueError("max_workers must be a positive integer") + if ( + isinstance(worker_timeout_sec, bool) + or not isinstance(worker_timeout_sec, (int, float)) + or not math.isfinite(worker_timeout_sec) + or worker_timeout_sec <= 0 + ): + raise ValueError("worker_timeout_sec must be positive") + if options.backend not in _DRAGON_GPU_BACKENDS: + raise ValueError( + "Dragon XPOIS workers require an explicit GPU backend" + ) + + phase_start = time.perf_counter() + manifest = load_image_pair_manifest(manifest_path) + coordinator_timings["manifest_load_sec"] = ( + time.perf_counter() - phase_start + ) + phase_start = time.perf_counter() + preflight_image_pair_manifest(manifest, options) + coordinator_timings["manifest_preflight_sec"] = ( + time.perf_counter() - phase_start + ) + phase_start = time.perf_counter() + api = _load_dragon_api() + system = api.System() + node_ids = tuple(system.nodes) + allocation_node_count = len(node_ids) + placements = discover_gpu_placements( + lambda: system, + api.Node, + node_ids=node_ids, + ) + coordinator_timings["dragon_discovery_sec"] = ( + time.perf_counter() - phase_start + ) + phase_start = time.perf_counter() + requested_workers = max_workers or len(placements) + worker_count = min( + requested_workers, len(placements), len(manifest.pairs) + ) + if worker_count <= 0: + raise RuntimeError("Dragon allocation exposes no usable GPUs") + selected = _select_gpu_placements(placements, worker_count) + distinct_host_count = len({placement.host for placement in selected}) + items = manifest.work_items() + shards = partition_byte_balanced(items, worker_count) + coordinator_timings["partition_sec"] = time.perf_counter() - phase_start + + phase_start = time.perf_counter() + join_timeout_sec = float(worker_timeout_sec + result_timeout_sec) + effective_run_id = run_id or new_run_id("dragon-xpois") + validate_identifier(effective_run_id, field="run_id") + run_dir = output_root.expanduser().resolve() / effective_run_id + run_dir.mkdir(parents=True, exist_ok=False) + for name in ("items", "records", "workers"): + (run_dir / name).mkdir() + atomic_write_json(run_dir / "manifest.json", manifest.canonical_payload()) + atomic_write_json( + run_dir / "input-identity.json", + manifest.input_identity_payload(), + ) + dragon_version = _distribution_version("dragonhpc") + atomic_write_json( + run_dir / "run.json", + { + "schema": "cuphoton.xpois.dragon-run/v1", + "record_type": "immutable-launch", + "run_id": effective_run_id, + "started_at_utc": started_at, + "manifest_sha256": manifest.sha256, + "dragonhpc_version": dragon_version, + "worker_count": worker_count, + "allocation_node_count": allocation_node_count, + "distinct_host_count": distinct_host_count, + "worker_timeout_sec": worker_timeout_sec, + "join_timeout_sec": join_timeout_sec, + "result_timeout_sec": result_timeout_sec, + "placements": [placement.to_dict() for placement in selected], + "options": options.to_payload(), + }, + ) + coordinator_timings["run_artifact_setup_sec"] = ( + time.perf_counter() - phase_start + ) + + lifecycle_errors: list[dict[str, str]] = [] + exit_status: list[dict[str, int]] = [] + shard_results: list[dict[str, Any]] = [] + queue_result_count = 0 + results_queue: Any | None = None + group: Any | None = None + start_attempted = False + join_completed_cleanly = False + lifecycle_phase = "queue_create" + process_setup_start = time.perf_counter() + try: + results_queue = api.Queue(maxsize=worker_count) + lifecycle_phase = "group_create" + group = api.ProcessGroup( + restart=False, + ignore_error_on_exit=False, + walltime=worker_timeout_sec, + ) + lifecycle_phase = "process_setup" + for placement, shard in zip(selected, shards): + policy = api.Policy( + placement=api.Policy.Placement.HOST_NAME, + host_name=placement.host, + gpu_affinity=[placement.gpu_id], + ) + template = api.ProcessTemplate( + target=_dragon_shard_worker, + args=( + effective_run_id, + str(run_dir), + placement.to_dict(), + [item.to_dict() for item in shard], + options.to_payload(), + results_queue, + allocation_node_count == 1, + ), + policy=policy, + ) + group.add_process(nproc=1, template=template) + + lifecycle_phase = "init" + group.init() + coordinator_timings["dragon_process_setup_sec"] = ( + time.perf_counter() - process_setup_start + ) + lifecycle_phase = "start" + phase_start = time.perf_counter() + start_attempted = True + try: + group.start() + finally: + coordinator_timings["dragon_launch_sec"] = ( + time.perf_counter() - phase_start + ) + lifecycle_phase = "join" + phase_start = time.perf_counter() + try: + group.join(timeout=join_timeout_sec) + join_completed_cleanly = True + except Exception as exc: + phase = ( + "join_timeout" if isinstance(exc, TimeoutError) else "join" + ) + lifecycle_errors.append({"phase": phase, **error_payload(exc)}) + try: + coordinator_timings["worker_join_sec"] = ( + time.perf_counter() - phase_start + ) + result_collection_start = time.perf_counter() + lifecycle_phase = "exit_status" + try: + # ``exit_status`` is exception-decorated. ``inactive_puids`` + # is the same public state without re-raising a worker + # exception, so it remains inspectable after a failed join. + exit_status = [ + {"puid": int(puid), "exit_code": int(code)} + for puid, code in group.inactive_puids + ] + except Exception as exc: + lifecycle_errors.append( + {"phase": "exit_status", **error_payload(exc)} + ) + deadline = time.monotonic() + result_timeout_sec + lifecycle_phase = "queue_result" + while queue_result_count < worker_count: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + try: + result = results_queue.get(timeout=remaining) + except (queue.Empty, TimeoutError): + break + queue_result_count += 1 + if isinstance(result, Mapping): + try: + shard_results.append( + json_mapping(result, field="Dragon shard result") + ) + except (TypeError, ValueError) as exc: + lifecycle_errors.append( + { + "phase": "queue_result", + "type": "InvalidShardResult", + "message": str(exc), + } + ) + else: + lifecycle_errors.append( + { + "phase": "queue_result", + "type": "InvalidShardResult", + "message": "worker result was not a mapping", + } + ) + coordinator_timings["result_collection_sec"] = ( + time.perf_counter() - result_collection_start + ) + finally: + coordinator_timings.setdefault( + "worker_join_sec", time.perf_counter() - phase_start + ) + except Exception as exc: + lifecycle_errors.append( + {"phase": lifecycle_phase, **error_payload(exc)} + ) + finally: + coordinator_timings.setdefault( + "dragon_process_setup_sec", + time.perf_counter() - process_setup_start, + ) + cleanup_start = time.perf_counter() + if group is not None: + if start_attempted and not join_completed_cleanly: + try: + group.stop(patience=5.0) + except Exception as exc: + lifecycle_errors.append( + {"phase": "stop_after_failure", **error_payload(exc)} + ) + try: + group.close(patience=5.0) + except Exception as exc: + lifecycle_errors.append( + {"phase": "close", **error_payload(exc)} + ) + cleanup = getattr(group, "_close_no_decorator", None) + if cleanup is not None: + try: + cleanup(patience=5.0) + except Exception as cleanup_exc: + lifecycle_errors.append( + { + "phase": "forced_close", + **error_payload(cleanup_exc), + } + ) + if results_queue is not None: + try: + results_queue.close() + except Exception as exc: + lifecycle_errors.append( + {"phase": "queue_close", **error_payload(exc)} + ) + coordinator_timings["dragon_cleanup_sec"] = ( + time.perf_counter() - cleanup_start + ) + + audit_start = time.perf_counter() + records, record_errors = _load_terminal_records(run_dir) + audit = audit_terminal_records([item.item_id for item in items], records) + expected_records = { + item.item_id: { + "worker_id": placement.worker_id, + "weight_bytes": item.weight_bytes, + "backend": options.backend, + } + for placement, shard in zip(selected, shards) + for item in shard + } + record_errors.extend( + _audit_terminal_record_contract( + run_id=effective_run_id, + expected=expected_records, + records=records, + ) + ) + shard_audit = _audit_shard_results( + shards, + shard_results, + records, + placements=selected, + allow_loopback_alias=allocation_node_count == 1, + backend=options.backend, + ) + item_timings, timing_errors = _aggregate_item_timings(records) + record_errors.extend(timing_errors) + process_audit = { + "ok": len(exit_status) == worker_count + and all(item["exit_code"] == 0 for item in exit_status), + "expected_count": worker_count, + "observed_count": len(exit_status), + "nonzero": [item for item in exit_status if item["exit_code"] != 0], + } + status = ( + "success" + if audit["ok"] + and not audit["failed_item_ids"] + and shard_audit["ok"] + and process_audit["ok"] + and not lifecycle_errors + and not record_errors + else "failed" + ) + coordinator_timings["terminal_audit_sec"] = ( + time.perf_counter() - audit_start + ) + setup_phases = ( + "manifest_load_sec", + "manifest_preflight_sec", + "dragon_discovery_sec", + "partition_sec", + "run_artifact_setup_sec", + "dragon_process_setup_sec", + "dragon_launch_sec", + ) + coordinator_setup_sec = sum( + coordinator_timings.get(name, 0.0) for name in setup_phases + ) + coordinator_wall_sec = time.perf_counter() - invocation_start + summary = { + "schema": "cuphoton.xpois.dragon-summary/v1", + "status": status, + "run_id": effective_run_id, + "started_at_utc": started_at, + "completed_at_utc": timestamp_utc(), + "coordinator_wall_sec": coordinator_wall_sec, + "coordinator_wall_definition": ( + "Function entry through worker cleanup and terminal audits; " + "the final summary.json atomic commit is excluded." + ), + "coordinator_timings_sec": coordinator_timings, + "coordinator_totals_sec": {"setup": coordinator_setup_sec}, + "manifest_sha256": manifest.sha256, + "dragonhpc_version": dragon_version, + "worker_timeout_sec": worker_timeout_sec, + "join_timeout_sec": join_timeout_sec, + "result_timeout_sec": result_timeout_sec, + "allocation_node_count": allocation_node_count, + "distinct_host_count": distinct_host_count, + "options": options.to_payload(), + "placements": [placement.to_dict() for placement in selected], + "shards": [ + { + "worker_id": placement.worker_id, + "item_count": len(shard), + "weight_bytes": sum(item.weight_bytes for item in shard), + "item_ids_sha256": _item_ids_sha256(shard), + } + for placement, shard in zip(selected, shards) + ], + "shard_results": sorted(shard_results, key=_shard_result_sort_key), + "queue_result_count": queue_result_count, + "shard_result_audit": shard_audit, + "terminal_record_audit": audit, + "terminal_record_errors": record_errors, + "process_exit_status": exit_status, + "process_exit_audit": process_audit, + "lifecycle_errors": lifecycle_errors, + "timings_sec": item_timings, + } + summary_path = run_dir / "summary.json" + atomic_write_json(summary_path, summary) + return DragonBatchResult( + run_id=effective_run_id, + run_dir=run_dir, + summary_path=summary_path, + status=status, + summary=summary, + ) + + +def discover_gpu_placements( + system_type: Callable[[], Any], + node_type: Callable[[Any], Any], + *, + node_ids: Sequence[Any] | None = None, +) -> tuple[Placement, ...]: + """Enumerate actual Dragon Node.gpus IDs in allocation order.""" + + if node_ids is None: + node_ids = tuple(system_type().nodes) + placements: list[Placement] = [] + seen: set[tuple[str, int]] = set() + for node_id in node_ids: + node = node_type(node_id) + host = str(node.hostname) + for gpu_id in node.gpus or []: + key = (host, gpu_id) + if key in seen: + raise RuntimeError( + f"duplicate Dragon GPU placement: {host}:{gpu_id}" + ) + seen.add(key) + placements.append( + Placement( + worker_id=len(placements), + host=host, + gpu_id=gpu_id, + ) + ) + return tuple(placements) + + +def _select_gpu_placements( + placements: Sequence[Placement], worker_count: int +) -> tuple[Placement, ...]: + """Select GPUs round-robin across hosts, then assign dense worker IDs.""" + + if isinstance(worker_count, bool) or not isinstance(worker_count, int): + raise TypeError("worker_count must be an integer") + if worker_count <= 0 or worker_count > len(placements): + raise ValueError("worker_count exceeds available Dragon placements") + by_host: dict[str, list[Placement]] = {} + for placement in placements: + by_host.setdefault(placement.host, []).append(placement) + selected: list[Placement] = [] + offset = 0 + while len(selected) < worker_count: + progressed = False + for host_placements in by_host.values(): + if offset >= len(host_placements): + continue + selected.append(host_placements[offset]) + progressed = True + if len(selected) == worker_count: + break + if not progressed: + raise RuntimeError("could not select requested Dragon placements") + offset += 1 + return tuple( + Placement( + worker_id=worker_id, + host=placement.host, + gpu_id=placement.gpu_id, + ) + for worker_id, placement in enumerate(selected) + ) + + +def _dragon_shard_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: validate placement, run one shard, report compactly.""" + + placement = Placement(**dict(placement_payload)) + items = tuple(WorkItem.from_dict(payload) for payload in item_payloads) + options = BatchFitOptions.from_payload(options_payload) + result = _execute_shard( + run_id=run_id, + run_dir=Path(run_dir_raw), + placement=placement, + items=items, + options=options, + item_runner=run_image_pair_item, + gpu_identity_loader=_collect_gpu_identity, + allow_loopback_alias=allow_loopback_alias, + ) + results_queue.put(result) + + +def _execute_shard( + *, + run_id: str, + run_dir: Path, + placement: Placement, + items: Sequence[WorkItem], + options: BatchFitOptions, + item_runner: Callable[ + [WorkItem, Path, BatchFitOptions], Mapping[str, Any] + ], + gpu_identity_loader: Callable[[str], Mapping[str, Any]], + require_clean_cuda_imports: bool = True, + allow_loopback_alias: bool = False, +) -> dict[str, Any]: + """Execute one shard with injectable CUDA-free test collaborators.""" + + worker_start = time.perf_counter() + started_at = timestamp_utc() + actual_host = socket.gethostname() + if not _hostnames_match( + placement.host, + actual_host, + allow_loopback_alias=allow_loopback_alias, + ): + raise RuntimeError( + "Dragon worker placement mismatch: requested " + f"{placement.host!r}, running on {actual_host!r}" + ) + visibility = _singleton_cuda_visibility(placement.gpu_id) + premature = sorted( + name + for name in ("cupy", "numba.cuda", "cuda.tile") + if name in sys.modules + ) + if require_clean_cuda_imports and premature: + raise RuntimeError( + "CUDA modules were imported before Dragon worker placement: " + + ", ".join(premature) + ) + gpu_identity = dict(gpu_identity_loader(options.backend)) + provenance = { + "worker_id": placement.worker_id, + "requested_host": placement.host, + "requested_gpu_id": placement.gpu_id, + "hostname": actual_host, + "pid": os.getpid(), + "cuda_visible_devices": visibility, + "gpu": gpu_identity, + } + success_count = 0 + failed_count = 0 + record_write_errors: list[dict[str, str]] = [] + worker_timings: dict[str, float] = {} + for item in items: + item_start = time.perf_counter() + record: dict[str, Any] = { + "schema": "cuphoton.xpois.dragon-item/v1", + "run_id": run_id, + "item_id": item.item_id, + "worker_id": placement.worker_id, + "weight_bytes": item.weight_bytes, + "started_at_utc": timestamp_utc(), + } + try: + metadata = dict( + item_runner(item, run_dir / "items" / item.item_id, options) + ) + for field in ("run_dir", "summary_path"): + if field in metadata: + metadata[field] = str( + Path(metadata[field]).relative_to(run_dir) + ) + if "timings_sec" in metadata: + metadata["timings_sec"] = _validated_timings( + metadata["timings_sec"], + field=f"item {item.item_id!r} timings_sec", + ) + if "wall_sec" in metadata: + metadata["wall_sec"] = _validated_timings( + metadata["wall_sec"], + field=f"item {item.item_id!r} wall_sec", + ) + record.update(metadata) + record["status"] = "success" + except Exception as exc: + record["status"] = "failed" + record["error"] = error_payload(exc) + notes = getattr(exc, "__notes__", ()) + if notes: + record["error"]["notes"] = "\n".join(map(str, notes)) + record["worker_seconds"] = time.perf_counter() - item_start + record["completed_at_utc"] = timestamp_utc() + for phase, value in record.get("timings_sec", {}).items(): + if isinstance(value, (int, float)) and not isinstance( + value, bool + ): + worker_timings[str(phase)] = worker_timings.get( + str(phase), 0.0 + ) + float(value) + try: + atomic_write_json( + run_dir / "records" / f"{item.item_id}.json", record + ) + except Exception as exc: + record_write_errors.append( + {"item_id": item.item_id, **error_payload(exc)} + ) + failed_count += 1 + else: + if record["status"] == "success": + success_count += 1 + else: + failed_count += 1 + + result = { + "schema": "cuphoton.xpois.dragon-shard/v1", + "worker_id": placement.worker_id, + "status": "success" if failed_count == 0 else "failed", + "item_count": len(items), + "success_count": success_count, + "failed_count": failed_count, + "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_start, + "timings_sec": worker_timings, + "provenance": provenance, + "record_write_errors": record_write_errors, + } + atomic_write_json( + run_dir / "workers" / f"worker-{placement.worker_id:04d}.json", + result, + ) + return result + + +def _singleton_cuda_visibility(expected_gpu_id: int | None = None) -> str: + raw = os.environ.get("CUDA_VISIBLE_DEVICES") + if raw is None: + raise RuntimeError("Dragon worker has no CUDA_VISIBLE_DEVICES") + tokens = [token.strip() for token in raw.split(",") if token.strip()] + if len(tokens) != 1 or tokens[0] == "-1": + raise RuntimeError( + f"Dragon worker must see exactly one CUDA device, got {raw!r}" + ) + if expected_gpu_id is not None and tokens[0] != str(expected_gpu_id): + raise RuntimeError( + "Dragon worker GPU placement mismatch: requested " + f"{expected_gpu_id}, got CUDA_VISIBLE_DEVICES={raw!r}" + ) + return tokens[0] + + +def _hostnames_match( + requested: str, + actual: str, + *, + allow_loopback_alias: bool = False, +) -> bool: + """Accept exact or equivalent short/FQDN scheduler hostnames.""" + + requested_normalized = requested.rstrip(".").lower() + actual_normalized = actual.rstrip(".").lower() + if allow_loopback_alias and requested_normalized in { + "localhost", + "localhost.localdomain", + }: + # Dragon 0.14.1 reports ``localhost`` for its single-node system + # descriptor even though socket.gethostname() exposes the machine + # hostname inside the launched worker. + return bool(actual_normalized) + if requested_normalized == actual_normalized: + return True + if "." not in requested_normalized: + return actual_normalized.startswith(requested_normalized + ".") + if "." not in actual_normalized: + return requested_normalized.startswith(actual_normalized + ".") + return False + + +def _collect_gpu_identity(backend: str) -> dict[str, Any]: + """Initialize the selected backend only after singleton placement.""" + + if backend in {"auto", "cupy", "cutile"}: + try: + return _collect_cupy_identity() + except (ImportError, OSError, RuntimeError): + if backend != "auto": + raise + return _collect_numba_identity() + + +def _collect_cupy_identity() -> dict[str, Any]: + import cupy as cp + + device_index = int(cp.cuda.runtime.getDevice()) + properties = cp.cuda.runtime.getDeviceProperties(device_index) + name = properties.get("name", "unknown") + if isinstance(name, bytes): + name = name.decode(errors="replace") + uuid_value = properties.get("uuid") + if isinstance(uuid_value, bytes): + uuid_value = uuid_value.hex() + elif uuid_value is not None and not isinstance(uuid_value, str): + uuid_value = str(uuid_value) + pci_bus_id = None + identity_error = None + try: + pci_bus_id = str(cp.cuda.runtime.deviceGetPCIBusId(device_index)) + except Exception as exc: # pragma: no cover - runtime-specific + identity_error = f"{type(exc).__name__}: {exc}" + return { + "backend": "cupy", + "device_index": device_index, + "name": str(name), + "uuid": uuid_value, + "pci_bus_id": pci_bus_id, + "identity_error": identity_error, + } + + +def _collect_numba_identity() -> dict[str, Any]: + from numba import cuda + + device = cuda.get_current_device() + name = device.name + if isinstance(name, bytes): + name = name.decode(errors="replace") + uuid_value = getattr(device, "uuid", None) + pci_bus_id = getattr(device, "pci_bus_id", None) + return { + "backend": "numba-cuda", + "device_index": int(getattr(device, "id", 0)), + "name": str(name), + "uuid": str(uuid_value) if uuid_value is not None else None, + "pci_bus_id": str(pci_bus_id) if pci_bus_id is not None else None, + "identity_error": None, + } + + +def _load_terminal_records( + run_dir: Path, +) -> tuple[list[dict[str, Any]], list[dict[str, str]]]: + records: list[dict[str, Any]] = [] + errors: list[dict[str, str]] = [] + for path in sorted((run_dir / "records").glob("*.json")): + try: + records.append(read_json_mapping(path)) + except Exception as exc: + errors.append( + { + "record_path": str(path.relative_to(run_dir)), + **error_payload(exc), + } + ) + return records, errors + + +def _audit_terminal_record_contract( + *, + run_id: str, + expected: Mapping[str, Mapping[str, Any]], + records: Sequence[Mapping[str, Any]], +) -> list[dict[str, str]]: + """Validate durable Dragon item records beyond identity/status counts.""" + + errors: list[dict[str, str]] = [] + for index, record in enumerate(records): + item_id_value = record.get("item_id") + item_id = item_id_value if isinstance(item_id_value, str) else "" + problems: list[str] = [] + if record.get("schema") != "cuphoton.xpois.dragon-item/v1": + problems.append("schema") + if record.get("run_id") != run_id: + problems.append("run_id") + assignment = expected.get(item_id) + if assignment is None: + problems.append("item_id") + else: + worker_id = _strict_integer(record.get("worker_id")) + if worker_id != assignment["worker_id"]: + problems.append("worker_id") + weight_bytes = _strict_integer(record.get("weight_bytes")) + if weight_bytes != assignment["weight_bytes"]: + problems.append("weight_bytes") + started_at = _parse_timezone_aware_timestamp( + record.get("started_at_utc") + ) + completed_at = _parse_timezone_aware_timestamp( + record.get("completed_at_utc") + ) + if started_at is None: + problems.append("started_at_utc") + if completed_at is None: + problems.append("completed_at_utc") + elif started_at is not None and completed_at < started_at: + problems.append("completed_at_utc") + worker_seconds = record.get("worker_seconds") + if ( + isinstance(worker_seconds, bool) + or not isinstance(worker_seconds, (int, float)) + or not math.isfinite(worker_seconds) + or worker_seconds < 0 + ): + problems.append("worker_seconds") + status = record.get("status") + if status == "success": + expected_run_dir = f"items/{item_id}" + if record.get("run_dir") != expected_run_dir: + problems.append("run_dir") + expected_summary_path = f"{expected_run_dir}/summary.json" + if record.get("summary_path") != expected_summary_path: + problems.append("summary_path") + expected_backend = ( + assignment.get("backend") if assignment is not None else None + ) + for field in ("requested_backend", "backend"): + if expected_backend is not None: + if record.get(field) != expected_backend: + problems.append(field) + elif ( + not isinstance(record.get(field), str) + or not record[field] + ): + problems.append(field) + if ( + not isinstance(record.get("device"), str) + or not record["device"] + ): + problems.append("device") + if not isinstance(record.get("runtime"), Mapping): + problems.append("runtime") + if not isinstance(record.get("timings_sec"), Mapping): + problems.append("timings_sec") + try: + _validated_timings( + record.get("wall_sec"), + field=f"item {item_id!r} wall_sec", + ) + except (TypeError, ValueError): + problems.append("wall_sec") + elif status == "failed": + error = record.get("error") + if not isinstance(error, Mapping) or not all( + isinstance(error.get(field), str) and error[field] + for field in ("type", "message") + ): + problems.append("error") + else: + problems.append("status") + if problems: + errors.append( + { + "item_id": item_id, + "type": "InvalidTerminalRecord", + "message": ( + f"record {index} has invalid field(s): " + + ", ".join(sorted(set(problems))) + ), + } + ) + return errors + + +def _audit_shard_results( + shards: Sequence[Sequence[WorkItem]], + shard_results: Sequence[Mapping[str, Any]], + terminal_records: Sequence[Mapping[str, Any]], + *, + placements: Sequence[Placement], + allow_loopback_alias: bool, + backend: str, +) -> dict[str, Any]: + worker_count = len(shards) + placement_by_worker = { + placement.worker_id: placement for placement in placements + } + valid_results: list[tuple[int, Mapping[str, Any]]] = [] + invalid_results: list[dict[str, Any]] = [] + for index, result in enumerate(shard_results): + worker_id = _strict_integer(result.get("worker_id")) + if worker_id is None: + invalid_results.append( + { + "result_index": index, + "field": "worker_id", + "message": "worker_id must be a non-boolean integer", + } + ) + continue + valid_results.append((worker_id, result)) + worker_ids = [worker_id for worker_id, _ in valid_results] + expected = set(range(worker_count)) + observed = set(worker_ids) + duplicates = sorted( + worker_id for worker_id in observed if worker_ids.count(worker_id) > 1 + ) + missing = sorted(expected - observed) + unexpected = sorted(observed - expected) + terminal_statuses: dict[str, list[Any]] = {} + for record in terminal_records: + item_id = record.get("item_id") + if isinstance(item_id, str): + terminal_statuses.setdefault(item_id, []).append( + record.get("status") + ) + mismatched: list[dict[str, Any]] = [] + physical_identities: list[ + tuple[int, str, frozenset[tuple[str, str]]] + ] = [] + for worker_id, result in valid_results: + if worker_id not in expected: + continue + shard = shards[worker_id] + expected_fields = { + "item_count": len(shard), + "weight_bytes": sum(item.weight_bytes for item in shard), + "item_ids_sha256": _item_ids_sha256(shard), + } + fields = { + field + for field in ("item_count", "weight_bytes") + if (value := _strict_integer(result.get(field))) is None + or value < 0 + or value != expected_fields[field] + } + if ( + result.get("item_ids_sha256") + != expected_fields["item_ids_sha256"] + ): + fields.add("item_ids_sha256") + if result.get("schema") != "cuphoton.xpois.dragon-shard/v1": + fields.add("schema") + started_at = _parse_timezone_aware_timestamp( + result.get("started_at_utc") + ) + completed_at = _parse_timezone_aware_timestamp( + result.get("completed_at_utc") + ) + if started_at is None: + fields.add("started_at_utc") + if completed_at is None: + fields.add("completed_at_utc") + elif started_at is not None and completed_at < started_at: + fields.add("completed_at_utc") + worker_wall_sec = result.get("worker_wall_sec") + if ( + isinstance(worker_wall_sec, bool) + or not isinstance(worker_wall_sec, (int, float)) + or not math.isfinite(worker_wall_sec) + or worker_wall_sec < 0 + ): + fields.add("worker_wall_sec") + status = result.get("status") + if status not in {"success", "failed"}: + fields.add("status") + success_count = _strict_integer(result.get("success_count")) + failed_count = _strict_integer(result.get("failed_count")) + if success_count is None or success_count < 0: + fields.add("success_count") + if failed_count is None or failed_count < 0: + fields.add("failed_count") + if ( + success_count is not None + and failed_count is not None + and success_count + failed_count != len(shard) + ): + fields.update({"success_count", "failed_count"}) + if success_count is not None and failed_count is not None: + expected_status = "success" if failed_count == 0 else "failed" + if status != expected_status: + fields.add("status") + terminal_success_count = sum( + status == "success" + for item in shard + for status in terminal_statuses.get(item.item_id, ()) + ) + terminal_failed_count = sum( + status == "failed" + for item in shard + for status in terminal_statuses.get(item.item_id, ()) + ) + if ( + success_count is not None + and success_count != terminal_success_count + ): + fields.add("success_count") + if failed_count is not None and failed_count != terminal_failed_count: + fields.add("failed_count") + if status != "success": + fields.add("status") + if failed_count != 0: + fields.add("failed_count") + placement = placement_by_worker.get(worker_id) + provenance = result.get("provenance") + if placement is None or not _valid_shard_provenance( + provenance, + placement=placement, + allow_loopback_alias=allow_loopback_alias, + backend=backend, + ): + fields.add("provenance") + else: + assert isinstance(provenance, Mapping) + physical_identity = _stable_gpu_physical_ids( + provenance.get("gpu") + ) + assert physical_identity is not None + physical_identities.append( + (worker_id, placement.host, physical_identity) + ) + record_write_errors = result.get("record_write_errors") + if not isinstance(record_write_errors, list) or any( + not isinstance(error, Mapping) for error in record_write_errors + ): + fields.add("record_write_errors") + elif record_write_errors: + fields.add("record_write_errors") + try: + _validated_timings( + result.get("timings_sec"), + field=f"worker {worker_id} timings_sec", + ) + except (TypeError, ValueError): + fields.add("timings_sec") + if fields: + mismatched.append( + {"worker_id": worker_id, "fields": sorted(fields)} + ) + duplicate_physical_gpu_worker_ids: set[int] = set() + for index, (worker_id, host, identity) in enumerate(physical_identities): + remaining_identities = physical_identities[index + 1 :] + for ( + other_worker_id, + other_host, + other_identity, + ) in remaining_identities: + if ( + worker_id != other_worker_id + and _hostnames_match(host, other_host) + and identity & other_identity + ): + duplicate_physical_gpu_worker_ids.update( + (worker_id, other_worker_id) + ) + marked_worker_ids: set[int] = set() + for mismatch in mismatched: + worker_id = mismatch["worker_id"] + if worker_id in duplicate_physical_gpu_worker_ids: + mismatch["fields"] = sorted( + set(mismatch["fields"]) | {"provenance"} + ) + marked_worker_ids.add(worker_id) + for worker_id in sorted( + duplicate_physical_gpu_worker_ids - marked_worker_ids + ): + mismatched.append({"worker_id": worker_id, "fields": ["provenance"]}) + return { + "ok": not missing + and not duplicates + and not unexpected + and not mismatched + and not invalid_results, + "expected_count": worker_count, + "observed_count": len(shard_results), + "missing_worker_ids": missing, + "duplicate_worker_ids": duplicates, + "unexpected_worker_ids": unexpected, + "duplicate_physical_gpu_worker_ids": sorted( + duplicate_physical_gpu_worker_ids + ), + "mismatched_shards": mismatched, + "invalid_results": invalid_results, + } + + +def _aggregate_item_timings( + records: Sequence[Mapping[str, Any]], +) -> tuple[dict[str, dict[str, float]], list[dict[str, str]]]: + valid: list[dict[str, float]] = [] + errors: list[dict[str, str]] = [] + for record in records: + item_id = str(record.get("item_id", "")) + raw = record.get("timings_sec") + if raw is None: + continue + try: + valid.append( + _validated_timings(raw, field=f"item {item_id!r} timings_sec") + ) + except (TypeError, ValueError) as exc: + errors.append( + { + "item_id": item_id, + "type": "InvalidTimingRecord", + "message": str(exc), + } + ) + phases = sorted({phase for timings in valid for phase in timings}) + result: dict[str, dict[str, float]] = {} + for phase in phases: + values = [timings[phase] for timings in valid if phase in timings] + result[phase] = { + "sum": sum(values), + "max": max(values, default=0.0), + "mean": sum(values) / len(values) if values else 0.0, + } + return result, errors + + +def _validated_timings(value: Any, *, field: str) -> dict[str, float]: + if not isinstance(value, Mapping): + raise TypeError(f"{field} must be a mapping") + timings: dict[str, float] = {} + for phase, elapsed in value.items(): + if not isinstance(phase, str) or not phase: + raise ValueError(f"{field} phase names must be non-empty strings") + if ( + isinstance(elapsed, bool) + or not isinstance(elapsed, (int, float)) + or not math.isfinite(elapsed) + or elapsed < 0 + ): + raise ValueError( + f"{field} phase {phase!r} must be a finite non-negative " + "number" + ) + timings[phase] = float(elapsed) + return timings + + +def _parse_timezone_aware_timestamp(value: Any) -> datetime | None: + if not isinstance(value, str) or not value: + return None + try: + parsed = datetime.fromisoformat(value) + except ValueError: + return None + if parsed.tzinfo is None or parsed.utcoffset() is None: + return None + return parsed + + +def _stable_gpu_physical_ids( + value: Any, +) -> frozenset[tuple[str, str]] | None: + if ( + not isinstance(value, Mapping) + or "identity_error" not in value + or value["identity_error"] is not None + ): + return None + identities = frozenset( + (field, identifier.strip().casefold()) + for field in ("uuid", "pci_bus_id") + if isinstance((identifier := value.get(field)), str) + and identifier.strip() + ) + return identities or None + + +def _valid_shard_provenance( + value: Any, + *, + placement: Placement, + allow_loopback_alias: bool, + backend: str, +) -> bool: + if not isinstance(value, Mapping): + return False + worker_id = _strict_integer(value.get("worker_id")) + requested_gpu_id = _strict_integer(value.get("requested_gpu_id")) + pid = _strict_integer(value.get("pid")) + requested_host = value.get("requested_host") + hostname = value.get("hostname") + visibility = value.get("cuda_visible_devices") + if ( + worker_id != placement.worker_id + or requested_host != placement.host + or requested_gpu_id != placement.gpu_id + or pid is None + or pid <= 0 + or not isinstance(hostname, str) + or not hostname + or not _hostnames_match( + requested_host, + hostname, + allow_loopback_alias=allow_loopback_alias, + ) + or not isinstance(visibility, str) + or [token.strip() for token in visibility.split(",") if token.strip()] + != [str(placement.gpu_id)] + ): + return False + gpu = value.get("gpu") + if not isinstance(gpu, Mapping) or not gpu: + return False + gpu_backend = gpu.get("backend") + expected_identity_backend = ( + "cupy" if backend in {"cupy", "cutile"} else backend + ) + return ( + isinstance(gpu_backend, str) + and bool(gpu_backend) + and gpu_backend == expected_identity_backend + and _stable_gpu_physical_ids(gpu) is not None + ) + + +def _strict_integer(value: Any) -> int | None: + if isinstance(value, bool) or not isinstance(value, int): + return None + return value + + +def _shard_result_sort_key(result: Mapping[str, Any]) -> tuple[int, int]: + worker_id = _strict_integer(result.get("worker_id")) + return (worker_id is None, worker_id if worker_id is not None else 0) + + +def _item_ids_sha256(items: Sequence[WorkItem]) -> str: + encoded = "\n".join(item.item_id for item in items).encode() + return hashlib.sha256(encoded).hexdigest() + + +def _load_dragon_api() -> _DragonAPI: + try: + from dragon.infrastructure.policy import Policy + from dragon.native.machine import Node, System + from dragon.native.process import ProcessTemplate + from dragon.native.process_group import ProcessGroup + from dragon.native.queue import Queue + except (ImportError, OSError) as exc: + raise RuntimeError( + "The Dragon executor requires the dragonhpc runtime; " + "install it in this Python environment on every node " + "and launch with dragon" + ) from exc + return _DragonAPI( + System=System, + Node=Node, + Policy=Policy, + ProcessGroup=ProcessGroup, + ProcessTemplate=ProcessTemplate, + Queue=Queue, + ) + + +def _distribution_version(name: str) -> str | None: + try: + return version(name) + except PackageNotFoundError: + return None + + +__all__ = [ + "DragonBatchResult", + "discover_gpu_placements", + "run_dragon_image_pair_batch", +] diff --git a/src/cuphoton/xpois/workflows.py b/src/cuphoton/xpois/workflows.py index 498cb38..fb7759b 100644 --- a/src/cuphoton/xpois/workflows.py +++ b/src/cuphoton/xpois/workflows.py @@ -119,12 +119,17 @@ def run_constant_kernel_fit( Run directory and JSON-compatible summary of saved artifacts. """ + workflow_start = time.perf_counter() + run_setup_start = time.perf_counter() run_dir = _resolve_run_dir(output_root, name, run_prefix) run_dir.mkdir(parents=True, exist_ok=False) artifacts_dir = run_dir / "artifacts" artifacts_dir.mkdir() + run_setup_sec = time.perf_counter() - run_setup_start try: + prepare_start = time.perf_counter() + input_read_sec = 0.0 if variance_hdu is not None and variance_path is None: raise ValueError("variance_hdu requires a variance image path") if fit_mask_path is not None and auto_stamp_mask: @@ -164,21 +169,25 @@ def run_constant_kernel_fit( "reference/target mask inputs require a non-'none' " "mask_policy" ) - reference, _, used_reference_hdu = load_image_with_wcs( - reference_path, - hdu=reference_hdu, - ) - target, _, used_target_hdu = load_image_with_wcs( - target_path, - hdu=target_hdu, - ) - variance = None - used_variance_hdu = None - if variance_path is not None: - variance, _, used_variance_hdu = load_variance_with_wcs( - variance_path, - hdu=variance_hdu, + input_read_start = time.perf_counter() + try: + reference, _, used_reference_hdu = load_image_with_wcs( + reference_path, + hdu=reference_hdu, + ) + target, _, used_target_hdu = load_image_with_wcs( + target_path, + hdu=target_hdu, ) + variance = None + used_variance_hdu = None + if variance_path is not None: + variance, _, used_variance_hdu = load_variance_with_wcs( + variance_path, + hdu=variance_hdu, + ) + finally: + input_read_sec += time.perf_counter() - input_read_start crop_metadata: dict[str, int] | None = None if crop_y0 is not None: crop_metadata = { @@ -199,18 +208,24 @@ def run_constant_kernel_fit( if normalized_mask_policy != MASK_POLICY_NONE: reference_mask_source = reference_mask_path or reference_path target_mask_source = target_mask_path or target_path - reference_mask, used_reference_mask_hdu, reference_plane_map = ( - load_mask_with_planes( + input_read_start = time.perf_counter() + try: + ( + reference_mask, + used_reference_mask_hdu, + reference_plane_map, + ) = load_mask_with_planes( reference_mask_source, hdu=reference_mask_hdu, ) - ) - target_mask, used_target_mask_hdu, target_plane_map = ( - load_mask_with_planes( - target_mask_source, - hdu=target_mask_hdu, + target_mask, used_target_mask_hdu, target_plane_map = ( + load_mask_with_planes( + target_mask_source, + hdu=target_mask_hdu, + ) ) - ) + finally: + input_read_sec += time.perf_counter() - input_read_start if crop_metadata is not None: reference_mask = apply_rectangular_cutout( reference_mask, @@ -267,10 +282,14 @@ def run_constant_kernel_fit( fit_mask_kind: str | None = None fit_mask_metadata: dict[str, Any] | None = None if fit_mask_path is not None: - fit_mask = _load_fit_mask( - fit_mask_path.expanduser().resolve(), - expected_shape=target.shape, - ) + input_read_start = time.perf_counter() + try: + fit_mask = _load_fit_mask( + fit_mask_path.expanduser().resolve(), + expected_shape=target.shape, + ) + finally: + input_read_sec += time.perf_counter() - input_read_start fit_mask_kind = "explicit_mask" elif auto_stamp_mask: selection_image = np.where( @@ -289,6 +308,9 @@ def run_constant_kernel_fit( fit_mask_kind = "auto_stamp_mask" fit_mask_metadata = auto_mask.to_metadata() + prepare_sec = time.perf_counter() - prepare_start + preprocess_sec = max(0.0, prepare_sec - input_read_sec) + solve_start = time.perf_counter() result = solve_constant_kernel( reference, target, @@ -300,11 +322,19 @@ def run_constant_kernel_fit( flux_conserve=flux_conserve, backend=backend, ) + solve_sec = time.perf_counter() - solve_start + postprocess_start = time.perf_counter() + artifact_write_sec = 0.0 + review_generation_and_write_sec = 0.0 + output_start = time.perf_counter() saved = _save_artifacts(artifacts_dir, result) + artifact_write_sec += time.perf_counter() - output_start if fit_mask_metadata is not None: fit_mask_metadata_path = artifacts_dir / "fit_mask_metadata.json" + output_start = time.perf_counter() _write_json(fit_mask_metadata_path, fit_mask_metadata) + artifact_write_sec += time.perf_counter() - output_start saved["fit_mask_metadata"] = str( fit_mask_metadata_path.relative_to(artifacts_dir.parent) ) @@ -312,7 +342,9 @@ def run_constant_kernel_fit( preprocessing_metadata_path = ( artifacts_dir / "input_mask_metadata.json" ) + output_start = time.perf_counter() _write_json(preprocessing_metadata_path, preprocessing_metadata) + artifact_write_sec += time.perf_counter() - output_start saved["input_mask_metadata"] = str( preprocessing_metadata_path.relative_to(artifacts_dir.parent) ) @@ -342,13 +374,16 @@ def run_constant_kernel_fit( ) ), } + review_start = time.perf_counter() review_saved, hotspots = write_review_metadata( artifacts_dir, run_name=run_dir.name, residual=result.residual, review_metrics=review_metrics, ) + review_generation_and_write_sec += time.perf_counter() - review_start saved.update(review_saved) + review_start = time.perf_counter() interactive_saved = write_interactive_review_artifact( artifacts_dir, run_name=run_dir.name, @@ -410,12 +445,36 @@ def run_constant_kernel_fit( else 1.0 ), ) + review_generation_and_write_sec += time.perf_counter() - review_start saved.update(interactive_saved) runtime = runtime_metadata( backend=result.backend, dtype=str(result.kernel.dtype), ) + postprocess_total_sec = time.perf_counter() - postprocess_start + postprocess_other_sec = max( + 0.0, + postprocess_total_sec + - artifact_write_sec + - review_generation_and_write_sec, + ) + timings_sec = { + "run_setup_sec": run_setup_sec, + "input_read_sec": input_read_sec, + "preprocess_sec": preprocess_sec, + "solve_sec": solve_sec, + "artifact_write_sec": artifact_write_sec, + "review_generation_and_write_sec": ( + review_generation_and_write_sec + ), + "postprocess_other_sec": postprocess_other_sec, + } + wall_sec = { + "workflow_before_summary_write": ( + time.perf_counter() - workflow_start + ) + } summary = { "workflow": workflow_name, "package_version": __version__, @@ -435,6 +494,8 @@ def run_constant_kernel_fit( "device": runtime["device"], "dtype": str(result.kernel.dtype), "runtime": runtime, + "timings_sec": timings_sec, + "wall_sec": wall_sec, "fit_pixel_count": result.fit_pixel_count, "chi2": result.chi2, "dof": result.dof, diff --git a/tests/core/test_bulk.py b/tests/core/test_bulk.py new file mode 100644 index 0000000..54528c4 --- /dev/null +++ b/tests/core/test_bulk.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import os + +import pytest + +from cuphoton.core.bulk import ( + Placement, + WorkItem, + atomic_write_json, + audit_terminal_records, + partition_byte_balanced, +) + + +def _item(item_id: str, weight: int) -> WorkItem: + return WorkItem(item_id, {"path": f"/{item_id}"}, weight) + + +def test_byte_balanced_partition_is_deterministic() -> None: + items = ( + _item("small", 1), + _item("large-b", 8), + _item("medium", 6), + _item("large-a", 8), + ) + + first = partition_byte_balanced(items, 2) + second = partition_byte_balanced(tuple(reversed(items)), 2) + + assert first == second + assert [[item.item_id for item in shard] for shard in first] == [ + ["large-a", "medium"], + ["large-b", "small"], + ] + + +def test_byte_balanced_partition_distributes_zero_weight_items() -> None: + shards = partition_byte_balanced( + tuple(_item(f"item-{index}", 0) for index in range(7)), + 3, + ) + + assert [[item.item_id for item in shard] for shard in shards] == [ + ["item-0", "item-3", "item-6"], + ["item-1", "item-4"], + ["item-2", "item-5"], + ] + + +def test_byte_balanced_partition_distributes_equal_weight_items() -> None: + shards = partition_byte_balanced( + tuple(_item(f"item-{index}", 5) for index in range(6)), + 3, + ) + + assert [[item.item_id for item in shard] for shard in shards] == [ + ["item-0", "item-3"], + ["item-1", "item-4"], + ["item-2", "item-5"], + ] + + +def test_partition_rejects_duplicate_item_ids() -> None: + with pytest.raises(ValueError, match="duplicate work item IDs"): + partition_byte_balanced((_item("same", 1), _item("same", 2)), 2) + + +@pytest.mark.parametrize("worker_count", [True, 2.0, "2"]) +def test_partition_rejects_non_integer_worker_count(worker_count) -> None: + with pytest.raises(TypeError, match="must be an integer"): + partition_byte_balanced((_item("one", 1),), worker_count) + + +@pytest.mark.parametrize("worker_id", [False, 1.0, "1"]) +def test_placement_rejects_non_integer_worker_id(worker_id) -> None: + with pytest.raises(TypeError, match="must be an integer"): + Placement(worker_id=worker_id, host="node", gpu_id=0) + + +def test_terminal_record_audit_reports_all_identity_failures() -> None: + audit = audit_terminal_records( + ["a", "b", "c"], + [ + {"item_id": "a", "status": "success"}, + {"item_id": "a", "status": "failed"}, + {"item_id": "outside", "status": "unknown"}, + ], + ) + + assert audit == { + "ok": False, + "expected_count": 3, + "terminal_record_count": 3, + "missing_item_ids": ["b", "c"], + "duplicate_item_ids": ["a"], + "unexpected_item_ids": ["outside"], + "invalid_status_item_ids": ["outside"], + "failed_item_ids": ["a"], + } + + +def test_terminal_audit_handles_failed_record_without_item_id() -> None: + audit = audit_terminal_records(["expected"], [{"status": "failed"}]) + + assert audit["ok"] is False + assert audit["missing_item_ids"] == ["expected"] + assert audit["unexpected_item_ids"] == [""] + assert audit["failed_item_ids"] == [] + + +def test_atomic_write_json_replaces_complete_mapping( + monkeypatch, tmp_path +) -> None: + path = tmp_path / "nested" / "record.json" + fsynced = [] + original_fsync = os.fsync + + def record_fsync(fd): + fsynced.append(fd) + original_fsync(fd) + + monkeypatch.setattr(os, "fsync", record_fsync) + + atomic_write_json(path, {"status": "starting"}) + atomic_write_json(path, {"status": "success", "count": 2}) + + assert json.loads(path.read_text()) == {"status": "success", "count": 2} + assert list(path.parent.glob(".*.tmp-*")) == [] + assert len(fsynced) == 4 + + +def test_work_item_rejects_non_json_payload() -> None: + with pytest.raises(ValueError, match="JSON-compatible"): + WorkItem("item", {"bad": object()}) diff --git a/tests/core/test_cli_contract.py b/tests/core/test_cli_contract.py index 7fe2b75..6771d24 100644 --- a/tests/core/test_cli_contract.py +++ b/tests/core/test_cli_contract.py @@ -72,16 +72,16 @@ def test_public_command_surface_counts_are_exact() -> None: assert per_group == [ ("xdr", 1, 0, 13, 0), ("xfit", 3, 3, 17, 1), - ("xpois", 6, 6, 95, 1), + ("xpois", 7, 7, 117, 1), ("xscan", 42, 42, 167, 1), ("xrep", 6, 6, 103, 1), ("xray", 33, 31, 378, 1), ] assert len(per_group) == 6 - assert sum(item[1] for item in per_group) == 91 - assert sum(item[1] + item[4] for item in per_group) == 96 - assert sum(item[2] for item in per_group) == 88 - assert sum(item[3] for item in per_group) == 773 + assert sum(item[1] for item in per_group) == 92 + assert sum(item[1] + item[4] for item in per_group) == 97 + assert sum(item[2] for item in per_group) == 89 + assert sum(item[3] for item in per_group) == 795 def test_public_registry_order_and_component_derivations() -> None: diff --git a/tests/xpois/test_batch.py b/tests/xpois/test_batch.py new file mode 100644 index 0000000..be12eae --- /dev/null +++ b/tests/xpois/test_batch.py @@ -0,0 +1,321 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json + +import numpy as np +import pytest + +from cuphoton.xpois.batch import ( + BatchFitOptions, + load_image_pair_manifest, + preflight_image_pair_manifest, + run_image_pair_item, +) + + +def _options(**overrides) -> BatchFitOptions: + values = { + "kernel_shape": (9, 9), + "basis_sigmas": (1.5,), + "basis_degrees": (0,), + "backend": "cupy", + } + values.update(overrides) + return BatchFitOptions(**values) + + +def _write_array(path, shape=(32, 32)) -> None: + np.save(path, np.zeros(shape, dtype=np.float64), allow_pickle=False) + + +def test_manifest_resolves_paths_hashes_and_preflights(tmp_path) -> None: + _write_array(tmp_path / "reference.npy") + _write_array(tmp_path / "target.npy") + manifest_path = tmp_path / "pairs.json" + manifest_path.write_text( + json.dumps( + { + "schema": "cuphoton.xpois.image-pairs/v1", + "pairs": [ + { + "id": "pair-1", + "reference": "reference.npy", + "target": "target.npy", + } + ], + } + ) + ) + + manifest = load_image_pair_manifest(manifest_path) + preflight_image_pair_manifest(manifest, _options()) + + assert manifest.pairs[0].reference == (tmp_path / "reference.npy") + assert len(manifest.sha256) == 64 + identity = manifest.input_identity_payload() + assert identity["policy"] == "size-mtime-ns" + assert identity["content_hashes"] is False + assert len(identity["files"]) == 2 + assert ( + manifest.work_items()[0].weight_bytes + == 2 * (tmp_path / "reference.npy").stat().st_size + ) + + +def test_manifest_work_item_counts_reused_path_once(tmp_path) -> None: + shared_path = tmp_path / "shared.npy" + _write_array(shared_path) + manifest_path = tmp_path / "pairs.json" + manifest_path.write_text( + json.dumps( + { + "schema": "cuphoton.xpois.image-pairs/v1", + "pairs": [ + { + "id": "shared", + "reference": "shared.npy", + "target": "shared.npy", + } + ], + } + ) + ) + + manifest = load_image_pair_manifest(manifest_path) + item = manifest.work_items()[0] + + assert item.weight_bytes == shared_path.stat().st_size + assert len(item.payload["_input_identity"]) == 1 + + +def test_preflight_rejects_input_changed_after_manifest_load( + tmp_path, +) -> None: + _write_array(tmp_path / "reference.npy") + _write_array(tmp_path / "target.npy") + manifest_path = tmp_path / "pairs.yaml" + manifest_path.write_text( + """schema: cuphoton.xpois.image-pairs/v1 +pairs: + - id: changed + reference: reference.npy + target: target.npy +""" + ) + manifest = load_image_pair_manifest(manifest_path) + with (tmp_path / "target.npy").open("ab") as handle: + handle.write(b"changed") + + with pytest.raises(RuntimeError, match="changed since manifest load"): + preflight_image_pair_manifest(manifest, _options()) + + +def test_preflight_ignores_fits_table_before_image(tmp_path) -> None: + from astropy.io import fits + + reference = tmp_path / "reference.fits" + target = tmp_path / "target.fits" + table = fits.BinTableHDU.from_columns( + [ + fits.Column( + name="VALUE", + format="D", + array=np.arange(3, dtype=np.float64), + ) + ] + ) + fits.HDUList( + [ + fits.PrimaryHDU(), + table, + fits.ImageHDU(np.zeros((32, 32), dtype=np.float64)), + ] + ).writeto(reference) + fits.PrimaryHDU(np.zeros((32, 32), dtype=np.float64)).writeto(target) + manifest_path = tmp_path / "pairs.json" + manifest_path.write_text( + json.dumps( + { + "schema": "cuphoton.xpois.image-pairs/v1", + "pairs": [ + { + "id": "table-before-image", + "reference": "reference.fits", + "target": "target.fits", + } + ], + } + ) + ) + + manifest = load_image_pair_manifest(manifest_path) + + preflight_image_pair_manifest(manifest, _options()) + + +def test_manifest_rejects_duplicate_json_keys(tmp_path) -> None: + manifest_path = tmp_path / "pairs.json" + manifest_path.write_text( + '{"schema":"cuphoton.xpois.image-pairs/v1",' + '"schema":"duplicate","pairs":[]}' + ) + + with pytest.raises(ValueError, match="duplicate key"): + load_image_pair_manifest(manifest_path) + + +def test_preflight_rejects_shape_mismatch(tmp_path) -> None: + _write_array(tmp_path / "reference.npy", (32, 32)) + _write_array(tmp_path / "target.npy", (16, 16)) + manifest_path = tmp_path / "pairs.yaml" + manifest_path.write_text( + """schema: cuphoton.xpois.image-pairs/v1 +pairs: + - id: mismatch + reference: reference.npy + target: target.npy +""" + ) + + manifest = load_image_pair_manifest(manifest_path) + with pytest.raises(ValueError, match="does not match"): + preflight_image_pair_manifest(manifest, _options()) + + +def test_batch_options_reject_unknown_backend() -> None: + with pytest.raises(ValueError, match="unsupported fit backend"): + _options(backend="unknown") + + +@pytest.mark.parametrize("sigma", [float("nan"), float("inf"), float("-inf")]) +def test_batch_options_reject_nonfinite_basis_sigmas(sigma) -> None: + with pytest.raises(ValueError, match="finite and positive"): + _options(basis_sigmas=(sigma,)) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"auto_stamp_size": 30}, "positive odd integer"), + ({"auto_stamp_size": True}, "positive odd integer"), + ({"auto_stamp_count": 0}, "positive integer"), + ({"auto_stamp_count": True}, "positive integer"), + ({"auto_peak_percentile": float("nan")}, "finite number"), + ({"auto_peak_percentile": 101.0}, "finite number"), + ], +) +def test_batch_options_reject_invalid_auto_stamp_values( + overrides, message +) -> None: + with pytest.raises(ValueError, match=message): + _options(**overrides) + + +def test_batch_options_round_trip_worker_payload() -> None: + options = _options(crop_y0=1, crop_x0=2, crop_height=16, crop_width=17) + + assert BatchFitOptions.from_payload(options.to_payload()) == options + + +def test_image_pair_item_cpu_smoke(tmp_path) -> None: + reference = np.zeros((64, 64), dtype=np.float64) + reference[20:40, 20:40] = 1.0 + target = reference.copy() + reference_path = tmp_path / "reference.npy" + target_path = tmp_path / "target.npy" + np.save(reference_path, reference, allow_pickle=False) + np.save(target_path, target, allow_pickle=False) + manifest_path = tmp_path / "cpu-pairs.json" + manifest_path.write_text( + json.dumps( + { + "schema": "cuphoton.xpois.image-pairs/v1", + "pairs": [ + { + "id": "cpu-pair", + "reference": str(reference_path), + "target": str(target_path), + } + ], + } + ) + ) + item = load_image_pair_manifest(manifest_path).work_items()[0] + output_dir = tmp_path / "items" / "cpu-pair" + output_dir.parent.mkdir() + + result = run_image_pair_item( + item, + output_dir, + _options(backend="cpu"), + ) + + assert result["backend"] == "cpu" + assert result["timings_sec"]["input_read_sec"] >= 0 + assert result["timings_sec"]["solve_sec"] >= 0 + assert result["timings_sec"]["artifact_write_sec"] >= 0 + assert result["wall_sec"]["item_runner"] >= sum( + result["timings_sec"][field] + for field in ("input_read_sec", "solve_sec", "artifact_write_sec") + ) + assert result["wall_sec"]["workflow_before_summary_write"] >= sum( + result["timings_sec"].values() + ) + assert "item_runner_sec" not in result["timings_sec"] + assert "workflow_before_summary_write_sec" not in result["timings_sec"] + assert (output_dir / "summary.json").is_file() + assert (output_dir / "artifacts" / "residual.npy").is_file() + + +def test_item_preserves_workflow_error_and_postcheck_failure( + monkeypatch, tmp_path +) -> None: + from cuphoton.xpois import workflows + + reference_path = tmp_path / "reference.npy" + target_path = tmp_path / "target.npy" + _write_array(reference_path) + _write_array(target_path) + manifest_path = tmp_path / "pairs.json" + manifest_path.write_text( + json.dumps( + { + "schema": "cuphoton.xpois.image-pairs/v1", + "pairs": [ + { + "id": "mutated-failure", + "reference": str(reference_path), + "target": str(target_path), + } + ], + } + ) + ) + item = load_image_pair_manifest(manifest_path).work_items()[0] + + def fail_after_mutation(**kwargs): + del kwargs + with target_path.open("ab") as handle: + handle.write(b"changed") + raise ValueError("injected workflow failure") + + monkeypatch.setattr( + workflows, "run_constant_kernel_fit", fail_after_mutation + ) + + with pytest.raises(ValueError, match="injected workflow failure") as info: + run_image_pair_item( + item, + tmp_path / "items" / item.item_id, + _options(backend="cpu"), + ) + + assert any( + "post-item input identity verification failed" in note + and "input changed since manifest load" in note + for note in info.value.__notes__ + ) diff --git a/tests/xpois/test_cli.py b/tests/xpois/test_cli.py index 85bce45..35a75cd 100644 --- a/tests/xpois/test_cli.py +++ b/tests/xpois/test_cli.py @@ -5,6 +5,7 @@ from __future__ import annotations import json +from types import SimpleNamespace import pytest @@ -56,6 +57,139 @@ def test_help_for_benchmark_backends_command_case(capsys) -> None: assert "--repeats" in captured.out +def test_help_for_fit_batch_dragon_command(capsys) -> None: + rc = _run_cli(["help", "fit-batch-dragon"]) + captured = capsys.readouterr() + + assert rc == 0 + assert "Usage: cuphoton xpois fit-batch-dragon" in captured.out + assert "--manifest" in captured.out + assert "--max-workers" in captured.out + assert "--result-timeout-sec" in captured.out + assert "--worker-timeout-sec" in captured.out + assert "[default: cupy]" in captured.out + assert "--reference" not in captured.out + assert "--target" not in captured.out + assert "--variance" not in captured.out + assert "--fit-mask" not in captured.out + + +def test_fit_batch_dragon_accepts_manifest_without_pair_arguments( + monkeypatch, tmp_path, capsys +) -> None: + from cuphoton.xpois import commands + + manifest = tmp_path / "pairs.yaml" + manifest.write_text("schema: cuphoton.xpois.image-pairs/v1\n") + summary = tmp_path / "run" / "summary.json" + seen = {} + + def _run_dragon_image_pair_batch(**kwargs): + seen.update(kwargs) + return SimpleNamespace( + status="success", + summary_path=summary, + to_dict=lambda: {"status": "success"}, + ) + + monkeypatch.setattr( + commands, + "run_dragon_image_pair_batch", + _run_dragon_image_pair_batch, + ) + + rc = _run_cli( + [ + "fit-batch-dragon", + "--manifest", + str(manifest), + "--output-dir", + str(tmp_path / "runs"), + ] + ) + captured = capsys.readouterr() + + assert rc == 0 + assert json.loads(captured.out) == {"status": "success"} + assert seen["manifest_path"] == manifest + assert seen["options"].backend == "cupy" + assert seen["worker_timeout_sec"] == 3600.0 + + +def test_fit_batch_dragon_wraps_invalid_options( + monkeypatch, tmp_path, capsys +) -> None: + from cuphoton.xpois import commands + + manifest = tmp_path / "pairs.yaml" + manifest.write_text("schema: cuphoton.xpois.image-pairs/v1\n") + called = False + + def should_not_run(**kwargs): + del kwargs + nonlocal called + called = True + + monkeypatch.setattr( + commands, + "run_dragon_image_pair_batch", + should_not_run, + ) + + rc = _run_cli( + [ + "fit-batch-dragon", + "--manifest", + str(manifest), + "--kernel-height", + "8", + ] + ) + captured = capsys.readouterr() + + assert rc == 1 + assert called is False + assert "kernel_shape must contain two positive odd values" in captured.err + assert "Traceback" not in captured.err + + +def test_fit_batch_dragon_rejects_auto_backend( + monkeypatch, tmp_path, capsys +) -> None: + from cuphoton.xpois import commands + + called = False + + def should_not_run(**kwargs): + del kwargs + nonlocal called + called = True + + monkeypatch.setattr( + commands, + "run_dragon_image_pair_batch", + should_not_run, + ) + + rc = _run_cli( + [ + "fit-batch-dragon", + "--manifest", + str(tmp_path / "pairs.yaml"), + "--backend", + "auto", + ] + ) + captured = capsys.readouterr() + + assert rc == 2 + assert called is False + assert "invalid choice: 'auto'" in captured.err + for backend in ("cupy", "cutile", "numba-cuda"): + assert backend in captured.err + assert "Traceback" not in captured.err + + def test_help_for_evaluate_subtraction_command(capsys) -> None: rc = _run_cli(["help", "evaluate-subtraction"]) captured = capsys.readouterr() diff --git a/tests/xpois/test_dragon.py b/tests/xpois/test_dragon.py new file mode 100644 index 0000000..32f99ee --- /dev/null +++ b/tests/xpois/test_dragon.py @@ -0,0 +1,1712 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import os +import queue +from types import SimpleNamespace + +import numpy as np +import pytest + +import cuphoton.xpois.dragon as dragon_module +from cuphoton.core.bulk import ( + Placement, + WorkItem, + atomic_write_json, + audit_terminal_records, +) +from cuphoton.xpois.batch import ( + BatchFitOptions, + load_image_pair_manifest, + run_image_pair_item, +) +from cuphoton.xpois.dragon import ( + _aggregate_item_timings, + _audit_shard_results, + _audit_terminal_record_contract, + _dragon_shard_worker, + _DragonAPI, + _execute_shard, + _hostnames_match, + _load_terminal_records, + _select_gpu_placements, + _singleton_cuda_visibility, + discover_gpu_placements, + run_dragon_image_pair_batch, +) + + +class _System: + nodes = (7, 11) + + +class _Node: + def __init__(self, node_id): + self.hostname = f"node-{node_id}" + self.gpus = [2, 5] if node_id == 7 else [4] + + +def _options(**overrides) -> BatchFitOptions: + values = { + "kernel_shape": (9, 9), + "basis_sigmas": (1.5,), + "basis_degrees": (0,), + "backend": "cupy", + } + values.update(overrides) + return BatchFitOptions(**values) + + +def _test_placements(worker_count: int) -> tuple[Placement, ...]: + return tuple( + Placement(worker_id=index, host=f"node-{index}", gpu_id=index + 3) + for index in range(worker_count) + ) + + +def _valid_shard_result( + shard: tuple[WorkItem, ...], + placement: Placement, + *, + backend: str = "cupy", +) -> dict[str, object]: + identity_backend = "cupy" if backend in {"cupy", "cutile"} else backend + return { + "schema": "cuphoton.xpois.dragon-shard/v1", + "worker_id": placement.worker_id, + "status": "success", + "item_count": len(shard), + "success_count": len(shard), + "failed_count": 0, + "weight_bytes": sum(item.weight_bytes for item in shard), + "item_ids_sha256": dragon_module._item_ids_sha256(shard), + "started_at_utc": "2026-08-21T00:00:00+00:00", + "completed_at_utc": "2026-08-21T00:00:01+00:00", + "worker_wall_sec": 1.0, + "timings_sec": {}, + "provenance": { + "worker_id": placement.worker_id, + "requested_host": placement.host, + "requested_gpu_id": placement.gpu_id, + "hostname": placement.host, + "pid": 1234, + "cuda_visible_devices": str(placement.gpu_id), + "gpu": { + "backend": identity_backend, + "name": "fake-gpu", + "uuid": f"GPU-worker-{placement.worker_id}", + "pci_bus_id": f"0000:{placement.gpu_id:02x}:00.0", + "identity_error": None, + }, + }, + "record_write_errors": [], + } + + +def test_discovery_uses_actual_noncontiguous_node_gpu_ids() -> None: + assert discover_gpu_placements(_System, _Node) == ( + Placement(worker_id=0, host="node-7", gpu_id=2), + Placement(worker_id=1, host="node-7", gpu_id=5), + Placement(worker_id=2, host="node-11", gpu_id=4), + ) + + +def test_worker_selection_round_robins_across_hosts() -> None: + placements = discover_gpu_placements(_System, _Node) + + assert _select_gpu_placements(placements, 2) == ( + Placement(worker_id=0, host="node-7", gpu_id=2), + Placement(worker_id=1, host="node-11", gpu_id=4), + ) + + +def test_singleton_visibility_rejects_multiple_devices(monkeypatch) -> None: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "2,5") + + with pytest.raises(RuntimeError, match="exactly one"): + _singleton_cuda_visibility() + + +def test_singleton_visibility_rejects_wrong_placed_device( + monkeypatch, +) -> None: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "5") + + with pytest.raises(RuntimeError, match="placement mismatch"): + _singleton_cuda_visibility(2) + + +def test_hostname_validation_accepts_short_fqdn_aliases() -> None: + assert _hostnames_match("gpu-node", "gpu-node.example.com") + assert _hostnames_match("GPU-NODE.EXAMPLE.COM.", "gpu-node") + assert not _hostnames_match( + "localhost", dragon_module.socket.gethostname() + ) + assert _hostnames_match( + "localhost", + dragon_module.socket.gethostname(), + allow_loopback_alias=True, + ) + assert not _hostnames_match("gpu-node-1", "gpu-node-2") + assert not _hostnames_match( + "gpu-node.dc1.example.com", "gpu-node.dc2.example.com" + ) + + +def test_shard_rejects_wrong_host_before_gpu_initialization( + monkeypatch, tmp_path +) -> None: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") + initialized = False + + def identity_loader(backend): + nonlocal initialized + initialized = True + return {"backend": backend} + + with pytest.raises(RuntimeError, match="placement mismatch"): + _execute_shard( + run_id="wrong-host", + run_dir=tmp_path, + placement=Placement( + worker_id=0, + host="definitely-not-this-host", + gpu_id=0, + ), + items=(), + options=_options(), + item_runner=lambda item, output, options: {}, + gpu_identity_loader=identity_loader, + require_clean_cuda_imports=False, + ) + assert initialized is False + + +def test_shard_rejects_wrong_gpu_before_gpu_initialization( + monkeypatch, tmp_path +) -> None: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "5") + initialized = False + + def identity_loader(backend): + nonlocal initialized + initialized = True + return {"backend": backend} + + with pytest.raises(RuntimeError, match="GPU placement mismatch"): + _execute_shard( + run_id="wrong-gpu", + run_dir=tmp_path, + placement=Placement( + worker_id=0, + host=dragon_module.socket.gethostname(), + gpu_id=2, + ), + items=(), + options=_options(), + item_runner=lambda item, output, options: {}, + gpu_identity_loader=identity_loader, + require_clean_cuda_imports=False, + ) + assert initialized is False + + +def test_shard_rejects_cuda_import_before_worker_preflight( + monkeypatch, tmp_path +) -> None: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "2") + monkeypatch.setitem(dragon_module.sys.modules, "cupy", object()) + initialized = False + + def identity_loader(backend): + nonlocal initialized + initialized = True + return {"backend": backend} + + with pytest.raises(RuntimeError, match="before Dragon worker placement"): + _execute_shard( + run_id="premature-cuda", + run_dir=tmp_path, + placement=Placement( + worker_id=0, + host=dragon_module.socket.gethostname(), + gpu_id=2, + ), + items=(), + options=_options(), + item_runner=lambda item, output, options: {}, + gpu_identity_loader=identity_loader, + ) + assert initialized is False + + +def test_shard_persists_success_and_failure_then_continues( + monkeypatch, tmp_path +) -> None: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "5") + run_dir = tmp_path / "run" + for name in ("items", "records", "workers"): + (run_dir / name).mkdir(parents=True, exist_ok=True) + items = ( + WorkItem("good", {"id": "good"}, 10), + WorkItem("bad", {"id": "bad"}, 20), + ) + visited = [] + + def item_runner(item, output_dir, options): + visited.append(item.item_id) + if item.item_id == "bad": + output_dir.mkdir(parents=True) + (output_dir / "partial.npy").write_bytes(b"partial") + error = ValueError("injected ordinary failure") + error.add_note("post-item input identity verification failed") + raise error + return { + "run_dir": str(output_dir), + "summary_path": str(output_dir / "summary.json"), + "timings_sec": {"solve": 1.25}, + } + + result = _execute_shard( + run_id="test-run", + run_dir=run_dir, + placement=Placement( + worker_id=0, + host=dragon_module.socket.gethostname(), + gpu_id=5, + ), + items=items, + options=_options(), + item_runner=item_runner, + gpu_identity_loader=lambda backend: { + "backend": backend, + "uuid": "GPU-test", + "pci_bus_id": "0000:01:00.0", + }, + require_clean_cuda_imports=False, + ) + + assert visited == ["good", "bad"] + assert result["status"] == "failed" + assert result["success_count"] == 1 + assert result["failed_count"] == 1 + good = json.loads((run_dir / "records" / "good.json").read_text()) + bad = json.loads((run_dir / "records" / "bad.json").read_text()) + assert good["status"] == "success" + assert good["summary_path"] == "items/good/summary.json" + assert bad["status"] == "failed" + assert bad["error"] == { + "type": "ValueError", + "message": "injected ordinary failure", + "notes": "post-item input identity verification failed", + } + assert (run_dir / "items" / "bad" / "partial.npy").is_file() + assert not list((run_dir / "records").glob(".*.tmp-*")) + assert result["provenance"]["cuda_visible_devices"] == "5" + assert result["provenance"]["gpu"]["uuid"] == "GPU-test" + + +def test_shard_result_audit_rejects_missing_and_duplicate_workers() -> None: + shards = ( + (WorkItem("zero", {}, 10),), + (WorkItem("one", {}, 20),), + (WorkItem("two", {}, 30),), + ) + placements = _test_placements(len(shards)) + incomplete = _valid_shard_result(shards[0], placements[0]) + for field in ("item_count", "weight_bytes", "item_ids_sha256"): + incomplete.pop(field) + audit = _audit_shard_results( + shards, + [ + incomplete, + dict(incomplete), + {"worker_id": 4}, + ], + [{"item_id": "zero", "status": "success"}], + placements=placements, + allow_loopback_alias=False, + backend="cupy", + ) + + assert audit["ok"] is False + assert audit["missing_worker_ids"] == [1, 2] + assert audit["duplicate_worker_ids"] == [0] + assert audit["unexpected_worker_ids"] == [4] + assert audit["mismatched_shards"] == [ + { + "worker_id": 0, + "fields": ["item_count", "item_ids_sha256", "weight_bytes"], + }, + { + "worker_id": 0, + "fields": ["item_count", "item_ids_sha256", "weight_bytes"], + }, + ] + + +def test_shard_result_audit_quarantines_invalid_worker_id() -> None: + shards = ((WorkItem("zero", {}, 10),),) + placements = _test_placements(1) + + audit = _audit_shard_results( + shards, + [{"worker_id": "not-an-integer"}], + [], + placements=placements, + allow_loopback_alias=False, + backend="cupy", + ) + + assert audit["ok"] is False + assert audit["missing_worker_ids"] == [0] + assert audit["invalid_results"] == [ + { + "result_index": 0, + "field": "worker_id", + "message": "worker_id must be a non-boolean integer", + } + ] + + +def test_shard_result_audit_rejects_status_count_mismatch() -> None: + shard = (WorkItem("zero", {}, 10),) + placements = _test_placements(1) + result = _valid_shard_result(shard, placements[0]) + result["status"] = "failed" + + audit = _audit_shard_results( + (shard,), + [result], + [{"item_id": "zero", "status": "success"}], + placements=placements, + allow_loopback_alias=False, + backend="cupy", + ) + + assert audit["ok"] is False + assert audit["mismatched_shards"] == [ + {"worker_id": 0, "fields": ["status"]} + ] + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("item_count", True), + ("item_count", 1.0), + ("weight_bytes", True), + ("weight_bytes", 10.0), + ], +) +def test_shard_result_audit_requires_strict_integral_totals( + field, value +) -> None: + shard = (WorkItem("zero", {}, 10),) + placements = _test_placements(1) + result = _valid_shard_result(shard, placements[0]) + result[field] = value + + audit = _audit_shard_results( + (shard,), + [result], + [{"item_id": "zero", "status": "success"}], + placements=placements, + allow_loopback_alias=False, + backend="cupy", + ) + + assert audit["ok"] is False + assert audit["mismatched_shards"] == [{"worker_id": 0, "fields": [field]}] + + +def test_shard_result_audit_rejects_reported_failure() -> None: + shard = (WorkItem("zero", {}, 10),) + placements = _test_placements(1) + result = _valid_shard_result(shard, placements[0]) + result.update(status="failed", success_count=0, failed_count=1) + + audit = _audit_shard_results( + (shard,), + [result], + [{"item_id": "zero", "status": "failed"}], + placements=placements, + allow_loopback_alias=False, + backend="cupy", + ) + + assert audit["ok"] is False + assert audit["mismatched_shards"] == [ + {"worker_id": 0, "fields": ["failed_count", "status"]} + ] + + +def test_shard_result_audit_cross_checks_terminal_status_counts() -> None: + shard = (WorkItem("zero", {}, 10),) + placements = _test_placements(1) + result = _valid_shard_result(shard, placements[0]) + + audit = _audit_shard_results( + (shard,), + [result], + [{"item_id": "zero", "status": "failed"}], + placements=placements, + allow_loopback_alias=False, + backend="cupy", + ) + + assert audit["ok"] is False + assert audit["mismatched_shards"] == [ + {"worker_id": 0, "fields": ["failed_count", "success_count"]} + ] + + +def test_shard_result_audit_rejects_record_write_errors() -> None: + shard = (WorkItem("zero", {}, 10),) + placements = _test_placements(1) + result = _valid_shard_result(shard, placements[0]) + result["record_write_errors"] = [{"item_id": "zero"}] + + audit = _audit_shard_results( + (shard,), + [result], + [{"item_id": "zero", "status": "success"}], + placements=placements, + allow_loopback_alias=False, + backend="cupy", + ) + + assert audit["ok"] is False + assert audit["mismatched_shards"] == [ + {"worker_id": 0, "fields": ["record_write_errors"]} + ] + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("started_at_utc", "2026-08-21T00:00:00"), + ("completed_at_utc", "not-a-timestamp"), + ("worker_wall_sec", True), + ("worker_wall_sec", -1.0), + ("worker_wall_sec", float("inf")), + ], +) +def test_shard_result_audit_rejects_invalid_worker_timing( + field, value +) -> None: + shard = (WorkItem("zero", {}, 10),) + placements = _test_placements(1) + result = _valid_shard_result(shard, placements[0]) + result[field] = value + + audit = _audit_shard_results( + (shard,), + [result], + [{"item_id": "zero", "status": "success"}], + placements=placements, + allow_loopback_alias=False, + backend="cupy", + ) + + assert audit["ok"] is False + assert audit["mismatched_shards"] == [{"worker_id": 0, "fields": [field]}] + + +def test_shard_result_audit_rejects_decreasing_timestamps() -> None: + shard = (WorkItem("zero", {}, 10),) + placements = _test_placements(1) + result = _valid_shard_result(shard, placements[0]) + result["completed_at_utc"] = "2026-08-20T23:59:59+00:00" + + audit = _audit_shard_results( + (shard,), + [result], + [{"item_id": "zero", "status": "success"}], + placements=placements, + allow_loopback_alias=False, + backend="cupy", + ) + + assert audit["mismatched_shards"] == [ + {"worker_id": 0, "fields": ["completed_at_utc"]} + ] + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("worker_id", 1), + ("requested_host", "different-node"), + ("requested_gpu_id", 99), + ("hostname", "different-node"), + ("pid", 0), + ("cuda_visible_devices", "3,7"), + ("gpu", {}), + ("gpu", {"backend": "", "name": "fake-gpu"}), + ("gpu", {"backend": "numba-cuda", "name": "fake-gpu"}), + ], +) +def test_shard_result_audit_rejects_invalid_provenance(field, value) -> None: + shard = (WorkItem("zero", {}, 10),) + placements = _test_placements(1) + result = _valid_shard_result(shard, placements[0]) + provenance = result["provenance"] + assert isinstance(provenance, dict) + provenance[field] = value + + audit = _audit_shard_results( + (shard,), + [result], + [{"item_id": "zero", "status": "success"}], + placements=placements, + allow_loopback_alias=False, + backend="cupy", + ) + + assert audit["mismatched_shards"] == [ + {"worker_id": 0, "fields": ["provenance"]} + ] + + +@pytest.mark.parametrize("backend", ["cupy", "numba-cuda", "cutile"]) +def test_shard_result_audit_accepts_backend_identity(backend) -> None: + shard = (WorkItem("zero", {}, 10),) + placement = Placement(worker_id=0, host="localhost", gpu_id=3) + result = _valid_shard_result(shard, placement, backend=backend) + provenance = result["provenance"] + assert isinstance(provenance, dict) + provenance["hostname"] = "actual-node.example.com" + + audit = _audit_shard_results( + (shard,), + [result], + [{"item_id": "zero", "status": "success"}], + placements=(placement,), + allow_loopback_alias=True, + backend=backend, + ) + + assert audit["ok"] is True + + +@pytest.mark.parametrize("backend", ["cupy", "numba-cuda", "cutile"]) +@pytest.mark.parametrize( + ("uuid", "pci_bus_id", "identity_error"), + [ + (None, None, None), + ("GPU-test", "0000:03:00.0", "identity lookup failed"), + ], +) +def test_shard_result_audit_requires_stable_gpu_identity( + backend, uuid, pci_bus_id, identity_error +) -> None: + shard = (WorkItem("zero", {}, 10),) + placement = Placement(worker_id=0, host="node", gpu_id=3) + result = _valid_shard_result(shard, placement, backend=backend) + provenance = result["provenance"] + assert isinstance(provenance, dict) + gpu = provenance["gpu"] + assert isinstance(gpu, dict) + gpu.update( + uuid=uuid, + pci_bus_id=pci_bus_id, + identity_error=identity_error, + ) + + audit = _audit_shard_results( + (shard,), + [result], + [{"item_id": "zero", "status": "success"}], + placements=(placement,), + allow_loopback_alias=False, + backend=backend, + ) + + assert audit["ok"] is False + assert audit["mismatched_shards"] == [ + {"worker_id": 0, "fields": ["provenance"]} + ] + + +@pytest.mark.parametrize("backend", ["cupy", "numba-cuda", "cutile"]) +@pytest.mark.parametrize("identity_field", ["uuid", "pci_bus_id"]) +def test_shard_result_audit_rejects_duplicate_physical_gpu_identity( + backend, identity_field +) -> None: + shards = ( + (WorkItem("zero", {}, 10),), + (WorkItem("one", {}, 20),), + ) + placements = ( + Placement(worker_id=0, host="node.example.com", gpu_id=3), + Placement(worker_id=1, host="node.example.com", gpu_id=7), + ) + results = [ + _valid_shard_result(shard, placement, backend=backend) + for shard, placement in zip(shards, placements) + ] + for result in results: + provenance = result["provenance"] + assert isinstance(provenance, dict) + gpu = provenance["gpu"] + assert isinstance(gpu, dict) + gpu[identity_field] = "shared-physical-id" + + audit = _audit_shard_results( + shards, + results, + [ + {"item_id": "zero", "status": "success"}, + {"item_id": "one", "status": "success"}, + ], + placements=placements, + allow_loopback_alias=False, + backend=backend, + ) + + assert audit["ok"] is False + assert audit["duplicate_physical_gpu_worker_ids"] == [0, 1] + assert audit["mismatched_shards"] == [ + {"worker_id": 0, "fields": ["provenance"]}, + {"worker_id": 1, "fields": ["provenance"]}, + ] + + +def test_shard_continues_after_terminal_record_write_failure( + monkeypatch, tmp_path +) -> None: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "5") + run_dir = tmp_path / "run" + for name in ("items", "records", "workers"): + (run_dir / name).mkdir(parents=True, exist_ok=True) + original_atomic_write_json = dragon_module.atomic_write_json + + def flaky_atomic_write_json(path, payload): + if path.parent.name == "records" and path.name == "first.json": + raise OSError("injected record write failure") + original_atomic_write_json(path, payload) + + monkeypatch.setattr( + dragon_module, "atomic_write_json", flaky_atomic_write_json + ) + visited = [] + + def item_runner(item, output_dir, options): + del options + visited.append(item.item_id) + return { + "run_dir": str(output_dir), + "summary_path": str(output_dir / "summary.json"), + } + + result = _execute_shard( + run_id="record-write-failure", + run_dir=run_dir, + placement=Placement( + worker_id=0, + host=dragon_module.socket.gethostname(), + gpu_id=5, + ), + items=( + WorkItem("first", {}, 10), + WorkItem("second", {}, 20), + ), + options=_options(), + item_runner=item_runner, + gpu_identity_loader=lambda backend: {"backend": backend}, + require_clean_cuda_imports=False, + ) + + assert visited == ["first", "second"] + assert result["status"] == "failed" + assert result["success_count"] == 1 + assert result["failed_count"] == 1 + assert result["record_write_errors"] == [ + { + "item_id": "first", + "type": "OSError", + "message": "injected record write failure", + } + ] + assert not (run_dir / "records" / "first.json").exists() + assert (run_dir / "records" / "second.json").is_file() + assert (run_dir / "workers" / "worker-0000.json").is_file() + + +def test_real_item_record_matches_coordinator_contract( + monkeypatch, tmp_path +) -> None: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") + reference = np.zeros((64, 64), dtype=np.float64) + reference[20:40, 20:40] = 1.0 + target = reference.copy() + reference_path = tmp_path / "reference.npy" + target_path = tmp_path / "target.npy" + np.save(reference_path, reference, allow_pickle=False) + np.save(target_path, target, allow_pickle=False) + manifest_path = tmp_path / "pairs.json" + manifest_path.write_text( + json.dumps( + { + "schema": "cuphoton.xpois.image-pairs/v1", + "pairs": [ + { + "id": "cpu-pair", + "reference": str(reference_path), + "target": str(target_path), + } + ], + } + ) + ) + item = load_image_pair_manifest(manifest_path).work_items()[0] + run_dir = tmp_path / "run" + for name in ("items", "records", "workers"): + (run_dir / name).mkdir(parents=True, exist_ok=True) + options = BatchFitOptions( + kernel_shape=(9, 9), + basis_sigmas=(1.5,), + basis_degrees=(0,), + backend="cpu", + ) + + result = _execute_shard( + run_id="real-item-contract", + run_dir=run_dir, + placement=Placement( + worker_id=0, + host=dragon_module.socket.gethostname(), + gpu_id=0, + ), + items=(item,), + options=options, + item_runner=run_image_pair_item, + gpu_identity_loader=lambda backend: {"backend": backend}, + require_clean_cuda_imports=False, + ) + records, load_errors = _load_terminal_records(run_dir) + + assert result["status"] == "success" + assert load_errors == [] + assert ( + _audit_terminal_record_contract( + run_id="real-item-contract", + expected={ + item.item_id: { + "worker_id": 0, + "weight_bytes": item.weight_bytes, + } + }, + records=records, + ) + == [] + ) + + +def test_item_timing_aggregation_quarantines_malformed_mapping() -> None: + timings, errors = _aggregate_item_timings( + [ + {"item_id": "good", "timings_sec": {"solve_sec": 1.0}}, + {"item_id": "bad", "timings_sec": "not-a-mapping"}, + ] + ) + + assert timings["solve_sec"] == {"sum": 1.0, "max": 1.0, "mean": 1.0} + assert errors == [ + { + "item_id": "bad", + "type": "InvalidTimingRecord", + "message": "item 'bad' timings_sec must be a mapping", + } + ] + + +def test_terminal_record_contract_rejects_incomplete_success() -> None: + errors = _audit_terminal_record_contract( + run_id="run", + expected={"item": {"worker_id": 0, "weight_bytes": 10}}, + records=[{"item_id": "item", "status": "success"}], + ) + + assert len(errors) == 1 + assert errors[0]["type"] == "InvalidTerminalRecord" + assert "schema" in errors[0]["message"] + assert "worker_id" in errors[0]["message"] + assert "summary_path" in errors[0]["message"] + assert "device" in errors[0]["message"] + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("started_at_utc", "2026-08-21T00:00:00"), + ("completed_at_utc", "not-a-timestamp"), + ], +) +def test_terminal_record_contract_rejects_invalid_timestamps( + field, value +) -> None: + record = { + "schema": "cuphoton.xpois.dragon-item/v1", + "run_id": "run", + "item_id": "item", + "worker_id": 0, + "weight_bytes": 10, + "started_at_utc": "2026-08-21T00:00:00+00:00", + "completed_at_utc": "2026-08-21T00:00:01+00:00", + "worker_seconds": 1.0, + "status": "failed", + "error": {"type": "RuntimeError", "message": "failed"}, + } + record[field] = value + + errors = _audit_terminal_record_contract( + run_id="run", + expected={"item": {"worker_id": 0, "weight_bytes": 10}}, + records=[record], + ) + + assert len(errors) == 1 + assert errors[0]["message"].endswith(field) + + +def test_terminal_record_contract_rejects_decreasing_timestamps() -> None: + errors = _audit_terminal_record_contract( + run_id="run", + expected={"item": {"worker_id": 0, "weight_bytes": 10}}, + records=[ + { + "schema": "cuphoton.xpois.dragon-item/v1", + "run_id": "run", + "item_id": "item", + "worker_id": 0, + "weight_bytes": 10, + "started_at_utc": "2026-08-21T00:00:01+00:00", + "completed_at_utc": "2026-08-21T00:00:00+00:00", + "worker_seconds": 1.0, + "status": "failed", + "error": {"type": "RuntimeError", "message": "failed"}, + } + ], + ) + + assert len(errors) == 1 + assert errors[0]["message"].endswith("completed_at_utc") + + +@pytest.mark.parametrize( + "wall_sec", + [ + {"fit": True}, + {"fit": -1.0}, + {"fit": float("inf")}, + {"fit": "one"}, + ], +) +def test_terminal_record_contract_validates_wall_timings(wall_sec) -> None: + errors = _audit_terminal_record_contract( + run_id="run", + expected={ + "item": { + "worker_id": 0, + "weight_bytes": 10, + "backend": "cupy", + } + }, + records=[ + { + "schema": "cuphoton.xpois.dragon-item/v1", + "run_id": "run", + "item_id": "item", + "worker_id": 0, + "weight_bytes": 10, + "started_at_utc": "2026-08-21T00:00:00+00:00", + "completed_at_utc": "2026-08-21T00:00:01+00:00", + "worker_seconds": 1.0, + "status": "success", + "run_dir": "items/item", + "summary_path": "items/item/summary.json", + "requested_backend": "cupy", + "backend": "cupy", + "device": "fake-gpu", + "runtime": {}, + "timings_sec": {}, + "wall_sec": wall_sec, + } + ], + ) + + assert len(errors) == 1 + assert errors[0]["message"].endswith("wall_sec") + + +def test_terminal_loader_exposes_unexpected_record_files(tmp_path) -> None: + records_dir = tmp_path / "records" + records_dir.mkdir() + atomic_write_json( + records_dir / "expected.json", + {"item_id": "expected", "status": "success"}, + ) + atomic_write_json( + records_dir / "unexpected.json", + {"item_id": "unexpected", "status": "success"}, + ) + + records, errors = _load_terminal_records(tmp_path) + audit = audit_terminal_records(["expected"], records) + + assert errors == [] + assert audit["ok"] is False + assert audit["unexpected_item_ids"] == ["unexpected"] + + +class _FakeQueue(queue.Queue): + def close(self): + return None + + +class _FakePolicy: + Placement = SimpleNamespace(HOST_NAME="host-name") + + def __init__(self, *, placement, host_name, gpu_affinity): + self.placement = placement + self.host_name = host_name + self.gpu_affinity = gpu_affinity + + +class _FakeTemplate: + def __init__(self, *, target, args, policy): + self.target = target + self.args = args + self.policy = policy + + +class _FakeGroup: + last_walltime = None + last_join_timeout = None + last_closed = False + last_stopped = False + + def __init__(self, *, restart, ignore_error_on_exit, walltime): + assert restart is False + assert ignore_error_on_exit is False + type(self).last_walltime = walltime + type(self).last_join_timeout = None + type(self).last_closed = False + type(self).last_stopped = False + self.templates = [] + self.exit_status = [] + + def add_process(self, *, nproc, template): + assert nproc == 1 + self.templates.append(template) + + def init(self): + return None + + def start(self): + 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=None): + type(self).last_join_timeout = timeout + return None + + def stop(self, patience=5.0): + assert patience == 5.0 + type(self).last_stopped = True + return None + + def close(self, patience=5.0): + assert patience == 5.0 + type(self).last_closed = True + return None + + @property + def inactive_puids(self): + return self.exit_status + + +class _FakeSystem: + nodes = (1,) + + +class _FakeNode: + hostname = "fake-node" + gpus = [3, 7] + + def __init__(self, node_id): + assert node_id == 1 + + +class _FakeFailedGroup(_FakeGroup): + def start(self): + self.exit_status = [(1000, -9)] + + +class _FakeSetupFailedGroup(_FakeGroup): + def add_process(self, *, nproc, template): + del nproc, template + raise RuntimeError("process setup failed") + + +class _FakeJoinTimedOutGroup(_FakeGroup): + def start(self): + return None + + def join(self, timeout=None): + type(self).last_join_timeout = timeout + raise TimeoutError("worker join deadline exceeded") + + +class _FakeStartFailedGroup(_FakeGroup): + def start(self): + raise RuntimeError("partial worker start failed") + + +class _FakeDecoratedJoinFailedGroup(_FakeGroup): + last_forced_closed = False + + def __init__(self, **kwargs): + super().__init__(**kwargs) + type(self).last_forced_closed = False + + def start(self): + return None + + def join(self, timeout=None): + type(self).last_join_timeout = timeout + raise RuntimeError("decorated worker failure") + + def stop(self, patience=5.0): + assert patience == 5.0 + type(self).last_stopped = True + raise RuntimeError("decorated worker failure") + + def close(self, patience=5.0): + assert patience == 5.0 + type(self).last_closed = True + raise RuntimeError("decorated worker failure") + + def _close_no_decorator(self, patience=5.0): + assert patience == 5.0 + type(self).last_forced_closed = True + + +def _fake_worker( + run_id, + run_dir_raw, + placement_payload, + item_payloads, + options_payload, + results_queue, + allow_loopback_alias, +): + del allow_loopback_alias + run_dir = dragon_module.Path(run_dir_raw) + worker_id = placement_payload["worker_id"] + requested_host = placement_payload["host"] + requested_gpu_id = placement_payload["gpu_id"] + backend = options_payload["backend"] + identity_backend = "cupy" if backend in {"cupy", "cutile"} else backend + for payload in item_payloads: + item = WorkItem.from_dict(payload) + atomic_write_json( + run_dir / "records" / f"{item.item_id}.json", + { + "schema": "cuphoton.xpois.dragon-item/v1", + "run_id": run_id, + "item_id": item.item_id, + "worker_id": worker_id, + "weight_bytes": item.weight_bytes, + "started_at_utc": "2026-08-21T00:00:00+00:00", + "completed_at_utc": "2026-08-21T00:00:01+00:00", + "worker_seconds": 1.0, + "status": "success", + "run_dir": f"items/{item.item_id}", + "summary_path": f"items/{item.item_id}/summary.json", + "requested_backend": "cupy", + "backend": "cupy", + "device": "fake-gpu", + "runtime": {}, + "timings_sec": {"solve": 0.5}, + "wall_sec": {"item_runner": 0.75}, + }, + ) + results_queue.put( + { + "schema": "cuphoton.xpois.dragon-shard/v1", + "worker_id": worker_id, + "status": "success", + "item_count": len(item_payloads), + "success_count": len(item_payloads), + "failed_count": 0, + "weight_bytes": sum( + WorkItem.from_dict(payload).weight_bytes + for payload in item_payloads + ), + "item_ids_sha256": dragon_module._item_ids_sha256( + [WorkItem.from_dict(payload) for payload in item_payloads] + ), + "started_at_utc": "2026-08-21T00:00:00+00:00", + "completed_at_utc": "2026-08-21T00:00:01+00:00", + "worker_wall_sec": 1.0, + "timings_sec": {"solve": 0.5 * len(item_payloads)}, + "provenance": { + "worker_id": worker_id, + "requested_host": requested_host, + "requested_gpu_id": requested_gpu_id, + "hostname": requested_host, + "pid": 1234, + "cuda_visible_devices": str(requested_gpu_id), + "gpu": { + "backend": identity_backend, + "name": "fake-gpu", + "uuid": f"GPU-worker-{worker_id}", + "pci_bus_id": f"0000:{requested_gpu_id:02x}:00.0", + "identity_error": None, + }, + }, + "record_write_errors": [], + } + ) + + +def _fake_corrupt_worker( + run_id, + run_dir_raw, + placement_payload, + item_payloads, + options_payload, + results_queue, + allow_loopback_alias, +): + local_results = _FakeQueue() + _fake_worker( + run_id, + run_dir_raw, + placement_payload, + item_payloads, + options_payload, + local_results, + allow_loopback_alias, + ) + local_results.get_nowait() + for payload in item_payloads: + item = WorkItem.from_dict(payload) + record_path = ( + dragon_module.Path(run_dir_raw) + / "records" + / f"{item.item_id}.json" + ) + record = json.loads(record_path.read_text()) + record["timings_sec"] = "corrupt" + atomic_write_json(record_path, record) + results_queue.put(["corrupt"]) + + +def _fake_inconsistent_shard_worker( + run_id, + run_dir_raw, + placement_payload, + item_payloads, + options_payload, + results_queue, + allow_loopback_alias, +): + local_results = _FakeQueue() + _fake_worker( + run_id, + run_dir_raw, + placement_payload, + item_payloads, + options_payload, + local_results, + allow_loopback_alias, + ) + result = local_results.get_nowait() + result["status"] = "failed" + result["success_count"] = 0 + result["failed_count"] = len(item_payloads) + results_queue.put(result) + + +def _write_fake_manifest(tmp_path, *, item_ids=("one",)): + reference = tmp_path / "reference.npy" + target = tmp_path / "target.npy" + np.save(reference, np.zeros((16, 16)), allow_pickle=False) + np.save(target, np.zeros((16, 16)), allow_pickle=False) + manifest = tmp_path / "pairs.json" + manifest.write_text( + json.dumps( + { + "schema": "cuphoton.xpois.image-pairs/v1", + "pairs": [ + { + "id": item_id, + "reference": str(reference), + "target": str(target), + } + for item_id in item_ids + ], + } + ) + ) + return manifest + + +def test_dragon_shard_worker_round_trips_wire_payloads( + monkeypatch, tmp_path +) -> None: + seen = {} + + def fake_execute_shard(**kwargs): + seen.update(kwargs) + return { + "worker_id": kwargs["placement"].worker_id, + "status": "success", + } + + monkeypatch.setattr(dragon_module, "_execute_shard", fake_execute_shard) + results = _FakeQueue() + placement = Placement(worker_id=3, host="node", gpu_id=7) + item = WorkItem("wire-item", {"id": "wire-item"}, 42) + + _dragon_shard_worker( + "wire-run", + str(tmp_path), + placement.to_dict(), + [item.to_dict()], + _options().to_payload(), + results, + True, + ) + + assert results.get_nowait() == {"worker_id": 3, "status": "success"} + assert seen["run_id"] == "wire-run" + assert seen["placement"] == placement + assert seen["items"] == (item,) + assert seen["options"] == _options() + assert seen["allow_loopback_alias"] is True + + +def test_coordinator_audits_fake_dragon_process_group( + monkeypatch, tmp_path +) -> None: + manifest = _write_fake_manifest(tmp_path) + api = _DragonAPI( + System=_FakeSystem, + Node=_FakeNode, + Policy=_FakePolicy, + ProcessGroup=_FakeGroup, + ProcessTemplate=_FakeTemplate, + Queue=_FakeQueue, + ) + monkeypatch.setattr(dragon_module, "_load_dragon_api", lambda: api) + monkeypatch.setattr(dragon_module, "_dragon_shard_worker", _fake_worker) + + result = run_dragon_image_pair_batch( + manifest_path=manifest, + output_root=tmp_path / "runs", + run_id="fake-dragon", + max_workers=2, + result_timeout_sec=1.0, + worker_timeout_sec=30.0, + options=_options(), + ) + + assert result.status == "success" + assert result.summary["terminal_record_audit"]["ok"] is True + assert result.summary["process_exit_audit"]["ok"] is True + assert result.summary["shard_result_audit"]["ok"] is True + assert result.summary["placements"] == [ + {"worker_id": 0, "host": "fake-node", "gpu_id": 3} + ] + assert result.summary["distinct_host_count"] == 1 + assert result.summary["coordinator_wall_sec"] >= 0 + coordinator_timings = result.summary["coordinator_timings_sec"] + assert "coordinator_setup_sec" not in coordinator_timings + assert result.summary["coordinator_totals_sec"]["setup"] >= 0 + assert _FakeGroup.last_walltime == 30.0 + assert _FakeGroup.last_join_timeout == 31.0 + assert _FakeGroup.last_stopped is False + assert _FakeGroup.last_closed is True + launch = json.loads((result.run_dir / "run.json").read_text()) + assert launch["record_type"] == "immutable-launch" + assert "status" not in launch + assert (result.run_dir / "input-identity.json").is_file() + + +def test_coordinator_audits_two_worker_fanout(monkeypatch, tmp_path) -> None: + manifest = _write_fake_manifest( + tmp_path, item_ids=("one", "two", "three") + ) + api = _DragonAPI( + System=_FakeSystem, + Node=_FakeNode, + Policy=_FakePolicy, + ProcessGroup=_FakeGroup, + ProcessTemplate=_FakeTemplate, + Queue=_FakeQueue, + ) + monkeypatch.setattr(dragon_module, "_load_dragon_api", lambda: api) + monkeypatch.setattr(dragon_module, "_dragon_shard_worker", _fake_worker) + + result = run_dragon_image_pair_batch( + manifest_path=manifest, + output_root=tmp_path / "multi-runs", + run_id="fake-dragon-multi", + max_workers=2, + result_timeout_sec=1.0, + worker_timeout_sec=30.0, + options=_options(), + ) + + assert result.status == "success" + assert result.summary["placements"] == [ + {"worker_id": 0, "host": "fake-node", "gpu_id": 3}, + {"worker_id": 1, "host": "fake-node", "gpu_id": 7}, + ] + assert result.summary["queue_result_count"] == 2 + assert result.summary["terminal_record_audit"]["expected_count"] == 3 + assert result.summary["terminal_record_audit"]["ok"] is True + assert result.summary["process_exit_audit"] == { + "ok": True, + "expected_count": 2, + "observed_count": 2, + "nonzero": [], + } + assert result.summary["shard_result_audit"]["ok"] is True + expected_shards = { + shard["worker_id"]: shard for shard in result.summary["shards"] + } + observed_shards = { + shard["worker_id"]: shard for shard in result.summary["shard_results"] + } + assert set(expected_shards) == set(observed_shards) == {0, 1} + assert sum(shard["item_count"] for shard in expected_shards.values()) == 3 + for worker_id, expected in expected_shards.items(): + observed = observed_shards[worker_id] + assert observed["item_count"] == expected["item_count"] + assert observed["weight_bytes"] == expected["weight_bytes"] + assert observed["item_ids_sha256"] == expected["item_ids_sha256"] + + +def test_coordinator_rejects_failed_shard_with_successful_records( + monkeypatch, tmp_path +) -> None: + manifest = _write_fake_manifest(tmp_path) + api = _DragonAPI( + System=_FakeSystem, + Node=_FakeNode, + Policy=_FakePolicy, + ProcessGroup=_FakeGroup, + ProcessTemplate=_FakeTemplate, + Queue=_FakeQueue, + ) + monkeypatch.setattr(dragon_module, "_load_dragon_api", lambda: api) + monkeypatch.setattr( + dragon_module, + "_dragon_shard_worker", + _fake_inconsistent_shard_worker, + ) + + result = run_dragon_image_pair_batch( + manifest_path=manifest, + output_root=tmp_path / "failed-shard-runs", + run_id="fake-failed-shard", + max_workers=1, + result_timeout_sec=1.0, + worker_timeout_sec=30.0, + options=_options(), + ) + + assert result.status == "failed" + assert result.summary["terminal_record_audit"]["ok"] is True + assert result.summary["shard_result_audit"]["mismatched_shards"] == [ + { + "worker_id": 0, + "fields": ["failed_count", "status", "success_count"], + } + ] + + +def test_coordinator_finalizes_corrupt_worker_payloads( + monkeypatch, tmp_path +) -> None: + manifest = _write_fake_manifest(tmp_path) + api = _DragonAPI( + System=_FakeSystem, + Node=_FakeNode, + Policy=_FakePolicy, + ProcessGroup=_FakeGroup, + ProcessTemplate=_FakeTemplate, + Queue=_FakeQueue, + ) + monkeypatch.setattr(dragon_module, "_load_dragon_api", lambda: api) + monkeypatch.setattr( + dragon_module, "_dragon_shard_worker", _fake_corrupt_worker + ) + + result = run_dragon_image_pair_batch( + manifest_path=manifest, + output_root=tmp_path / "corrupt-runs", + run_id="fake-corrupt", + max_workers=1, + result_timeout_sec=1.0, + worker_timeout_sec=30.0, + options=_options(), + ) + + assert result.status == "failed" + assert result.summary_path.is_file() + assert result.summary["shard_result_audit"]["missing_worker_ids"] == [0] + assert result.summary["queue_result_count"] == 1 + assert result.summary["lifecycle_errors"] == [ + { + "phase": "queue_result", + "type": "InvalidShardResult", + "message": "worker result was not a mapping", + } + ] + assert result.summary["terminal_record_errors"] == [ + { + "item_id": "one", + "type": "InvalidTerminalRecord", + "message": "record 0 has invalid field(s): timings_sec", + }, + { + "item_id": "one", + "type": "InvalidTimingRecord", + "message": "item 'one' timings_sec must be a mapping", + }, + ] + + +def test_coordinator_rejects_abrupt_worker_exit( + monkeypatch, tmp_path +) -> None: + manifest = _write_fake_manifest(tmp_path) + api = _DragonAPI( + System=_FakeSystem, + Node=_FakeNode, + Policy=_FakePolicy, + ProcessGroup=_FakeFailedGroup, + ProcessTemplate=_FakeTemplate, + Queue=_FakeQueue, + ) + monkeypatch.setattr(dragon_module, "_load_dragon_api", lambda: api) + + result = run_dragon_image_pair_batch( + manifest_path=manifest, + output_root=tmp_path / "failed-runs", + run_id="fake-worker-killed", + max_workers=1, + result_timeout_sec=0.001, + worker_timeout_sec=30.0, + options=_options(), + ) + + assert result.status == "failed" + assert result.summary["process_exit_audit"]["nonzero"] == [ + {"puid": 1000, "exit_code": -9} + ] + assert result.summary["terminal_record_audit"]["missing_item_ids"] == [ + "one" + ] + assert result.summary["shard_result_audit"]["missing_worker_ids"] == [0] + + +def test_coordinator_finalizes_process_setup_failure( + monkeypatch, tmp_path +) -> None: + manifest = _write_fake_manifest(tmp_path) + api = _DragonAPI( + System=_FakeSystem, + Node=_FakeNode, + Policy=_FakePolicy, + ProcessGroup=_FakeSetupFailedGroup, + ProcessTemplate=_FakeTemplate, + Queue=_FakeQueue, + ) + monkeypatch.setattr(dragon_module, "_load_dragon_api", lambda: api) + + result = run_dragon_image_pair_batch( + manifest_path=manifest, + output_root=tmp_path / "setup-failed-runs", + run_id="fake-setup-failed", + max_workers=1, + result_timeout_sec=1.0, + worker_timeout_sec=30.0, + options=_options(), + ) + + assert result.status == "failed" + assert result.summary_path.is_file() + assert result.summary["lifecycle_errors"] == [ + { + "phase": "process_setup", + "type": "RuntimeError", + "message": "process setup failed", + } + ] + assert result.summary["terminal_record_audit"]["missing_item_ids"] == [ + "one" + ] + + +def test_coordinator_finalizes_join_timeout(monkeypatch, tmp_path) -> None: + manifest = _write_fake_manifest(tmp_path) + api = _DragonAPI( + System=_FakeSystem, + Node=_FakeNode, + Policy=_FakePolicy, + ProcessGroup=_FakeJoinTimedOutGroup, + ProcessTemplate=_FakeTemplate, + Queue=_FakeQueue, + ) + monkeypatch.setattr(dragon_module, "_load_dragon_api", lambda: api) + + result = run_dragon_image_pair_batch( + manifest_path=manifest, + output_root=tmp_path / "join-timeout-runs", + run_id="fake-join-timeout", + max_workers=1, + result_timeout_sec=0.001, + worker_timeout_sec=0.01, + options=_options(), + ) + + assert result.status == "failed" + assert result.summary_path.is_file() + assert result.summary["lifecycle_errors"] == [ + { + "phase": "join_timeout", + "type": "TimeoutError", + "message": "worker join deadline exceeded", + } + ] + assert _FakeJoinTimedOutGroup.last_join_timeout == pytest.approx(0.011) + assert _FakeJoinTimedOutGroup.last_stopped is True + assert _FakeJoinTimedOutGroup.last_closed is True + + +def test_coordinator_stops_after_partial_start_failure( + monkeypatch, tmp_path +) -> None: + manifest = _write_fake_manifest(tmp_path) + api = _DragonAPI( + System=_FakeSystem, + Node=_FakeNode, + Policy=_FakePolicy, + ProcessGroup=_FakeStartFailedGroup, + ProcessTemplate=_FakeTemplate, + Queue=_FakeQueue, + ) + monkeypatch.setattr(dragon_module, "_load_dragon_api", lambda: api) + + result = run_dragon_image_pair_batch( + manifest_path=manifest, + output_root=tmp_path / "start-failed-runs", + run_id="fake-start-failed", + max_workers=1, + result_timeout_sec=0.001, + worker_timeout_sec=30.0, + options=_options(), + ) + + assert result.status == "failed" + assert result.summary["lifecycle_errors"] == [ + { + "phase": "start", + "type": "RuntimeError", + "message": "partial worker start failed", + } + ] + assert _FakeStartFailedGroup.last_stopped is True + assert _FakeStartFailedGroup.last_closed is True + + +def test_coordinator_preserves_join_and_cleanup_errors( + monkeypatch, tmp_path +) -> None: + manifest = _write_fake_manifest(tmp_path) + api = _DragonAPI( + System=_FakeSystem, + Node=_FakeNode, + Policy=_FakePolicy, + ProcessGroup=_FakeDecoratedJoinFailedGroup, + ProcessTemplate=_FakeTemplate, + Queue=_FakeQueue, + ) + monkeypatch.setattr(dragon_module, "_load_dragon_api", lambda: api) + + result = run_dragon_image_pair_batch( + manifest_path=manifest, + output_root=tmp_path / "join-failed-runs", + run_id="fake-join-failed", + max_workers=1, + result_timeout_sec=0.001, + worker_timeout_sec=30.0, + options=_options(), + ) + + assert result.status == "failed" + assert result.summary["lifecycle_errors"] == [ + { + "phase": "join", + "type": "RuntimeError", + "message": "decorated worker failure", + }, + { + "phase": "stop_after_failure", + "type": "RuntimeError", + "message": "decorated worker failure", + }, + { + "phase": "close", + "type": "RuntimeError", + "message": "decorated worker failure", + }, + ] + assert _FakeDecoratedJoinFailedGroup.last_stopped is True + assert _FakeDecoratedJoinFailedGroup.last_closed is True + assert _FakeDecoratedJoinFailedGroup.last_forced_closed is True + + +@pytest.mark.parametrize("backend", ["auto", "cpu"]) +def test_coordinator_requires_explicit_gpu_backend( + monkeypatch, tmp_path, backend +) -> None: + monkeypatch.setattr( + dragon_module, + "_load_dragon_api", + lambda: pytest.fail("Dragon must not load before backend validation"), + ) + + with pytest.raises(ValueError, match="explicit GPU backend"): + run_dragon_image_pair_batch( + manifest_path=tmp_path / "missing.json", + output_root=tmp_path / "runs", + run_id="invalid-backend", + max_workers=1, + result_timeout_sec=1.0, + worker_timeout_sec=30.0, + options=_options(backend=backend), + ) From ecbdc4cb0060a57b1d06a9a154f6ea8e30d7b25f Mon Sep 17 00:00:00 2001 From: Carmelo Gonzales Date: Tue, 22 Sep 2026 11:02:39 -0600 Subject: [PATCH 2/4] Separate record write failures from shard evidence mismatches Signed-off-by: Carmelo Gonzales --- docs/components/xpois.md | 5 +- src/cuphoton/xpois/dragon.py | 22 +++---- tests/xpois/test_dragon.py | 113 ++++++++++++++++++++++++++++++++--- 3 files changed, 121 insertions(+), 19 deletions(-) diff --git a/docs/components/xpois.md b/docs/components/xpois.md index 69c625d..4f30ea9 100644 --- a/docs/components/xpois.md +++ b/docs/components/xpois.md @@ -222,7 +222,10 @@ Every attempt uses a new, immutable run directory. Each item has an atomic terminal record under `records/`; an ordinary item error does not prevent the remaining items in its shard from running. The final `summary.json` checks for missing, duplicate, unexpected, malformed, or assignment-inconsistent -item and worker results, as well as nonzero worker exits. Worker wall time +item and worker results, as well as nonzero worker exits. A worker that could +not write a terminal record still fails the run, but its declared errors are +reported under `shard_result_audit.write_failed_shards` rather than as an +evidence mismatch. Worker wall time defaults to one hour. The coordinator attempts bounded stop and close cleanup after failed starts or joins. diff --git a/src/cuphoton/xpois/dragon.py b/src/cuphoton/xpois/dragon.py index 6c70156..638e446 100644 --- a/src/cuphoton/xpois/dragon.py +++ b/src/cuphoton/xpois/dragon.py @@ -748,12 +748,8 @@ def _hostnames_match( def _collect_gpu_identity(backend: str) -> dict[str, Any]: """Initialize the selected backend only after singleton placement.""" - if backend in {"auto", "cupy", "cutile"}: - try: - return _collect_cupy_identity() - except (ImportError, OSError, RuntimeError): - if backend != "auto": - raise + if backend in {"cupy", "cutile"}: + return _collect_cupy_identity() return _collect_numba_identity() @@ -972,6 +968,7 @@ def _audit_shard_results( record.get("status") ) mismatched: list[dict[str, Any]] = [] + write_failed: list[dict[str, Any]] = [] physical_identities: list[ tuple[int, str, frozenset[tuple[str, str]]] ] = [] @@ -1054,10 +1051,6 @@ def _audit_shard_results( fields.add("success_count") if failed_count is not None and failed_count != terminal_failed_count: fields.add("failed_count") - if status != "success": - fields.add("status") - if failed_count != 0: - fields.add("failed_count") placement = placement_by_worker.get(worker_id) provenance = result.get("provenance") if placement is None or not _valid_shard_provenance( @@ -1082,7 +1075,12 @@ def _audit_shard_results( ): fields.add("record_write_errors") elif record_write_errors: - fields.add("record_write_errors") + write_failed.append( + { + "worker_id": worker_id, + "record_write_errors": record_write_errors, + } + ) try: _validated_timings( result.get("timings_sec"), @@ -1127,6 +1125,7 @@ def _audit_shard_results( and not duplicates and not unexpected and not mismatched + and not write_failed and not invalid_results, "expected_count": worker_count, "observed_count": len(shard_results), @@ -1137,6 +1136,7 @@ def _audit_shard_results( duplicate_physical_gpu_worker_ids ), "mismatched_shards": mismatched, + "write_failed_shards": write_failed, "invalid_results": invalid_results, } diff --git a/tests/xpois/test_dragon.py b/tests/xpois/test_dragon.py index 32f99ee..aaa9eb9 100644 --- a/tests/xpois/test_dragon.py +++ b/tests/xpois/test_dragon.py @@ -424,7 +424,7 @@ def test_shard_result_audit_requires_strict_integral_totals( assert audit["mismatched_shards"] == [{"worker_id": 0, "fields": [field]}] -def test_shard_result_audit_rejects_reported_failure() -> None: +def test_shard_result_audit_accepts_honest_reported_failure() -> None: shard = (WorkItem("zero", {}, 10),) placements = _test_placements(1) result = _valid_shard_result(shard, placements[0]) @@ -439,10 +439,9 @@ def test_shard_result_audit_rejects_reported_failure() -> None: backend="cupy", ) - assert audit["ok"] is False - assert audit["mismatched_shards"] == [ - {"worker_id": 0, "fields": ["failed_count", "status"]} - ] + assert audit["ok"] is True + assert audit["mismatched_shards"] == [] + assert audit["write_failed_shards"] == [] def test_shard_result_audit_cross_checks_terminal_status_counts() -> None: @@ -465,7 +464,7 @@ def test_shard_result_audit_cross_checks_terminal_status_counts() -> None: ] -def test_shard_result_audit_rejects_record_write_errors() -> None: +def test_shard_result_audit_separates_record_write_errors() -> None: shard = (WorkItem("zero", {}, 10),) placements = _test_placements(1) result = _valid_shard_result(shard, placements[0]) @@ -481,6 +480,32 @@ def test_shard_result_audit_rejects_record_write_errors() -> None: ) assert audit["ok"] is False + assert audit["mismatched_shards"] == [] + assert audit["write_failed_shards"] == [ + {"worker_id": 0, "record_write_errors": [{"item_id": "zero"}]} + ] + + +@pytest.mark.parametrize("value", [None, {"item_id": "zero"}, ["zero"]]) +def test_shard_result_audit_rejects_malformed_record_write_errors( + value, +) -> None: + shard = (WorkItem("zero", {}, 10),) + placements = _test_placements(1) + result = _valid_shard_result(shard, placements[0]) + result["record_write_errors"] = value + + audit = _audit_shard_results( + (shard,), + [result], + [{"item_id": "zero", "status": "success"}], + placements=placements, + allow_loopback_alias=False, + backend="cupy", + ) + + assert audit["ok"] is False + assert audit["write_failed_shards"] == [] assert audit["mismatched_shards"] == [ {"worker_id": 0, "fields": ["record_write_errors"]} ] @@ -1247,6 +1272,36 @@ def _fake_inconsistent_shard_worker( results_queue.put(result) +def _fake_record_write_failure_shard_worker( + run_id, + run_dir_raw, + placement_payload, + item_payloads, + options_payload, + results_queue, + allow_loopback_alias, +): + local_results = _FakeQueue() + _fake_worker( + run_id, + run_dir_raw, + placement_payload, + item_payloads, + options_payload, + local_results, + allow_loopback_alias, + ) + result = local_results.get_nowait() + result["record_write_errors"] = [ + { + "item_id": item_payloads[0]["item_id"], + "type": "OSError", + "message": "injected record write failure", + } + ] + results_queue.put(result) + + def _write_fake_manifest(tmp_path, *, item_ids=("one",)): reference = tmp_path / "reference.npy" target = tmp_path / "target.npy" @@ -1440,9 +1495,53 @@ def test_coordinator_rejects_failed_shard_with_successful_records( assert result.status == "failed" assert result.summary["terminal_record_audit"]["ok"] is True assert result.summary["shard_result_audit"]["mismatched_shards"] == [ + {"worker_id": 0, "fields": ["failed_count", "success_count"]} + ] + + +def test_coordinator_rejects_declared_record_write_errors( + monkeypatch, tmp_path +) -> None: + manifest = _write_fake_manifest(tmp_path) + api = _DragonAPI( + System=_FakeSystem, + Node=_FakeNode, + Policy=_FakePolicy, + ProcessGroup=_FakeGroup, + ProcessTemplate=_FakeTemplate, + Queue=_FakeQueue, + ) + monkeypatch.setattr(dragon_module, "_load_dragon_api", lambda: api) + monkeypatch.setattr( + dragon_module, + "_dragon_shard_worker", + _fake_record_write_failure_shard_worker, + ) + + result = run_dragon_image_pair_batch( + manifest_path=manifest, + output_root=tmp_path / "write-failed-shard-runs", + run_id="fake-write-failed-shard", + max_workers=1, + result_timeout_sec=1.0, + worker_timeout_sec=30.0, + options=_options(), + ) + + assert result.status == "failed" + shard_audit = result.summary["shard_result_audit"] + assert shard_audit["ok"] is False + assert shard_audit["mismatched_shards"] == [] + assert shard_audit["write_failed_shards"] == [ { "worker_id": 0, - "fields": ["failed_count", "status", "success_count"], + "record_write_errors": [ + { + "item_id": "one", + "type": "OSError", + "message": "injected record write failure", + } + ], } ] From 319b0b937502d20456137ba105e53e77a14c7045 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 22 Sep 2026 10:13:17 -0700 Subject: [PATCH 3/4] Reconcile declared record-write failures in shard audits Count declared record-write failures against their assigned items, including records visible after a directory fsync failure. Reject malformed, unknown and duplicate error identities without labeling honest I/O failures as contradictory worker counts. Signed-off-by: Trent Nelson --- src/cuphoton/xpois/dragon.py | 46 ++++++++---- tests/xpois/test_dragon.py | 133 ++++++++++++++++++++++++++++++++++- 2 files changed, 164 insertions(+), 15 deletions(-) diff --git a/src/cuphoton/xpois/dragon.py b/src/cuphoton/xpois/dragon.py index 638e446..f42f347 100644 --- a/src/cuphoton/xpois/dragon.py +++ b/src/cuphoton/xpois/dragon.py @@ -1034,14 +1034,46 @@ def _audit_shard_results( expected_status = "success" if failed_count == 0 else "failed" if status != expected_status: fields.add("status") + record_write_errors = result.get("record_write_errors") + write_error_ids: set[str] = set() + write_errors_valid = isinstance(record_write_errors, list) + if write_errors_valid: + shard_item_ids = {item.item_id for item in shard} + for error in record_write_errors: + if ( + not isinstance(error, Mapping) + or not isinstance(error.get("item_id"), str) + or error["item_id"] not in shard_item_ids + or error["item_id"] in write_error_ids + or not isinstance(error.get("type"), str) + or not error["type"] + or not isinstance(error.get("message"), str) + ): + write_errors_valid = False + break + write_error_ids.add(error["item_id"]) + if not write_errors_valid: + fields.add("record_write_errors") + write_error_ids.clear() + elif write_error_ids: + write_failed.append( + { + "worker_id": worker_id, + "record_write_errors": record_write_errors, + } + ) + # Directory fsync can fail after a record becomes visible. The + # worker counts that item as a failed write, not a durable outcome. terminal_success_count = sum( status == "success" for item in shard + if item.item_id not in write_error_ids for status in terminal_statuses.get(item.item_id, ()) ) - terminal_failed_count = sum( + terminal_failed_count = len(write_error_ids) + sum( status == "failed" for item in shard + if item.item_id not in write_error_ids for status in terminal_statuses.get(item.item_id, ()) ) if ( @@ -1069,18 +1101,6 @@ def _audit_shard_results( physical_identities.append( (worker_id, placement.host, physical_identity) ) - record_write_errors = result.get("record_write_errors") - if not isinstance(record_write_errors, list) or any( - not isinstance(error, Mapping) for error in record_write_errors - ): - fields.add("record_write_errors") - elif record_write_errors: - write_failed.append( - { - "worker_id": worker_id, - "record_write_errors": record_write_errors, - } - ) try: _validated_timings( result.get("timings_sec"), diff --git a/tests/xpois/test_dragon.py b/tests/xpois/test_dragon.py index aaa9eb9..b54485d 100644 --- a/tests/xpois/test_dragon.py +++ b/tests/xpois/test_dragon.py @@ -468,7 +468,15 @@ def test_shard_result_audit_separates_record_write_errors() -> None: shard = (WorkItem("zero", {}, 10),) placements = _test_placements(1) result = _valid_shard_result(shard, placements[0]) - result["record_write_errors"] = [{"item_id": "zero"}] + errors = [ + {"item_id": "zero", "type": "OSError", "message": "write failed"} + ] + result.update( + status="failed", + success_count=0, + failed_count=1, + record_write_errors=errors, + ) audit = _audit_shard_results( (shard,), @@ -482,7 +490,7 @@ def test_shard_result_audit_separates_record_write_errors() -> None: assert audit["ok"] is False assert audit["mismatched_shards"] == [] assert audit["write_failed_shards"] == [ - {"worker_id": 0, "record_write_errors": [{"item_id": "zero"}]} + {"worker_id": 0, "record_write_errors": errors} ] @@ -1292,6 +1300,9 @@ def _fake_record_write_failure_shard_worker( allow_loopback_alias, ) result = local_results.get_nowait() + result["status"] = "failed" + result["success_count"] -= 1 + result["failed_count"] += 1 result["record_write_errors"] = [ { "item_id": item_payloads[0]["item_id"], @@ -1809,3 +1820,121 @@ def test_coordinator_requires_explicit_gpu_backend( worker_timeout_sec=30.0, options=_options(backend=backend), ) + + +@pytest.mark.parametrize("visible_status", [None, "success", "failed"]) +def test_shard_result_audit_separates_declared_write_failure( + visible_status, +) -> None: + shard = (WorkItem("first", {}, 10), WorkItem("second", {}, 20)) + placements = _test_placements(1) + result = _valid_shard_result(shard, placements[0]) + errors = [ + { + "item_id": "first", + "type": "OSError", + "message": "[Errno 28] No space left on device", + } + ] + result.update( + status="failed", + success_count=1, + failed_count=1, + record_write_errors=errors, + ) + records = [{"item_id": "second", "status": "success"}] + if visible_status is not None: + records.append({"item_id": "first", "status": visible_status}) + + audit = _audit_shard_results( + (shard,), + [result], + records, + placements=placements, + allow_loopback_alias=False, + backend="cupy", + ) + + assert audit["ok"] is False + assert audit["mismatched_shards"] == [] + assert audit["write_failed_shards"] == [ + {"worker_id": 0, "record_write_errors": errors} + ] + + +@pytest.mark.parametrize( + "errors", + [ + None, + {}, + ["not an error mapping"], + [{"type": "OSError", "message": "write failed"}], + [{"item_id": "other", "type": "OSError", "message": "failed"}], + [{"item_id": "first", "type": "OSError"}], + [ + {"item_id": "first", "type": "OSError", "message": "failed"}, + {"item_id": "first", "type": "OSError", "message": "failed"}, + ], + ], +) +def test_shard_result_audit_rejects_invalid_write_failure_evidence( + errors, +) -> None: + shard = (WorkItem("first", {}, 10),) + placements = _test_placements(1) + result = _valid_shard_result(shard, placements[0]) + result["record_write_errors"] = errors + + audit = _audit_shard_results( + (shard,), + [result], + [{"item_id": "first", "status": "success"}], + placements=placements, + allow_loopback_alias=False, + backend="cupy", + ) + + assert audit["ok"] is False + assert audit["mismatched_shards"] == [ + {"worker_id": 0, "fields": ["record_write_errors"]} + ] + assert audit["write_failed_shards"] == [] + + +@pytest.mark.parametrize( + ("overrides", "mismatch"), + [ + ({"status": "success"}, "status"), + ({"failed_count": 0}, "failed_count"), + ({"success_count": 2, "failed_count": 0}, "success_count"), + ], +) +def test_declared_write_failure_does_not_hide_inconsistent_counts( + overrides, mismatch +) -> None: + shard = (WorkItem("first", {}, 10), WorkItem("second", {}, 20)) + placements = _test_placements(1) + result = _valid_shard_result(shard, placements[0]) + result.update( + status="failed", + success_count=1, + failed_count=1, + record_write_errors=[ + {"item_id": "first", "type": "OSError", "message": "failed"} + ], + ) + result.update(overrides) + + audit = _audit_shard_results( + (shard,), + [result], + [{"item_id": "second", "status": "success"}], + placements=placements, + allow_loopback_alias=False, + backend="cupy", + ) + + assert audit["ok"] is False + assert len(audit["mismatched_shards"]) == 1 + assert mismatch in audit["mismatched_shards"][0]["fields"] + assert len(audit["write_failed_shards"]) == 1 From 72ddd0ef9dcb377c609fd84e69b3504024e25f5c Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 22 Sep 2026 11:37:23 -0700 Subject: [PATCH 4/4] Validate automatic mask options before batch workers start Match the mask builder constraints for stamp size and peak percentile before launching workers. Keep a one-pixel stamp valid when automatic masking is disabled. Signed-off-by: Trent Nelson --- src/cuphoton/xpois/batch.py | 10 ++++++++-- tests/xpois/test_batch.py | 10 ++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/src/cuphoton/xpois/batch.py b/src/cuphoton/xpois/batch.py index b23c5c5..a1f23b0 100644 --- a/src/cuphoton/xpois/batch.py +++ b/src/cuphoton/xpois/batch.py @@ -325,6 +325,11 @@ def __post_init__(self) -> None: or self.auto_stamp_size % 2 == 0 ): raise ValueError("auto_stamp_size must be a positive odd integer") + if self.auto_stamp_mask and self.auto_stamp_size < 3: + raise ValueError( + "auto_stamp_size must be at least 3 " + "when auto_stamp_mask is enabled" + ) if ( isinstance(self.auto_stamp_count, bool) or not isinstance(self.auto_stamp_count, int) @@ -335,10 +340,11 @@ def __post_init__(self) -> None: isinstance(self.auto_peak_percentile, bool) or not isinstance(self.auto_peak_percentile, (int, float)) or not math.isfinite(self.auto_peak_percentile) - or not 0 <= self.auto_peak_percentile <= 100 + or not 0 < self.auto_peak_percentile < 100 ): raise ValueError( - "auto_peak_percentile must be a finite number from 0 to 100" + "auto_peak_percentile must be a finite number " + "strictly between 0 and 100" ) if self.mask_policy not in _MASK_POLICIES: raise ValueError(f"unsupported mask policy: {self.mask_policy}") diff --git a/tests/xpois/test_batch.py b/tests/xpois/test_batch.py index be12eae..bc2482e 100644 --- a/tests/xpois/test_batch.py +++ b/tests/xpois/test_batch.py @@ -202,10 +202,13 @@ def test_batch_options_reject_nonfinite_basis_sigmas(sigma) -> None: [ ({"auto_stamp_size": 30}, "positive odd integer"), ({"auto_stamp_size": True}, "positive odd integer"), + ({"auto_stamp_mask": True, "auto_stamp_size": 1}, "at least 3"), ({"auto_stamp_count": 0}, "positive integer"), ({"auto_stamp_count": True}, "positive integer"), ({"auto_peak_percentile": float("nan")}, "finite number"), ({"auto_peak_percentile": 101.0}, "finite number"), + ({"auto_stamp_mask": True, "auto_peak_percentile": 0}, "strictly"), + ({"auto_stamp_mask": True, "auto_peak_percentile": 100}, "strictly"), ], ) def test_batch_options_reject_invalid_auto_stamp_values( @@ -215,6 +218,13 @@ def test_batch_options_reject_invalid_auto_stamp_values( _options(**overrides) +@pytest.mark.parametrize("enabled, size", [(False, 1), (True, 3)]) +def test_batch_options_accept_minimum_stamp_size(enabled, size) -> None: + options = _options(auto_stamp_mask=enabled, auto_stamp_size=size) + + assert BatchFitOptions.from_payload(options.to_payload()) == options + + def test_batch_options_round_trip_worker_payload() -> None: options = _options(crop_y0=1, crop_x0=2, crop_height=16, crop_width=17)