diff --git a/docs/cli.md b/docs/cli.md index 26b0eb54..66d19192 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -32,6 +32,9 @@ See [xDataReader](components/xdr.md). `data-inspect` and `data-validate` check pickle-free NPZ dipole batches whose arrays are numeric or Unicode; `fit-dipoles` fits sampled-stamp or analytic Gaussian models and writes portable fit and uncertainty artifacts. +`fit-dipoles --executor dragon|mpi` distributes candidate chunks; local +execution remains the default. `--warmup-rounds` or `--measure-rounds` opts +into persistent workers and a separate artifact directory for each round. See [xFit](components/xfit.md). ## XPOIS: `cuphoton xpois` @@ -97,6 +100,12 @@ standalone raw-comparison and Alard--Lupton review servers are available as `cuphoton xscan --help`, then `cuphoton xscan help ` for command-specific contracts. See [XScan](components/xscan.md). +`infer-real-bogus --executor dragon|mpi` distributes complete minibatches; +`run-pipeline --executor dragon|mpi` distributes complete image pairs through +xPOIS, xFit and xScan. Optional round flags retain the workers across passes. +These commands require a shared input/output filesystem and the matching +runtime launcher. + ## xRep: `cuphoton xrep` `inspect-image`, `reproject-image`, `reproject-stack`, `compare-backends`, diff --git a/docs/components/xfit.md b/docs/components/xfit.md index e1468252..1be478d6 100644 --- a/docs/components/xfit.md +++ b/docs/components/xfit.md @@ -142,6 +142,50 @@ acceptable. See [Data and artifact contracts](../data-artifacts.md#xfit-dipole-batches) for the stable shapes and output fields. +## Distributed fitting + +`fit-dipoles --executor dragon|mpi` distributes independent candidate chunks +across GPUs. The default `--executor local` retains the original batch fit. +Distributed fitting requires `--backend cupy` or `--backend cutile`, a shared +filesystem for the input and output, and the same installed environment on +every worker. `--chunk-size` sets candidates per task independently of worker +count. Keep it fixed for matched comparisons; candidate IDs and input order +are restored in the merged artifacts. +The task count must be at least the MPI rank count. Dragon uses the smaller +of the requested worker count and task count. + +Under an allocation with Dragon configured, run one warmup and two measured +passes with workers retained across all three passes: + +```bash +dragon .venv/bin/cuphoton xfit fit-dipoles \ + --executor dragon --max-workers 8 \ + --input /shared/dipoles.npz --model gaussian --backend cupy \ + --chunk-size 256 --output-dir /shared/results/xfit-dragon \ + --warmup-rounds 1 --measure-rounds 2 +``` + +With Open MPI, use the installed rank wrapper to narrow GPU visibility before +Python starts. The parent mask must list allocated GPUs in local-rank order: + +```bash +: "${CUDA_VISIBLE_DEVICES:?must enumerate the allocated GPUs}" +mpirun -n 8 --map-by slot --bind-to none -x CUDA_VISIBLE_DEVICES \ + .venv/bin/cuphoton-openmpi-rank-exec -- \ + .venv/bin/cuphoton xfit fit-dipoles \ + --executor mpi --input /shared/dipoles.npz --model gaussian --backend cupy \ + --chunk-size 256 --output-dir /shared/results/xfit-mpi \ + --warmup-rounds 1 --measure-rounds 2 +``` + +The output directory must be new. Its basename is the run ID: 1–128 ASCII +letters, digits, dots, underscores or hyphens, starting with a letter or digit. +Each pass retains normal xFit artifacts +under `rounds//scientific/`, including warmup passes. Without round +flags, a single pass writes them under `scientific/`. The execution summary +records placement, item receipts and round timing; scientific merging and +validation occur after the timed worker phase. + ## Opt-in observational-data checks The observational checks read caller-supplied FITS files and create every diff --git a/docs/components/xscan.md b/docs/components/xscan.md index 54493928..4f0c90e5 100644 --- a/docs/components/xscan.md +++ b/docs/components/xscan.md @@ -236,6 +236,64 @@ The `*.blackwell.example.yaml` files demonstrate throughput-oriented settings for recent NVIDIA GPUs. Tune these starting points for your hardware and dataset. +## Distributed inference + +`infer-real-bogus --executor dragon|mpi` scores a fixed split across GPUs. +Each worker loads the checkpoint once and retains it across optional warmup +and measured passes. Training and evaluation commands retain their existing +local behavior. The default `--executor local` also preserves the original +inference output location. + +Inputs, checkpoint, optional xFit feature bundle and output directory must +reside on a filesystem shared by all workers. `--batch-size` retains its +ordinary inference meaning. `--task-batches` groups whole minibatches into +tasks; only the final task can contain the original final partial batch. +Keep both values fixed when comparing worker counts. The merge restores +selected-split order, original sample indices, labels and candidate metadata, +and uses the same host probability calculation as local inference. +The task count must be at least the MPI rank count. Dragon uses the smaller +of the requested worker count and task count. + +Distributed inference defaults to `--num-workers 0`, independently of the +checkpoint's loader setting. An explicit positive value creates persistent +loader processes during worker setup; the loader and parsed metadata are +reused across tasks and rounds. Workers retain inputs and model state, so +their initial load is outside the timed rounds. + +Under a configured Dragon allocation: + +```bash +dragon .venv/bin/cuphoton xscan infer-real-bogus \ + --executor dragon --max-workers 8 \ + --run-dir /shared/model --dataset-dir /shared/dataset --split test \ + --batch-size 32 --task-batches 16 --num-workers 0 \ + --output-dir /shared/results/inference-dragon \ + --warmup-rounds 1 --measure-rounds 2 +``` + +With Open MPI, the rank wrapper narrows GPU visibility before Python starts. +The parent mask must list allocated GPUs in local-rank order: + +```bash +: "${CUDA_VISIBLE_DEVICES:?must enumerate the allocated GPUs}" +mpirun -n 8 --map-by slot --bind-to none -x CUDA_VISIBLE_DEVICES \ + .venv/bin/cuphoton-openmpi-rank-exec -- \ + .venv/bin/cuphoton xscan infer-real-bogus \ + --executor mpi --run-dir /shared/model --dataset-dir /shared/dataset \ + --split test --batch-size 32 --task-batches 16 --num-workers 0 \ + --output-dir /shared/results/inference-mpi \ + --warmup-rounds 1 --measure-rounds 2 +``` + +Distributed inference requires a new `--output-dir`. Its basename is the run +ID: 1–128 ASCII letters, digits, dots, underscores or hyphens, starting with +a letter or digit. Each pass writes merged +logits, labels, probabilities, sample indices and a summary under +`rounds//scientific/`; warmup outputs are retained too. Without +round flags, the single pass uses `scientific/`. The model directory remains +unchanged. Execution receipts and timing are separate from these scientific +outputs, and merging and validation occur after the timed worker phase. + ## Persistent XPOIS, xFit and XScan pipeline The Python API in `cuphoton.xscan.device_pipeline` runs complete image pairs @@ -356,6 +414,48 @@ scientific evidence; Dragon saves each result in `items//summary.json` and checks it again in the coordinator. The evidence decoder returns the 22 named arrays for comparison. +The `run-pipeline` command accepts the same descriptors with either Dragon or +MPI. Save the configuration and items above as a manifest: + +```python +import json + +Path("pipeline.json").write_text(json.dumps({ + "schema": "cuphoton.xscan.pipeline-manifest/v1", + "configuration": config.to_payload(), + "items": [item.to_payload() for item in items], +}, indent=2) + "\n") +``` + +Paths may be absolute or relative to the manifest. All nodes must see the +same source, inputs and output directory. Use your site's launcher settings; +these examples run one warmup pass and two measured passes: + +```bash +dragon cuphoton xscan run-pipeline --executor dragon \ + --manifest pipeline.json --output-dir runs --name dragon-pipeline \ + --max-workers 8 --warmup-rounds 1 --measure-rounds 2 + +mpiexec -n 8 -x CUDA_VISIBLE_DEVICES cuphoton-openmpi-rank-exec -- \ + cuphoton xscan run-pipeline --executor mpi \ + --manifest pipeline.json --output-dir runs --name mpi-pipeline \ + --warmup-rounds 1 --measure-rounds 2 +``` + +The MPI example requires Open MPI and the allocation's visible GPU list on +each node. Other MPI launchers must bind each rank to exactly one GPU before +starting Python; do not use the Open MPI wrapper with MPICH. Each rank uses +local `cuda:0`, and duplicate physical GPU assignments fail validation. + +Both executors initialize one worker context and reuse it across all items +and rounds. Warmup outputs remain under `rounds/warmup-*`; measured outputs +remain under `rounds/measure-*`. Each round retains ordinary item results and +audit evidence. The parent `summary.json` separates readiness, batch time, +artifact validation and cleanup. Batch time includes dispatch, input loading, +numerical work and worker output publication; coordinator audits follow that +timer. It is not kernel-only time. Failed rounds or cleanup invalidate the +reported statistics. Omitting both round flags runs a single pass. + The pipeline retains device owners through the blocking terminal copy and synchronizes failed work before reuse. Failed cleanup makes the context unusable. Transfer receipts count pipeline-owned uploads and the packed diff --git a/docs/data-artifacts.md b/docs/data-artifacts.md index 5353072a..6831e4c5 100644 --- a/docs/data-artifacts.md +++ b/docs/data-artifacts.md @@ -47,6 +47,12 @@ Covariance uses supplied variances when present and residual scaling otherwise. Non-converged or rank-deficient fits retain status information and mark their uncertainties invalid. +For distributed xFit and XScan inference, `summary.json` is the terminal +execution summary. Merged scientific artifacts live in `scientific/` for +ordinary runs or `rounds//scientific/` for benchmark runs. Round +summaries describe individual passes; only the root summary includes the +executor's final lifecycle status. Warmup artifacts are retained. + ## XScan HSC NPY inputs XScan accepts an HSC NPY directory directly or as `HSC_npy` beneath a supplied diff --git a/src/cuphoton/xfit/commands.py b/src/cuphoton/xfit/commands.py index ad79edb3..28c829e3 100644 --- a/src/cuphoton/xfit/commands.py +++ b/src/cuphoton/xfit/commands.py @@ -12,6 +12,7 @@ import numpy as np +from cuphoton.core.bulk import validate_identifier from cuphoton.core.cli import ( BoolInvariant, CommandError, @@ -22,6 +23,7 @@ SetInvariant, StringInvariant, ) +from cuphoton.core.cli.executor import ExecutorOptions from ._types import ( BACKEND_REQUESTS, @@ -186,7 +188,7 @@ def run(self) -> None: self._emit_json(payload) -class FitDipolesCommand(_ValidatedDatasetCommand): +class FitDipolesCommand(ExecutorOptions, _ValidatedDatasetCommand): """Fit a batch of astronomical dipoles and persist safe artifacts.""" output_dir = None @@ -199,6 +201,13 @@ class FitDipolesCommand(_ValidatedDatasetCommand): g_tol = None max_evaluations = None use_finite_difference = None + chunk_size = None + + class ChunkSizeArg(PositiveIntegerInvariant): + _arg = "--chunk-size" + _help = "Candidates per distributed task. [default: 256]" + _mandatory = False + _default = None class OutputDirArg(PathSpecInvariant): _arg = "--output-dir" @@ -283,6 +292,59 @@ def _model_specification( ) def run(self) -> None: + executor_options = self.executor_options() + if self.executor != "local": + from cuphoton.core.executors import run_workload + + from .executor import prepare_xfit_workload + + if self.executor == "mpi": + executor_options["prepare_on_root"] = True + output_dir = Path(self.output_dir).expanduser().resolve() + self._call( + validate_identifier, + output_dir.name, + field="--output-dir basename", + ) + fit_options = { + name: getattr(self, name) + for name in ( + "model", + "mode", + "backend", + "compute_dtype", + "stamp_evaluation", + "stamp_scale", + "f_tol", + "x_tol", + "g_tol", + "max_evaluations", + "use_finite_difference", + ) + } + result = self._call( + run_workload, + executor=self.executor, + prepare_workload=lambda rank: prepare_xfit_workload( + input_path=self.input, + chunk_size=self.chunk_size or 256, + fit_options=fit_options, + retain_input=rank == 0, + ), + output_root=output_dir.parent, + run_id=output_dir.name, + **executor_options, + ) + if result is not None: + self._emit_json(result.to_dict()) + if result.status != "success": + raise CommandError("distributed xFit execution failed") + return + if self.chunk_size is not None: + raise CommandError( + "--chunk-size requires --executor dragon or mpi" + ) + from . import LMConfig, fit_dipoles dataset = self._load_dataset() diff --git a/src/cuphoton/xfit/executor.py b/src/cuphoton/xfit/executor.py new file mode 100644 index 00000000..73a90779 --- /dev/null +++ b/src/cuphoton/xfit/executor.py @@ -0,0 +1,546 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Candidate-sharded execution of the ordinary xFit numerical workflow.""" + +from __future__ import annotations + +import hashlib +import json +import time +from collections.abc import Mapping, Sequence +from dataclasses import fields, replace +from pathlib import Path +from typing import Any + +import numpy as np + +from cuphoton.core.artifacts import array_sha256, file_sha256 +from cuphoton.core.bulk import ( + WorkItem, + atomic_write_json, + collect_gpu_identity, + json_mapping, + read_json_mapping, +) + +from .api import DipoleFitResult, _floating_dtype, fit_dipoles +from .io import ( + XFitDataset, + _broadcast_image_auxiliary, + load_xfit_dataset, + write_fit_artifacts, +) +from .models import GaussianDipoleModel, StampDipoleModel +from .solver import LMConfig + +_SCHEMA = "cuphoton.xfit.executor-chunk/v1" +_META_FIELDS = { + "parameter_names", + "backend", + "device", + "dtype", + "model", + "mode", +} +_ARRAY_FIELDS = tuple( + field.name + for field in fields(DipoleFitResult) + if field.name not in _META_FIELDS +) +_DEFAULTS = { + "model": "gaussian", + "mode": "difference", + "backend": "cupy", + "compute_dtype": "input", + "stamp_evaluation": "bilinear", + "stamp_scale": 1.0, + "f_tol": None, + "x_tol": None, + "g_tol": None, + "max_evaluations": None, + "use_finite_difference": False, +} + + +def _digest(payload: Mapping[str, Any]) -> str: + return hashlib.sha256( + json.dumps( + payload, sort_keys=True, separators=(",", ":"), allow_nan=False + ).encode() + ).hexdigest() + + +def plan_xfit_chunks( + input_path: Path, *, chunk_size: int, fit_options: Mapping[str, Any] +) -> tuple[tuple[WorkItem, ...], dict[str, Any]]: + """Freeze candidate ranges independently of the eventual worker count.""" + + items, options, _ = _plan_xfit_chunks( + input_path, chunk_size=chunk_size, fit_options=fit_options + ) + return items, options + + +def _plan_xfit_chunks( + input_path: Path, *, chunk_size: int, fit_options: Mapping[str, Any] +) -> tuple[tuple[WorkItem, ...], dict[str, Any], XFitDataset]: + + if ( + isinstance(chunk_size, bool) + or not isinstance(chunk_size, int) + or chunk_size < 1 + ): + raise ValueError("chunk_size must be a positive integer") + if set(fit_options) - set(_DEFAULTS): + raise ValueError("unsupported xFit execution option") + settings = json_mapping( + {**_DEFAULTS, **fit_options}, field="xFit options" + ) + if settings["compute_dtype"] not in {"input", "float32", "float64"}: + raise ValueError("unsupported xFit compute dtype") + dataset = load_xfit_dataset( + input_path, model=settings["model"], mode=settings["mode"] + ) + if not dataset.batch_size: + raise ValueError("distributed xFit requires at least one candidate") + options = { + "input_path": str(dataset.path), + "input_sha256": dataset.input_archive_sha256, + "candidate_count": dataset.batch_size, + "chunk_size": chunk_size, + "result_dtype": ( + str(_floating_dtype(dataset.images)) + if settings["compute_dtype"] == "input" + else settings["compute_dtype"] + ), + "fit_options": settings, + } + stat = dataset.path.stat() + options["input_stat"] = { + "size": stat.st_size, + "mtime_ns": stat.st_mtime_ns, + } + options["configuration_sha256"] = _digest(options) + items = tuple( + WorkItem( + f"candidates-{start:08d}-{stop:08d}", + { + "start": start, + "stop": stop, + "candidate_ids_sha256": array_sha256( + dataset.candidate_id[start:stop] + ), + "configuration_sha256": options["configuration_sha256"], + }, + int(dataset.images[start:stop].nbytes), + ) + for start in range(0, dataset.batch_size, chunk_size) + for stop in (min(start + chunk_size, dataset.batch_size),) + ) + if file_sha256(dataset.path) != dataset.input_archive_sha256: + raise ValueError("xFit input archive changed while planning") + _check_input_stat(options) + return items, options, dataset + + +def _load_input(options: Mapping[str, Any]) -> XFitDataset: + settings = options["fit_options"] + dataset = load_xfit_dataset( + options["input_path"], model=settings["model"], mode=settings["mode"] + ) + if dataset.input_archive_sha256 != options["input_sha256"]: + raise ValueError("xFit input archive changed after planning") + if file_sha256(dataset.path) != options["input_sha256"]: + raise ValueError("xFit input archive changed while loading") + return dataset + + +def _chunk_dataset( + dataset: XFitDataset, start: int, stop: int +) -> XFitDataset: + auxiliaries = { + name: ( + None + if (array := getattr(dataset, name)) is None + else _broadcast_image_auxiliary(array, dataset.images)[start:stop] + ) + for name in ("mask", "variance") + } + initial = dataset.initial + if initial is not None and initial.ndim == 2: + initial = initial[start:stop] + return replace( + dataset, + candidate_id=dataset.candidate_id[start:stop], + images=dataset.images[start:stop], + initial=initial, + **auxiliaries, + ) + + +class XFitWorker: + """Retain validated input and solver configuration per worker.""" + + def __init__(self, options: Mapping[str, Any]) -> None: + self.options = dict(options) + self.dataset = _load_input(options) + self.settings = options["fit_options"] + self.gpu_identity = collect_gpu_identity(self.settings["backend"]) + self.config = LMConfig( + **{ + name: self.settings[name] + for name in ("f_tol", "x_tol", "g_tol", "max_evaluations") + }, + use_finite_difference=( + self.settings["use_finite_difference"] + or self.settings["model"] == "stamp" + ), + ) + self.model: Any = "gaussian" + if self.settings["model"] == "stamp": + basis = np.asarray(self.dataset.stamp_basis) + self.model = StampDipoleModel( + basis[0] if basis.ndim == 3 else basis, + image_shape=tuple(self.dataset.images.shape[-2:]), + evaluation=self.settings["stamp_evaluation"], + scale=self.settings["stamp_scale"], + ) + + def run_item(self, item: WorkItem, item_dir: Path) -> Mapping[str, Any]: + started = time.perf_counter() + start, stop = _range(item, self.options) + dataset = _chunk_dataset(self.dataset, start, stop) + if ( + array_sha256(dataset.candidate_id) + != item.payload["candidate_ids_sha256"] + ): + raise ValueError("xFit candidate identity differs from work item") + _check_input_stat(self.options) + images = dataset.images + if self.settings["compute_dtype"] != "input": + images = images.astype(self.settings["compute_dtype"], copy=False) + solve_start = time.perf_counter() + result = fit_dipoles( + images, + model=self.model, + initial=dataset.initial, + mask=dataset.mask, + variance=dataset.variance, + mode=self.settings["mode"], + backend=self.settings["backend"], + config=self.config, + ) + solve_seconds = time.perf_counter() - solve_start + _check_input_stat(self.options) + item_dir.mkdir(parents=True, exist_ok=False) + archive = item_dir / "result.npz" + np.savez_compressed( + archive, + candidate_id=dataset.candidate_id, + **{ + name: np.asarray(getattr(result, name)) + for name in _ARRAY_FIELDS + }, + ) + summary = { + "schema": _SCHEMA, + "item_id": item.item_id, + **dict(item.payload), + "input_sha256": self.options["input_sha256"], + "result_sha256": file_sha256(archive), + "result_metadata": { + name: list(result.parameter_names) + if name == "parameter_names" + else getattr(result, name) + for name in sorted(_META_FIELDS) + }, + } + atomic_write_json(item_dir / "summary.json", summary) + return { + "run_dir": str(item_dir), + "summary_path": str(item_dir / "summary.json"), + "requested_backend": self.settings["backend"], + "backend": result.backend, + "device": result.device, + "runtime": { + "configuration_sha256": self.options["configuration_sha256"] + }, + "timings_sec": {"fit_sec": solve_seconds}, + "wall_sec": {"item_runner": time.perf_counter() - started}, + } + + def close(self) -> None: + """Release host input ownership; CUDA caches belong to the process.""" + self.dataset = None + self.model = None + + +def create_xfit_worker(options: Mapping[str, Any]) -> XFitWorker: + """Construct a worker after its executor has bound the GPU.""" + return XFitWorker(options) + + +def _check_input_stat(options: Mapping[str, Any]) -> None: + stat = Path(options["input_path"]).stat() + if {"size": stat.st_size, "mtime_ns": stat.st_mtime_ns} != options[ + "input_stat" + ]: + raise ValueError("xFit input archive changed during execution") + + +def _range(item: WorkItem, options: Mapping[str, Any]) -> tuple[int, int]: + start, stop = item.payload.get("start"), item.payload.get("stop") + if ( + type(start) is not int + or type(stop) is not int + or not 0 <= start < stop <= options["candidate_count"] + or item.payload.get("configuration_sha256") + != options["configuration_sha256"] + ): + raise ValueError("invalid xFit chunk identity") + return start, stop + + +def _read_result( + item: WorkItem, round_dir: Path, options: Mapping[str, Any] +) -> tuple[dict[str, Any], dict[str, np.ndarray]]: + item_dir = round_dir / "items" / item.item_id + summary = read_json_mapping(item_dir / "summary.json") + expected = { + "schema": _SCHEMA, + "item_id": item.item_id, + **dict(item.payload), + "input_sha256": options["input_sha256"], + } + if any(summary.get(key) != value for key, value in expected.items()): + raise ValueError("xFit chunk summary identity mismatch") + archive = item_dir / "result.npz" + if file_sha256(archive) != summary.get("result_sha256"): + raise ValueError("xFit chunk result checksum mismatch") + with np.load(archive, allow_pickle=False) as loaded: + if set(loaded.files) != {*_ARRAY_FIELDS, "candidate_id"}: + raise ValueError("xFit chunk result fields mismatch") + arrays = {name: loaded[name] for name in loaded.files} + start, stop = _range(item, options) + if any( + array.ndim < 1 or array.shape[0] != stop - start + for array in arrays.values() + ): + raise ValueError("xFit chunk result row count mismatch") + if ( + array_sha256(arrays["candidate_id"]) + != item.payload["candidate_ids_sha256"] + ): + raise ValueError("xFit chunk result candidate identity mismatch") + metadata = summary.get("result_metadata") + if not isinstance(metadata, dict) or set(metadata) != _META_FIELDS: + raise ValueError("xFit chunk result metadata mismatch") + settings = options["fit_options"] + model_type = ( + GaussianDipoleModel + if settings["model"] == "gaussian" + else StampDipoleModel + ) + expected_metadata = { + "model": settings["model"], + "mode": settings["mode"], + "backend": settings["backend"], + "dtype": options["result_dtype"], + "parameter_names": list(model_type.parameter_names), + } + if any( + metadata.get(key) != value for key, value in expected_metadata.items() + ): + raise ValueError("xFit chunk result configuration mismatch") + return metadata, arrays + + +def _effective_config( + options: Mapping[str, Any], + dataset: XFitDataset, + result: DipoleFitResult, + output: Path, +) -> dict[str, Any]: + settings = options["fit_options"] + tolerance = float(np.sqrt(np.finfo(np.dtype(result.dtype)).eps)) + finite_difference = ( + settings["use_finite_difference"] or settings["model"] == "stamp" + ) + config = LMConfig( + **{ + name: settings[name] + for name in ("f_tol", "x_tol", "g_tol", "max_evaluations") + }, + use_finite_difference=finite_difference, + ) + return { + "schema_version": 1, + "command": "fit-dipoles", + "input": str(dataset.path), + "output_dir": str(output), + "model": settings["model"], + "mode": settings["mode"], + "backend": settings["backend"], + "compute_dtype": { + "requested": settings["compute_dtype"], + "resolved": result.dtype, + }, + "stamp": ( + { + "evaluation": settings["stamp_evaluation"], + "scale": settings["stamp_scale"], + "basis": "input:stamp_basis", + } + if settings["model"] == "stamp" + else None + ), + "solver": { + **{ + name: tolerance if settings[name] is None else settings[name] + for name in ("f_tol", "x_tol", "g_tol") + }, + "max_evaluations": config.resolved_max_evaluations( + len(result.parameter_names) + ), + "initial_damping": config.initial_damping, + "damping_increase": config.damping_increase, + "damping_decrease": config.damping_decrease, + "finite_difference_step": tolerance + if config.finite_difference_step is None + else config.finite_difference_step, + "use_finite_difference": finite_difference, + }, + } + + +def finalize_xfit_round( + round_dir: Path, + records: Sequence[Mapping[str, Any]], + *, + items: Sequence[WorkItem], + options: Mapping[str, Any], + _dataset: XFitDataset | None = None, +) -> Mapping[str, Any]: + """Restore candidate order and publish ordinary xFit artifacts.""" + expected = {item.item_id for item in items} + observed = [record.get("item_id") for record in records] + if ( + len(observed) != len(expected) + or set(observed) != expected + or any(record.get("status") != "success" for record in records) + ): + raise ValueError( + "xFit merge requires one successful result per chunk" + ) + if _dataset is None: + dataset = _load_input(options) + else: + dataset = _dataset + _check_input_stat(options) + if file_sha256(dataset.path) != options["input_sha256"]: + raise ValueError("xFit input archive changed after planning") + _check_input_stat(options) + ordered = sorted(items, key=lambda item: item.payload["start"]) + chunks = [] + metadata = None + next_row = 0 + for item in ordered: + start, stop = _range(item, options) + if start != next_row: + raise ValueError("xFit chunks overlap or omit candidate rows") + current_metadata, arrays = _read_result(item, round_dir, options) + if metadata is not None and current_metadata != metadata: + raise ValueError("xFit result metadata differs across chunks") + metadata = current_metadata + if not np.array_equal( + arrays["candidate_id"], dataset.candidate_id[start:stop] + ): + raise ValueError("xFit merged candidate order differs from input") + chunks.append(arrays) + next_row = stop + if next_row != dataset.batch_size or metadata is None: + raise ValueError("xFit chunks do not cover the complete input") + values = { + name: np.concatenate([chunk[name] for chunk in chunks]) + for name in _ARRAY_FIELDS + } + values["uncertainty_reason"] = tuple( + values["uncertainty_reason"].tolist() + ) + result = DipoleFitResult( + **values, + **{**metadata, "parameter_names": tuple(metadata["parameter_names"])}, + ) + output = round_dir / "scientific" + summary = write_fit_artifacts( + output, + dataset=dataset, + result=result, + effective_config=_effective_config(options, dataset, result, output), + ) + return { + "summary_path": "scientific/summary.json", + "candidate_count": dataset.batch_size, + "artifact_sha256": summary["artifact_sha256"], + } + + +def prepare_xfit_workload( + *, + input_path: Path, + chunk_size: int = 256, + fit_options: Mapping[str, Any], + retain_input: bool = True, +): + """Preflight one standalone fit for either shared GPU executor.""" + from cuphoton.core.execution import WorkloadSpec + + settings = {**_DEFAULTS, **fit_options} + backend = settings["backend"] + if backend not in {"cupy", "cutile"}: + raise ValueError("distributed xFit requires backend cupy or cutile") + if backend == "cutile" and settings["model"] == "stamp": + raise ValueError( + "backend='cutile' currently supports only the Gaussian model" + ) + if backend == "cutile" and settings["use_finite_difference"]: + raise ValueError( + "backend='cutile' does not support finite-difference fitting" + ) + items, options, dataset = _plan_xfit_chunks( + input_path, chunk_size=chunk_size, fit_options=fit_options + ) + if not retain_input: + dataset = None + by_id = {item.item_id: item for item in items} + + def validate(record, round_dir): + try: + _read_result(by_id[record["item_id"]], round_dir, options) + except (KeyError, OSError, TypeError, ValueError) as exc: + return [str(exc)] + return [] + + manifest = { + "schema": "cuphoton.xfit.executor-manifest/v1", + "items": [item.to_dict() for item in items], + "options": options, + } + return WorkloadSpec( + items=items, + options_payload=options, + manifest_payload=manifest, + input_identity_payload={ + "input": options["input_path"], + "sha256": options["input_sha256"], + }, + manifest_sha256=_digest(manifest), + backend=backend, + worker_factory=create_xfit_worker, + success_record_validator=validate, + finalize_round=lambda round_dir, records: finalize_xfit_round( + round_dir, records, items=items, options=options, _dataset=dataset + ), + ) diff --git a/src/cuphoton/xscan/commands.py b/src/cuphoton/xscan/commands.py index 51cadc8e..22669ebc 100644 --- a/src/cuphoton/xscan/commands.py +++ b/src/cuphoton/xscan/commands.py @@ -11,6 +11,7 @@ from pathlib import Path from typing import Any, Callable, TypeVar +from cuphoton.core.bulk import validate_identifier from cuphoton.core.cli import ( BoolInvariant, CommandError, @@ -20,6 +21,7 @@ SetInvariant, StringInvariant, ) +from cuphoton.core.cli.executor import ExecutorOptions T = TypeVar("T") _DEFAULT_REVIEW_DIR = os.environ.get("CUPHOTON_XSCAN_REVIEW_DIR") @@ -888,7 +890,57 @@ def run(self) -> None: self._emit_json({"run_dir": str(result.run_dir), **result.summary}) -class InferRealBogusCommand(XScanCommand): +class RunPipelineCommand(ExecutorOptions, XScanCommand): + """Run complete XPOIS, xFit and XScan jobs in persistent GPU workers.""" + + manifest = None + output_dir = None + run_name = None + + class ExecutorArg(ExecutorOptions.ExecutorArg): + _set = {"dragon", "mpi"} + _default = None + _mandatory = True + _help = ( + "Distributed runtime: dragon or mpi, with its matching launcher." + ) + + class ManifestArg(PathSpecInvariant): + _arg = "--manifest" + _mandatory = True + _help = "JSON pipeline manifest with configuration and image pairs." + + class OutputDirArg(PathSpecInvariant): + _arg = "--output-dir" + _mandatory = True + _help = "Parent directory for immutable pipeline runs." + + class RunNameArg(StringInvariant): + _arg = "--name" + _default = None + _help = "Optional unique run identifier." + + def run(self) -> None: + options = self.executor_options() + from .pipeline_executor import run_pipeline_manifest + + result = self._call( + run_pipeline_manifest, + executor=self.executor, + manifest_path=Path(self.manifest), + output_root=Path(self.output_dir), + run_id=self.run_name, + **options, + ) + if result is not None: + self._emit_json(result.to_dict()) + if result.status != "success": + raise CommandError( + f"Pipeline failed; inspect {result.summary_path}" + ) + + +class InferRealBogusCommand(ExecutorOptions, XScanCommand): """Run inference for a trained real-bogus model on one split.""" run_dir = None @@ -898,6 +950,22 @@ class InferRealBogusCommand(XScanCommand): num_workers = None use_x_fit_features = None x_fit_feature_dir = None + output_dir = None + task_batches = None + + class OutputDirArg(PathSpecInvariant): + _arg = "--output-dir" + _help = "New execution directory; required with Dragon or MPI." + _mandatory = False + _default = None + + class TaskBatchesArg(PositiveIntegerInvariant): + _arg = "--task-batches" + _help = ( + "Whole inference minibatches per distributed task. [default: 16]" + ) + _mandatory = False + _default = None class RunDirArg(PathSpecInvariant): _arg = "--run-dir" @@ -923,7 +991,10 @@ class BatchSizeArg(PositiveIntegerInvariant): class NumWorkersArg(NonNegativeIntegerInvariant): _arg = "--num-workers" - _help = "Loader workers; 0 disables workers, omitted uses checkpoint." + _help = ( + "Loader workers; 0 disables workers. Distributed inference " + "defaults to 0; local inference uses the checkpoint setting." + ) _mandatory = False _default = None @@ -943,6 +1014,51 @@ class XFitFeatureDirArg(PathSpecInvariant): _default = None def run(self) -> None: + executor_options = self.executor_options() + if self.executor != "local": + from cuphoton.core.executors import run_workload + + from .executor import prepare_inference_workload + + if self.output_dir is None: + raise CommandError( + "--output-dir is required for distributed inference" + ) + if self.executor == "mpi": + executor_options["prepare_on_root"] = True + output_dir = Path(self.output_dir).expanduser().resolve() + self._call( + validate_identifier, + output_dir.name, + field="--output-dir basename", + ) + result = self._call( + run_workload, + executor=self.executor, + prepare_workload=lambda rank: prepare_inference_workload( + run_dir=Path(self.run_dir).expanduser(), + dataset_dir=Path(self.dataset_dir).expanduser(), + split=self.split, + batch_size=self.batch_size, + task_batches=self.task_batches or 16, + num_workers=self.num_workers, + use_xfit_features=bool(self.use_x_fit_features), + xfit_feature_dir=self._path(self.x_fit_feature_dir), + ), + output_root=output_dir.parent, + run_id=output_dir.name, + **executor_options, + ) + if result is not None: + self._emit_json(result.to_dict()) + if result.status != "success": + raise CommandError("distributed xScan inference failed") + return + if self.output_dir is not None or self.task_batches is not None: + raise CommandError( + "--output-dir and --task-batches require " + "--executor dragon or mpi" + ) result = self._call( infer_workflow, run_dir=Path(self.run_dir).expanduser(), diff --git a/src/cuphoton/xscan/dragon_pipeline.py b/src/cuphoton/xscan/dragon_pipeline.py index f83bf077..c4bf9802 100644 --- a/src/cuphoton/xscan/dragon_pipeline.py +++ b/src/cuphoton/xscan/dragon_pipeline.py @@ -545,6 +545,7 @@ def _observe_content_file( expected_sha256: str, description: str, cache: dict[Path, tuple[str, int]], + validate_content: bool = True, ) -> int: try: if not path.is_file(): @@ -556,7 +557,10 @@ def _observe_content_file( ) from exc observed = cache.get(path) if observed is None: - observed = (file_sha256(path), size_bytes) + observed = ( + file_sha256(path) if validate_content else expected_sha256, + size_bytes, + ) cache[path] = observed if observed != (expected_sha256, size_bytes): raise RuntimeError(f"{description} changed before Dragon launch") @@ -565,6 +569,8 @@ def _observe_content_file( def _preflight_items( items: Sequence[DevicePipelineItem], + *, + validate_content: bool = True, ) -> tuple[tuple[WorkItem, ...], list[dict[str, Any]]]: if not items: raise ValueError("Dragon device pipeline items must not be empty") @@ -586,6 +592,7 @@ def _preflight_items( expected_sha256=descriptor.sha256, description=f"{item.item_id} {role} input", cache=observed_files, + validate_content=validate_content, ) descriptor_payload = descriptor.to_payload() existing = by_path.get(descriptor.path) @@ -623,6 +630,7 @@ def _preflight( items: Sequence[DevicePipelineItem], *, config: DevicePipelineConfig, + validate_content: bool = True, ) -> tuple[ tuple[WorkItem, ...], dict[str, Any], @@ -644,6 +652,7 @@ def _preflight( expected_sha256=config.checkpoint_sha256, description="XScan checkpoint", cache=observed_files, + validate_content=validate_content, ) feature_schema_path = Path(config.feature_schema_path) feature_schema_size = _observe_content_file( @@ -651,8 +660,11 @@ def _preflight( expected_sha256=config.feature_schema_sha256, description="xFit feature schema", cache=observed_files, + validate_content=validate_content, + ) + work_items, item_identities = _preflight_items( + items, validate_content=validate_content ) - work_items, item_identities = _preflight_items(items) manifest_payload = { "schema": DRAGON_DEVICE_PIPELINE_MANIFEST_SCHEMA, "backend": _BACKEND, @@ -918,6 +930,7 @@ def _success_record_problems( output_root: Path, config: DevicePipelineConfig, items: Mapping[str, DevicePipelineItem], + run_dir: Path | None = None, ) -> tuple[str, ...]: problems: set[str] = set() item_id = record.get("item_id") @@ -994,7 +1007,7 @@ def _success_record_problems( problems.add("summary_path") return tuple(sorted(problems)) assert isinstance(summary_relative, str) - summary_path = output_root / run_id / summary_relative + summary_path = (run_dir or output_root / run_id) / summary_relative try: summary = read_json_mapping(summary_path) if file_sha256(summary_path) != metadata["summary_sha256"]: @@ -1034,6 +1047,7 @@ def _failed_record_problems( *, output_root: Path, items: Mapping[str, DevicePipelineItem], + run_dir: Path | None = None, ) -> tuple[str, ...]: item_id = record.get("item_id") run_id = record.get("run_id") @@ -1043,7 +1057,7 @@ def _failed_record_problems( or not isinstance(run_id, str) ): return () - item_dir = output_root / run_id / "items" / item_id + item_dir = (run_dir or output_root / run_id) / "items" / item_id return ("failed_item_output",) if os.path.lexists(item_dir) else () diff --git a/src/cuphoton/xscan/executor.py b/src/cuphoton/xscan/executor.py new file mode 100644 index 00000000..5e802d22 --- /dev/null +++ b/src/cuphoton/xscan/executor.py @@ -0,0 +1,565 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Ordered inference shards with one checkpoint load per placed worker.""" + +from __future__ import annotations + +import copy +import hashlib +import json +import time +from collections.abc import Mapping, Sequence +from dataclasses import asdict, replace +from pathlib import Path +from typing import Any + +import numpy as np + +from cuphoton.core.artifacts import array_sha256, file_sha256 +from cuphoton.core.bulk import WorkItem, atomic_write_json, read_json_mapping + +from .hsc import prediction_group_indices +from .metrics import sigmoid + +_SCHEMA = "cuphoton.xscan.executor-chunk/v1" +_OUTPUT_ARRAYS = ("logits", "labels", "probabilities", "sample_index") + + +def _digest(payload: Mapping[str, Any]) -> str: + return hashlib.sha256( + json.dumps( + payload, sort_keys=True, separators=(",", ":"), allow_nan=False + ).encode() + ).hexdigest() + + +def _input_files(options: Mapping[str, Any]) -> list[Path]: + dataset = Path(options["dataset_dir"]) + files = [Path(options["run_dir"]) / "checkpoint.pt"] + files.extend( + dataset / name + for name in ( + "search.npy", + "template.npy", + "labels.npy", + "split.npy", + "metadata.jsonl", + ) + ) + if (dataset / "difference.npy").exists(): + files.append(dataset / "difference.npy") + if options["xfit_feature_dir"] is not None: + feature_root = Path(options["xfit_feature_dir"]) + files.extend( + feature_root / name + for name in ( + "candidate-id.npy", + "features.npy", + "input-image-sha256.npy", + "schema.json", + ) + ) + return files + + +def _check_inputs( + options: Mapping[str, Any], *, hash_content: bool = False +) -> None: + for identity in options["input_files"]: + path = Path(identity["path"]) + stat = path.stat() + if ( + stat.st_size != identity["size"] + or stat.st_mtime_ns != identity["mtime_ns"] + ): + raise ValueError(f"XScan input changed: {path.name}") + if hash_content and file_sha256(path) != identity["sha256"]: + raise ValueError(f"XScan input checksum changed: {path.name}") + + +def _indices(options: Mapping[str, Any]) -> np.ndarray: + return prediction_group_indices( + np.load( + Path(options["dataset_dir"]) / "split.npy", allow_pickle=False + ), + options["split"], + ) + + +def plan_inference_chunks( + *, + run_dir: Path, + dataset_dir: Path, + split: str, + batch_size: int = 32, + task_batches: int = 16, + num_workers: int | None = None, + xfit_feature_dir: Path | None = None, + use_xfit_features: bool = False, +) -> tuple[tuple[WorkItem, ...], dict[str, Any]]: + """Describe complete original minibatches without importing Torch.""" + + for name, value in ( + ("batch_size", batch_size), + ("task_batches", task_batches), + ): + if type(value) is not int or value < 1: + raise ValueError(f"{name} must be a positive integer") + if num_workers is not None and ( + type(num_workers) is not int or num_workers < 0 + ): + raise ValueError("num_workers must be a non-negative integer") + if type(use_xfit_features) is not bool or use_xfit_features != ( + xfit_feature_dir is not None + ): + raise ValueError( + "xFit fusion requires both enablement and a feature directory" + ) + options = { + "run_dir": str(run_dir.expanduser().resolve()), + "dataset_dir": str(dataset_dir.expanduser().resolve()), + "split": split, + "batch_size": batch_size, + "task_batches": task_batches, + "num_workers": 0 if num_workers is None else num_workers, + "use_xfit_features": use_xfit_features, + "xfit_feature_dir": str(xfit_feature_dir.expanduser().resolve()) + if xfit_feature_dir is not None + else None, + } + selected = _indices(options) + if not selected.size: + raise ValueError( + "distributed XScan inference requires a non-empty split" + ) + options["sample_count"] = int(selected.size) + options["sample_indices_sha256"] = array_sha256(selected) + identities = [] + for path in _input_files(options): + before = path.stat() + checksum = file_sha256(path) + after = path.stat() + if (before.st_size, before.st_mtime_ns) != ( + after.st_size, + after.st_mtime_ns, + ): + raise ValueError("XScan input changed during planning") + identities.append( + { + "path": str(path), + "size": after.st_size, + "mtime_ns": after.st_mtime_ns, + "sha256": checksum, + } + ) + options["input_files"] = identities + options["configuration_sha256"] = _digest(options) + task_size = batch_size * task_batches + items = tuple( + WorkItem( + f"samples-{start:08d}-{stop:08d}", + { + "start": start, + "stop": stop, + "sample_indices_sha256": array_sha256(selected[start:stop]), + "configuration_sha256": options["configuration_sha256"], + }, + stop - start, + ) + for start in range(0, int(selected.size), task_size) + for stop in (min(start + task_size, int(selected.size)),) + ) + return items, options + + +def _device(): + import torch + + return torch.device("cuda:0") + + +def _gpu_identity() -> Mapping[str, Any]: + import torch + + index = torch.cuda.current_device() + properties = torch.cuda.get_device_properties(index) + uuid = getattr(properties, "uuid", None) + if not uuid: + raise RuntimeError("Torch did not expose a stable GPU UUID") + return { + "backend": "torch", + "device_index": index, + "name": properties.name, + "uuid": str(uuid), + "pci_bus_id": None, + "identity_error": None, + } + + +class XScanInferenceWorker: + """Retain the model, performance policy and dataset in one process.""" + + def __init__(self, options: Mapping[str, Any]) -> None: + from .dataset import load_metadata_rows + from .training import ( + load_model_from_checkpoint, + make_dataloader, + normalize_performance_config, + xfit_fit_coverage, + ) + from .workflows import ( + _checkpoint_dataset, + _checkpoint_xfit_feature_matrix, + ) + + self.options = dict(options) + _check_inputs(options, hash_content=True) + self.device = _device() + self.gpu_identity = _gpu_identity() + self.model, checkpoint, self.performance = load_model_from_checkpoint( + Path(options["run_dir"]), device=self.device + ) + self.performance = normalize_performance_config( + replace( + self.performance, + num_workers=options["num_workers"], + persistent_workers=options["num_workers"] > 0, + ), + device=self.device, + ) + feature_dir = ( + Path(options["xfit_feature_dir"]) + if options["xfit_feature_dir"] is not None + else None + ) + feature_matrix = _checkpoint_xfit_feature_matrix( + checkpoint=checkpoint, + dataset_dir=Path(options["dataset_dir"]), + xfit_feature_dir=feature_dir, + use_xfit_features=options["use_xfit_features"], + ) + self.dataset = _checkpoint_dataset( + checkpoint=checkpoint, + dataset_dir=Path(options["dataset_dir"]), + split=options["split"], + xfit_feature_matrix=feature_matrix, + ) + if ( + array_sha256(self.dataset.indices) + != options["sample_indices_sha256"] + ): + raise ValueError( + "XScan worker split differs from planned sample order" + ) + self.metadata_rows = load_metadata_rows(Path(options["dataset_dir"])) + self.summary = { + "workflow": "infer", + "run_dir": options["run_dir"], + "dataset_dir": options["dataset_dir"], + "split": options["split"], + "batch_size": options["batch_size"], + "performance": asdict(self.performance), + "xfit_features": { + "enabled": options["use_xfit_features"], + "feature_dir": options["xfit_feature_dir"], + "feature_names": list(self.dataset.xfit_feature_names), + "bundle_identity": self.dataset.xfit_feature_bundle_identity, + "join_diagnostics": self.dataset.xfit_join_diagnostics, + "fit_coverage": xfit_fit_coverage( + self.dataset, split=options["split"] + ) + if options["use_xfit_features"] + else None, + }, + } + _check_inputs(options, hash_content=True) + # The sampler lives in this process; loader children retain the full + # immutable dataset while each task selects its original minibatches. + self.sample_positions: list[int] = [] + self.loader = make_dataloader( + self.dataset, + batch_size=options["batch_size"], + shuffle=False, + performance=self.performance, + sampler=self.sample_positions, + ) + if self.performance.num_workers: + # The first iterator starts children; resetting it waits for their + # startup acknowledgements. Neither pass reads samples or runs the + # model, and child initialization completes before READY. + for _ in self.loader: + pass + for _ in self.loader: + pass + + def run_item(self, item: WorkItem, item_dir: Path) -> Mapping[str, Any]: + from .training import predict_dataset + + started = time.perf_counter() + start, stop = _range(item, self.options) + _check_inputs(self.options) + dataset = copy.copy(self.dataset) + dataset.indices = self.dataset.indices[start:stop] + if ( + array_sha256(dataset.indices) + != item.payload["sample_indices_sha256"] + ): + raise ValueError( + "XScan item sample order differs from its descriptor" + ) + self.sample_positions[:] = range(start, stop) + prediction_started = time.perf_counter() + payload = predict_dataset( + model=self.model, + dataset=dataset, + batch_size=self.options["batch_size"], + device=self.device, + performance=self.performance, + _loader=self.loader, + _metadata_rows=self.metadata_rows, + ) + prediction_seconds = time.perf_counter() - prediction_started + _check_inputs(self.options) + item_dir.mkdir(parents=True, exist_ok=False) + arrays = { + name: payload[name] + for name in ("logits", "labels", "probabilities") + } + arrays["sample_index"] = dataset.indices + checksums = {} + for name, array in arrays.items(): + path = item_dir / f"{name}.npy" + np.save(path, array, allow_pickle=False) + checksums[name] = file_sha256(path) + atomic_write_json( + item_dir / "summary.json", + { + "schema": _SCHEMA, + "item_id": item.item_id, + **dict(item.payload), + "artifact_sha256": checksums, + "metadata_rows": payload["metadata_rows"], + "inference_summary": self.summary, + }, + ) + return { + "run_dir": str(item_dir), + "summary_path": str(item_dir / "summary.json"), + "requested_backend": "torch", + "backend": "torch", + "device": str(self.device), + "runtime": { + "configuration_sha256": self.options["configuration_sha256"] + }, + "timings_sec": {"predict_sec": prediction_seconds}, + "wall_sec": {"item_runner": time.perf_counter() - started}, + } + + def close(self) -> None: + """Release model and dataset ownership after the final round.""" + iterator = getattr(self.loader, "_iterator", None) + if iterator is not None: + iterator._shutdown_workers() + self.loader = None + self.model = None + self.dataset = None + self.metadata_rows = None + + +def create_xscan_worker(options: Mapping[str, Any]) -> XScanInferenceWorker: + """Construct the Torch worker after executor device binding.""" + return XScanInferenceWorker(options) + + +def _range(item: WorkItem, options: Mapping[str, Any]) -> tuple[int, int]: + start, stop = item.payload.get("start"), item.payload.get("stop") + batch_size = options["batch_size"] + if ( + type(start) is not int + or type(stop) is not int + or not 0 <= start < stop <= options["sample_count"] + or start % batch_size + or (stop != options["sample_count"] and stop % batch_size) + or item.payload.get("configuration_sha256") + != options["configuration_sha256"] + ): + raise ValueError("invalid XScan chunk or minibatch boundary") + return start, stop + + +def _read_result(item: WorkItem, round_dir: Path, options: Mapping[str, Any]): + item_dir = round_dir / "items" / item.item_id + summary = read_json_mapping(item_dir / "summary.json") + expected = { + "schema": _SCHEMA, + "item_id": item.item_id, + **dict(item.payload), + } + if any(summary.get(key) != value for key, value in expected.items()): + raise ValueError("XScan chunk summary identity mismatch") + start, stop = _range(item, options) + arrays = {} + for name in _OUTPUT_ARRAYS: + path = item_dir / f"{name}.npy" + if file_sha256(path) != summary["artifact_sha256"].get(name): + raise ValueError("XScan chunk result checksum mismatch") + values = np.load(path, allow_pickle=False) + if values.shape != (stop - start,): + raise ValueError("XScan chunk row count mismatch") + arrays[name] = values + if ( + array_sha256(arrays["sample_index"]) + != item.payload["sample_indices_sha256"] + ): + raise ValueError("XScan chunk sample identity mismatch") + probabilities = arrays["probabilities"] + if ( + arrays["logits"].dtype != np.dtype("float64") + or probabilities.dtype != np.dtype("float64") + or not np.all(np.isfinite(arrays["logits"])) + or not np.all(np.isfinite(probabilities)) + ): + raise ValueError("XScan chunk scores must be finite float64 arrays") + try: + # NumPy exp dispatch can differ by a few ULP between hosts. Retain + # the worker's original probabilities and bound only this audit. + np.testing.assert_array_max_ulp( + probabilities, sigmoid(arrays["logits"]), maxulp=4 + ) + except AssertionError as exc: + raise ValueError( + "XScan chunk probabilities differ from the host sigmoid " + "by more than 4 ULP" + ) from exc + return summary, arrays + + +def finalize_inference_round( + round_dir: Path, + records: Sequence[Mapping[str, Any]], + *, + items: Sequence[WorkItem], + options: Mapping[str, Any], +) -> Mapping[str, Any]: + """Merge ordered predictions without fitting thresholds or metrics.""" + observed = [record.get("item_id") for record in records] + if ( + len(observed) != len(items) + or set(observed) != {item.item_id for item in items} + or any(record.get("status") != "success" for record in records) + ): + raise ValueError( + "XScan merge requires one successful result per chunk" + ) + _check_inputs(options, hash_content=True) + selected = _indices(options) + metadata_path = Path(options["dataset_dir"]) / "metadata.jsonl" + metadata = [ + json.loads(line) + for line in metadata_path.read_text(encoding="utf-8").splitlines() + if line.strip() + ] + labels = np.load( + Path(options["dataset_dir"]) / "labels.npy", allow_pickle=False + ) + chunks = [] + inference_summary = None + next_row = 0 + for item in sorted(items, key=lambda item: item.payload["start"]): + start, stop = _range(item, options) + if start != next_row: + raise ValueError("XScan chunks overlap or omit sample rows") + summary, arrays = _read_result(item, round_dir, options) + expected_rows = [ + {**metadata[int(index)], "sample_index": int(index)} + for index in selected[start:stop] + ] + if ( + summary["metadata_rows"] != expected_rows + or not np.array_equal( + arrays["sample_index"], selected[start:stop] + ) + or not np.array_equal( + arrays["labels"], + labels[selected[start:stop]].astype(np.int64), + ) + ): + raise ValueError( + "XScan chunk sample identities or labels differ from input" + ) + current_summary = summary["inference_summary"] + if ( + inference_summary is not None + and current_summary != inference_summary + ): + raise ValueError( + "XScan checkpoint, feature or performance settings " + "differ across workers" + ) + inference_summary = current_summary + chunks.append(arrays) + next_row = stop + if next_row != options["sample_count"] or inference_summary is None: + raise ValueError("XScan chunks do not cover the selected split") + output = round_dir / "scientific" + output.mkdir(exist_ok=False) + checksums = {} + for name in _OUTPUT_ARRAYS: + path = output / f"{name}.npy" + np.save( + path, + np.concatenate([chunk[name] for chunk in chunks]), + allow_pickle=False, + ) + checksums[name] = file_sha256(path) + # Paths in this execution-owned scientific summary are relative to itself, + # rather than to the read-only training run directory. + final_summary = { + **inference_summary, + "output_dir": str(output), + "saved": {name: f"{name}.npy" for name in _OUTPUT_ARRAYS}, + "artifact_sha256": checksums, + "configuration_sha256": options["configuration_sha256"], + } + atomic_write_json(output / "summary.json", final_summary) + return { + "summary_path": "scientific/summary.json", + "sample_count": next_row, + "artifact_sha256": checksums, + } + + +def prepare_inference_workload(**kwargs): + """Preflight inference descriptors without importing Torch.""" + from cuphoton.core.execution import WorkloadSpec + + items, options = plan_inference_chunks(**kwargs) + by_id = {item.item_id: item for item in items} + + def validate(record, round_dir): + try: + _read_result(by_id[record["item_id"]], round_dir, options) + except (KeyError, OSError, TypeError, ValueError) as exc: + return [str(exc)] + return [] + + manifest = { + "schema": "cuphoton.xscan.executor-manifest/v1", + "items": [item.to_dict() for item in items], + "options": options, + } + return WorkloadSpec( + items=items, + options_payload=options, + manifest_payload=manifest, + input_identity_payload={"files": options["input_files"]}, + manifest_sha256=_digest(manifest), + backend="torch", + worker_factory=create_xscan_worker, + success_record_validator=validate, + finalize_round=lambda round_dir, records: finalize_inference_round( + round_dir, records, items=items, options=options + ), + ) diff --git a/src/cuphoton/xscan/pipeline_executor.py b/src/cuphoton/xscan/pipeline_executor.py new file mode 100644 index 00000000..2caf0a75 --- /dev/null +++ b/src/cuphoton/xscan/pipeline_executor.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Manifest-driven device pipeline execution through Dragon or MPI.""" + +from __future__ import annotations + +import importlib +import json +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from cuphoton.core.bulk import WorkItem, collect_gpu_identity, json_mapping + +from .device_pipeline import DevicePipelineConfig, DevicePipelineItem +from .dragon_pipeline import ( + DRAGON_DEVICE_PIPELINE_OPTIONS_SCHEMA, + _failed_record_problems, + _preflight, + _publish_item_result, + _strict_worker_options, + _success_record_problems, +) + +if TYPE_CHECKING: + from cuphoton.core.execution import ExecutionResult, WorkloadSpec + +PIPELINE_MANIFEST_SCHEMA = "cuphoton.xscan.pipeline-manifest/v1" + + +def _unique_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"duplicate manifest key: {key}") + result[key] = value + return result + + +def load_pipeline_manifest( + path: Path, +) -> tuple[DevicePipelineConfig, tuple[DevicePipelineItem, ...]]: + """Read config/items and resolve paths beside the manifest.""" + + path = path.expanduser().resolve() + payload = json_mapping( + json.loads( + path.read_text(encoding="utf-8"), + object_pairs_hook=_unique_mapping, + ), + field="device pipeline manifest", + ) + if set(payload) != {"schema", "configuration", "items"}: + raise ValueError( + "pipeline manifest requires schema, configuration, items" + ) + if payload["schema"] != PIPELINE_MANIFEST_SCHEMA: + raise ValueError("unsupported pipeline manifest schema") + + def resolve(raw: Any) -> str: + if not isinstance(raw, str) or not raw: + raise ValueError("manifest paths must be nonempty strings") + value = Path(raw).expanduser() + nested = path.parent / value + return str(nested.parent.resolve() / nested.name) + + configuration = json_mapping( + payload["configuration"], field="configuration" + ) + for field in ("checkpoint_dir", "feature_schema_path"): + if field in configuration: + configuration[field] = resolve(configuration[field]) + config = DevicePipelineConfig.from_payload(configuration) + if config.device != "cuda:0": + raise ValueError("device pipeline manifest device must be cuda:0") + if not isinstance(payload["items"], list): + raise ValueError("pipeline items must be a list") + items = [] + for raw in payload["items"]: + item = json_mapping(raw, field="pipeline item") + for role in ("reference", "target", "variance", "fit_mask"): + if item.get(role) is not None: + descriptor = json_mapping(item[role], field=role) + if "path" in descriptor: + descriptor["path"] = resolve(descriptor["path"]) + item[role] = descriptor + items.append(DevicePipelineItem.from_payload(item)) + return config, tuple(items) + + +class _PipelineWorker: + def __init__(self, options: Mapping[str, Any]) -> None: + self.config = _strict_worker_options(options) + # Torch must establish its CUDA library load order before CuPy calls + # CUDA. The executor has already bound this process to one GPU. + torch = importlib.import_module("torch") + cp = importlib.import_module("cupy") + + from .device_pipeline import DeviceWorkerContext + + torch.cuda.set_device(0) + cp.cuda.Device(0).use() + self.gpu_identity = collect_gpu_identity("cupy") + self.context = DeviceWorkerContext.initialize(self.config) + self.closed = False + + def run_item(self, item: WorkItem, output_dir: Path) -> Mapping[str, Any]: + if self.closed: + raise RuntimeError("pipeline worker is closed") + return _publish_item_result( + work_item=item, + item_dir=output_dir, + config=self.config, + context=self.context, + ) + + def close(self) -> None: + self.closed = True + failure = None + for stream in ( + self.context.producer_stream, + self.context.consumer_stream, + ): + try: + stream.synchronize() + except Exception as exc: + if failure is None: + failure = exc + else: + failure.add_note(str(exc)) + if failure is not None: + raise RuntimeError( + "pipeline worker stream cleanup failed" + ) from failure + + +def create_pipeline_worker(options: Mapping[str, Any]) -> _PipelineWorker: + """Create a process-local context after executor placement.""" + + return _PipelineWorker(options) + + +def prepare_pipeline_workload( + items: Sequence[DevicePipelineItem], + config: DevicePipelineConfig, + *, + rank: int = 0, +) -> WorkloadSpec: + """Prepare descriptors with component-owned scientific validation.""" + + from cuphoton.core.execution import WorkloadSpec + + items = tuple(items) + work, manifest, identities, digest = _preflight( + items, config=config, validate_content=rank == 0 + ) + by_id = {item.item_id: item for item in items} + + def validate_success(record: Mapping[str, Any], round_dir: Path): + return _success_record_problems( + record, + output_root=round_dir.parent, + run_dir=round_dir, + config=config, + items=by_id, + ) + + def validate_failure(record: Mapping[str, Any], round_dir: Path): + return _failed_record_problems( + record, + output_root=round_dir.parent, + run_dir=round_dir, + items=by_id, + ) + + return WorkloadSpec( + items=work, + options_payload={ + "schema": DRAGON_DEVICE_PIPELINE_OPTIONS_SCHEMA, + "backend": "cupy", + "configuration": config.to_payload(), + "configuration_sha256": config.configuration_sha256, + }, + manifest_payload=manifest, + input_identity_payload=identities, + manifest_sha256=digest, + backend="cupy", + worker_factory=create_pipeline_worker, + success_record_validator=validate_success, + failed_record_validator=validate_failure, + ) + + +def run_pipeline_manifest( + *, + executor: str, + manifest_path: Path, + output_root: Path, + run_id: str | None = None, + **options: Any, +) -> ExecutionResult | None: + """Run an XPOIS/xFit/XScan manifest through the selected executor.""" + + from cuphoton.core.executors import run_workload + + def prepare(rank: int): + config, items = load_pipeline_manifest(manifest_path) + return prepare_pipeline_workload(items, config, rank=rank) + + return run_workload( + executor=executor, + prepare_workload=prepare, + output_root=output_root, + run_id=run_id, + **options, + ) diff --git a/src/cuphoton/xscan/training.py b/src/cuphoton/xscan/training.py index 97fffbc1..178cbdfa 100644 --- a/src/cuphoton/xscan/training.py +++ b/src/cuphoton/xscan/training.py @@ -11,7 +11,7 @@ import os import random import time -from collections.abc import Iterator +from collections.abc import Iterator, Sequence from contextlib import contextmanager, nullcontext from dataclasses import asdict, dataclass from dataclasses import field as dataclass_field @@ -374,6 +374,7 @@ def make_dataloader( batch_size: int, shuffle: bool, performance: PerformanceConfig, + sampler: Sequence[int] | None = None, ) -> DataLoader[CollatedDatasetBatch]: kwargs: dict[str, Any] = { "batch_size": batch_size, @@ -382,6 +383,8 @@ def make_dataloader( "pin_memory": performance.pin_memory, "persistent_workers": performance.persistent_workers, } + if sampler is not None: + kwargs["sampler"] = sampler if ( performance.num_workers > 0 and performance.worker_start_method is not None @@ -1613,17 +1616,22 @@ def predict_dataset( device: torch.device, performance: PerformanceConfig | None = None, _loader: DataLoader[CollatedDatasetBatch] | None = None, + _metadata_rows: Sequence[dict[str, Any]] | None = None, ) -> dict[str, Any]: resolved_performance = normalize_performance_config( performance or PerformanceConfig(), device=device, ) amp_dtype = resolve_amp_dtype(resolved_performance, device=device) - loader = _loader or make_dataloader( - dataset, - batch_size=batch_size, - shuffle=False, - performance=resolved_performance, + loader = ( + _loader + if _loader is not None + else make_dataloader( + dataset, + batch_size=batch_size, + shuffle=False, + performance=resolved_performance, + ) ) logits_chunks = [] label_chunks = [] @@ -1652,7 +1660,11 @@ def predict_dataset( if label_chunks else np.zeros((0,)) ) - rows = load_metadata_rows(dataset.dataset_dir) + rows = ( + load_metadata_rows(dataset.dataset_dir) + if _metadata_rows is None + else _metadata_rows + ) selected_rows = [] for idx in dataset.indices.tolist(): row = dict(rows[int(idx)]) diff --git a/tests/core/test_cli_contract.py b/tests/core/test_cli_contract.py index 4be2b23d..74d369dd 100644 --- a/tests/core/test_cli_contract.py +++ b/tests/core/test_cli_contract.py @@ -71,17 +71,17 @@ def test_public_command_surface_counts_are_exact() -> None: assert per_group == [ ("xdr", 1, 0, 14, 0), - ("xfit", 3, 3, 17, 1), + ("xfit", 3, 3, 25, 1), ("xpois", 7, 7, 146, 1), - ("xscan", 42, 42, 171, 1), + ("xscan", 43, 43, 190, 1), ("xrep", 6, 6, 103, 1), ("xray", 33, 31, 390, 1), ] assert len(per_group) == 6 - assert sum(item[1] for item in per_group) == 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) == 841 + assert sum(item[1] for item in per_group) == 93 + assert sum(item[1] + item[4] for item in per_group) == 98 + assert sum(item[2] for item in per_group) == 90 + assert sum(item[3] for item in per_group) == 868 def test_public_registry_order_and_component_derivations() -> None: diff --git a/tests/core/test_mpi_adapters.py b/tests/core/test_mpi_adapters.py new file mode 100644 index 00000000..4b91450a --- /dev/null +++ b/tests/core/test_mpi_adapters.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +import pickle +import sys +from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace +from threading import Barrier, local +from types import SimpleNamespace + +import numpy as np +import pytest + +from cuphoton.core import mpi +from cuphoton.core._mpi_runtime import _MPIAPI +from cuphoton.core.benchmark import BenchmarkOptions + + +@pytest.mark.parametrize("component", ["xfit", "xscan"]) +@pytest.mark.parametrize("repeated", [False, True]) +def test_two_rank_adapters_publish_scientific_results( + tmp_path, monkeypatch, component, repeated +): + for name in tuple(os.environ): + if name.startswith( + ( + "OMPI_", + "PMI_", + "PMIX_", + "SLURM_", + "MPI_LOCAL", + "CUPHOTON_ALLOCATED", + ) + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") + for name in ("cupy", "numba.cuda", "cuda.tile"): + monkeypatch.delitem(sys.modules, name, raising=False) + state = local() + barrier = Barrier(2, timeout=10) + messages = [None, None] + planning_ranks = [] + input_load_ranks = [] + + class Comm: + def Get_rank(self): + return state.rank + + def Get_size(self): + return 2 + + def Split_type(self, split_type, key): + return SimpleNamespace( + Get_rank=lambda: state.rank, Free=lambda: None + ) + + def exchange(self, kind, value, root): + messages[state.rank] = (kind, root, pickle.dumps(value)) + barrier.wait() + assert all(message[:2] == (kind, root) for message in messages) + values = [pickle.loads(message[2]) for message in messages] + barrier.wait() + return values + + def gather(self, value, root): + values = self.exchange("gather", value, root) + return values if state.rank == root else None + + def bcast(self, value, root): + return self.exchange("bcast", value, root)[root] + + monkeypatch.setattr( + mpi, + "_load_mpi_api", + lambda: _MPIAPI( + SimpleNamespace(COMM_TYPE_SHARED=1), + Comm(), + "test-mpi4py", + "test-MPI", + ), + ) + + def gpu_identity(backend): + return { + "backend": backend, + "device_index": 0, + "uuid": f"GPU-{state.rank}", + "identity_error": None, + } + + if component == "xfit": + from cuphoton.xfit import executor, fit_dipoles + + from ..xfit.test_executor import _input + + path = _input(tmp_path) + load_input = executor.load_xfit_dataset + + def counted_load(*args, **kwargs): + input_load_ranks.append(state.rank) + return load_input(*args, **kwargs) + + monkeypatch.setattr(executor, "load_xfit_dataset", counted_load) + + def cpu_fit(*args, **kwargs): + result = fit_dipoles(*args, **{**kwargs, "backend": "numpy"}) + return replace(result, backend="cupy", device="cuda:0") + + monkeypatch.setattr(executor, "fit_dipoles", cpu_fit) + monkeypatch.setattr(executor, "collect_gpu_identity", gpu_identity) + + def prepare(rank): + planning_ranks.append(rank) + return executor.prepare_xfit_workload( + input_path=path, + chunk_size=2, + fit_options={"backend": "cupy", "max_evaluations": 2}, + retain_input=rank == 0, + ) + + else: + torch = pytest.importorskip("torch") + from cuphoton.xscan import executor + + from ..xscan.test_executor import _inputs + + run_dir, dataset_dir = _inputs(tmp_path) + monkeypatch.setattr(executor, "_device", lambda: torch.device("cpu")) + monkeypatch.setattr( + executor, "_gpu_identity", lambda: gpu_identity("torch") + ) + + def prepare(rank): + planning_ranks.append(rank) + return executor.prepare_inference_workload( + run_dir=run_dir, + dataset_dir=dataset_dir, + split="test", + batch_size=2, + task_batches=1, + ) + + def run(rank): + state.rank = rank + return mpi.run_mpi_work_items( + prepare_workload=prepare, + prepare_on_root=True, + output_root=tmp_path, + run_id="distributed", + rank_setup_timeout_sec=1, + benchmark=BenchmarkOptions(warmup_rounds=1, measure_rounds=2) + if repeated + else None, + ) + + with ThreadPoolExecutor(max_workers=2) as pool: + root, peer = list(pool.map(run, (0, 1))) + assert planning_ranks == [0] + if component == "xfit": + # Rank zero plans once; each rank loads its worker input once. + assert sorted(input_load_ranks) == [0, 0, 1] + assert peer is None + assert root.status == "success", root.summary + run_dir = tmp_path / "distributed" + rounds = ( + [ + run_dir / "rounds" / name + for name in ("warmup-0000", "measure-0000", "measure-0001") + ] + if repeated + else [run_dir] + ) + for round_dir in rounds: + scientific = round_dir / "scientific" + assert (scientific / "summary.json").is_file() + if component == "xfit": + import pyarrow.parquet as pq + + assert pq.read_table(scientific / "fits.parquet")[ + "candidate_id" + ].to_pylist() == ["z", "a", "m"] + else: + np.testing.assert_array_equal( + np.load(scientific / "sample_index.npy"), [0, 2, 3, 4, 6] + ) diff --git a/tests/xfit/test_commands.py b/tests/xfit/test_commands.py index b26bdb2d..f975d85c 100644 --- a/tests/xfit/test_commands.py +++ b/tests/xfit/test_commands.py @@ -754,3 +754,55 @@ def test_command_help_exposes_stable_safe_input_options(capsys) -> None: assert "vignetted" in captured.out assert "finite-volume" in captured.out assert "--use-finite-difference" in captured.out + + +@pytest.mark.parametrize("runtime", ["dragon", "mpi"]) +def test_distributed_output_name_is_validated_before_launch( + tmp_path, monkeypatch, capsys, runtime +): + from cuphoton.core import executors + + input_path = tmp_path / "input.npz" + _write_gaussian_input(input_path) + monkeypatch.setattr( + executors, "run_workload", lambda **kwargs: pytest.fail("launched") + ) + rc = run_component( + "xfit", + [ + "fit-dipoles", + "--input", + str(input_path), + "--model", + "gaussian", + "--backend", + "cupy", + "--executor", + runtime, + "--output-dir", + str(tmp_path / "invalid name"), + ], + ) + assert rc != 0 + assert "--output-dir basename" in capsys.readouterr().err + + +def test_local_fitting_rejects_distributed_chunk_option(tmp_path, capsys): + input_path = tmp_path / "input.npz" + _write_gaussian_input(input_path) + rc = run_component( + "xfit", + [ + "fit-dipoles", + "--input", + str(input_path), + "--model", + "gaussian", + "--output-dir", + str(tmp_path / "output"), + "--chunk-size", + "1", + ], + ) + assert rc != 0 + assert "--chunk-size requires --executor" in capsys.readouterr().err diff --git a/tests/xfit/test_executor.py b/tests/xfit/test_executor.py new file mode 100644 index 00000000..4c00db8d --- /dev/null +++ b/tests/xfit/test_executor.py @@ -0,0 +1,450 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import gc +import json +import os +import weakref +from dataclasses import fields, replace + +import numpy as np +import pyarrow.parquet as pq +import pytest + +from cuphoton.core.artifacts import file_sha256 +from cuphoton.xfit import ( + DipoleFitResult, + GaussianDipoleModel, + StampDipoleModel, + executor, + fit_dipoles, +) +from cuphoton.xfit.io import load_xfit_dataset + + +def _input(tmp_path, *, mode="difference", auxiliary="candidate"): + model = GaussianDipoleModel((7, 9), dtype=np.float64) + truth = np.array( + [ + [4, 1.2, 1.4, 0.1, 2, 2, 4, 2.5], + [5, 1.3, 1.1, -0.1, 2.1, 1.9, 4.1, 2.6], + [3, 1.1, 1.3, 0.2, 2.2, 2.2, 4.2, 2.4], + ], + dtype=np.float64, + ) + images = np.asarray(model.evaluate(truth, mode=mode)) + variance = np.full((3, 7, 9), 2.0) + variance[1] = 3.0 + mask = np.ones_like(variance, dtype=bool) + mask[1, 0, 0] = False + if mode == "split" and auxiliary == "plane": + variance = variance[None, :] + mask = mask[None, :] + elif auxiliary == "scalar": + variance = np.asarray(2.0) + mask = np.asarray(True) + path = tmp_path / "input.npz" + np.savez_compressed( + path, + candidate_id=np.array(["z", "a", "m"]), + images=images, + initial=truth + 0.01, + variance=variance, + mask=mask, + ) + return path + + +@pytest.mark.parametrize( + "model,finite_difference,message", + [ + ("stamp", False, "only the Gaussian model"), + ("gaussian", True, "does not support finite-difference fitting"), + ], +) +def test_preflight_rejects_unsupported_cutile_settings( + tmp_path, monkeypatch, model, finite_difference, message +): + def unexpected_load(*args, **kwargs): + pytest.fail("invalid settings must fail before reading input") + + monkeypatch.setattr(executor, "load_xfit_dataset", unexpected_load) + with pytest.raises(ValueError, match=message): + executor.prepare_xfit_workload( + input_path=tmp_path / "missing.npz", + fit_options={ + "backend": "cutile", + "model": model, + "use_finite_difference": finite_difference, + }, + ) + + +@pytest.mark.parametrize("backend", ["auto", "numpy"]) +def test_preflight_rejects_cpu_backend_before_input_reads(tmp_path, backend): + with pytest.raises(ValueError, match="requires backend cupy or cutile"): + executor.prepare_xfit_workload( + input_path=tmp_path / "missing.npz", + fit_options={"backend": backend}, + ) + + +def test_chunk_size_is_part_of_configuration_identity(tmp_path): + path = _input(tmp_path) + _, first = executor.plan_xfit_chunks(path, chunk_size=1, fit_options={}) + _, second = executor.plan_xfit_chunks(path, chunk_size=2, fit_options={}) + assert first["chunk_size"] == 1 + assert second["chunk_size"] == 2 + assert first["configuration_sha256"] != second["configuration_sha256"] + + +def test_finalizer_retains_planning_input_and_still_checks_content( + tmp_path, monkeypatch +): + path = _input(tmp_path) + loads = [] + original_load = executor.load_xfit_dataset + + def load(*args, **kwargs): + loads.append(True) + return original_load(*args, **kwargs) + + def cpu_fit(*args, **kwargs): + result = fit_dipoles(*args, **{**kwargs, "backend": "numpy"}) + return replace(result, backend="cupy", device="cuda:0") + + monkeypatch.setattr(executor, "load_xfit_dataset", load) + monkeypatch.setattr(executor, "fit_dipoles", cpu_fit) + monkeypatch.setattr(executor, "collect_gpu_identity", lambda _: {}) + spec = executor.prepare_xfit_workload( + input_path=path, + chunk_size=2, + fit_options={"backend": "cupy", "max_evaluations": 2}, + ) + worker = spec.worker_factory(spec.options_payload) + for index in range(2): + round_dir = tmp_path / f"round-{index}" + records = [] + for item in spec.items: + worker.run_item(item, round_dir / "items" / item.item_id) + records.append({"item_id": item.item_id, "status": "success"}) + assert spec.finalize_round(round_dir, records)["candidate_count"] == 3 + assert len(loads) == 2 # One planning load and one worker load. + before = path.stat() + changed = bytearray(path.read_bytes()) + changed[-1] ^= 1 + path.write_bytes(changed) + os.utime(path, ns=(before.st_atime_ns, before.st_mtime_ns)) + with pytest.raises(ValueError, match="archive changed after planning"): + spec.finalize_round(tmp_path / "unused", records) + assert len(loads) == 2 + worker.close() + + +def test_preflight_accepts_analytic_gaussian_cutile(tmp_path): + spec = executor.prepare_xfit_workload( + input_path=_input(tmp_path), fit_options={"backend": "cutile"} + ) + assert spec.backend == "cutile" + + +@pytest.mark.parametrize( + "runtime,rank", [("dragon", 0), ("mpi", 0), ("mpi", 1)] +) +def test_cli_dispatches_collective_preflight_and_preserves_options( + tmp_path, monkeypatch, capsys, runtime, rank +): + from types import SimpleNamespace + + from cuphoton.core import executors + from cuphoton.core.cli import run_component + + path = _input(tmp_path) + output_dir = tmp_path / "distributed" + calls = [] + planning_arrays = [] + original_load = executor.load_xfit_dataset + + def load(*args, **kwargs): + dataset = original_load(*args, **kwargs) + planning_arrays.append(weakref.ref(dataset.images)) + return dataset + + def run(**kwargs): + spec = kwargs["prepare_workload"](rank) + calls.append((kwargs, spec)) + gc.collect() + assert len(planning_arrays) == 1 + assert (planning_arrays[0]() is not None) == (rank == 0) + return SimpleNamespace( + status="success", to_dict=lambda: {"status": "success"} + ) + + monkeypatch.setattr(executors, "run_workload", run) + monkeypatch.setattr(executor, "load_xfit_dataset", load) + assert ( + run_component( + "xfit", + [ + "fit-dipoles", + "--input", + str(path), + "--model", + "gaussian", + "--backend", + "cupy", + "--executor", + runtime, + "--chunk-size", + "2", + "--output-dir", + str(output_dir), + "--max-evaluations", + "10", + "--warmup-rounds", + "1", + "--measure-rounds", + "2", + ], + ) + == 0 + ) + kwargs, spec = calls[0] + assert kwargs.get("prepare_on_root", False) == (runtime == "mpi") + assert kwargs["run_id"] == "distributed" + assert kwargs["output_root"] == tmp_path + assert kwargs["benchmark"].warmup_rounds == 1 + assert kwargs["benchmark"].measure_rounds == 2 + assert len(spec.items) == 2 + assert spec.options_payload["fit_options"]["max_evaluations"] == 10 + assert json.loads(capsys.readouterr().out)["status"] == "success" + + +@pytest.mark.parametrize( + "mode,auxiliary", + [ + ("difference", "candidate"), + ("split", "candidate"), + ("split", "plane"), + ("difference", "scalar"), + ], +) +def test_chunks_preserve_numerics_broadcasting_and_order( + tmp_path, monkeypatch, mode, auxiliary +): + path = _input(tmp_path, mode=mode, auxiliary=auxiliary) + items, options = executor.plan_xfit_chunks( + path, + chunk_size=2, + fit_options={"backend": "numpy", "mode": mode, "max_evaluations": 10}, + ) + identities = [] + monkeypatch.setattr( + executor, + "collect_gpu_identity", + lambda backend: identities.append(backend) or {"backend": "numpy"}, + ) + worker = executor.create_xfit_worker(options) + dataset = load_xfit_dataset(path, mode=mode, model="gaussian") + ordinary = fit_dipoles( + dataset.images, + initial=dataset.initial, + mask=dataset.mask, + variance=dataset.variance, + mode=mode, + model="gaussian", + backend="numpy", + config=worker.config, + ) + for round_index in range(2): + round_dir = tmp_path / f"round-{round_index}" + records = [] + for item in reversed(items): + worker.run_item(item, round_dir / "items" / item.item_id) + records.append({"item_id": item.item_id, "status": "success"}) + receipt = executor.finalize_xfit_round( + round_dir, records, items=items, options=options + ) + assert receipt["candidate_count"] == 3 + chunks = [ + executor._read_result(item, round_dir, options)[1] + for item in items + ] + for field in fields(DipoleFitResult): + if field.name in executor._META_FIELDS: + continue + actual = np.concatenate([chunk[field.name] for chunk in chunks]) + expected = np.asarray(getattr(ordinary, field.name)) + if expected.dtype.kind in "biuUS": + np.testing.assert_array_equal(actual, expected) + else: + np.testing.assert_allclose( + actual, expected, rtol=1e-10, atol=1e-12, equal_nan=True + ) + rows = pq.read_table( + round_dir / "scientific" / "fits.parquet" + ).to_pylist() + assert [row["candidate_id"] for row in rows] == ["z", "a", "m"] + assert [row["candidate_index"] for row in rows] == [0, 1, 2] + with np.load( + round_dir / "scientific" / "fit-arrays.npz", allow_pickle=False + ) as saved: + np.testing.assert_allclose( + saved["residuals"], ordinary.residuals, atol=1e-12 + ) + summary = json.loads( + (round_dir / "scientific" / "summary.json").read_text() + ) + assert ( + summary["inputs"]["input_archive_sha256"] + == dataset.input_archive_sha256 + ) + assert summary["metrics"]["candidate_count"] == 3 + assert identities == ["numpy"] + worker.close() + + +@pytest.mark.parametrize( + "defect,message", + [ + ("missing", "one successful result"), + ("duplicate", "one successful result"), + ("failed", "one successful result"), + ("corrupt", "checksum mismatch"), + ("checksum", "checksum mismatch"), + ("wrong_candidates", "candidate identity mismatch"), + ("wrong_summary", "summary identity mismatch"), + ("wrong_model", "configuration mismatch"), + ], +) +def test_merge_rejects_incomplete_or_changed_scientific_results( + tmp_path, monkeypatch, defect, message +): + path = _input(tmp_path) + items, options = executor.plan_xfit_chunks( + path, + chunk_size=2, + fit_options={"backend": "numpy", "max_evaluations": 2}, + ) + monkeypatch.setattr(executor, "collect_gpu_identity", lambda _: {}) + worker = executor.create_xfit_worker(options) + round_dir = tmp_path / "round" + records = [] + for item in items: + worker.run_item(item, round_dir / "items" / item.item_id) + records.append({"item_id": item.item_id, "status": "success"}) + if defect == "missing": + records.pop() + elif defect == "duplicate": + records[1] = dict(records[0]) + elif defect == "failed": + records[1]["status"] = "failed" + elif defect == "corrupt": + (round_dir / "items" / items[0].item_id / "result.npz").write_bytes( + b"changed" + ) + else: + item_dir = round_dir / "items" / items[0].item_id + summary_path = item_dir / "summary.json" + summary = json.loads(summary_path.read_text()) + if defect == "wrong_summary": + summary["configuration_sha256"] = "changed" + elif defect == "wrong_model": + summary["result_metadata"]["model"] = "stamp" + else: + archive = item_dir / "result.npz" + with np.load(archive, allow_pickle=False) as loaded: + values = dict(loaded) + if defect == "checksum": + values["parameters"][0, 0] += 1 + else: + values["candidate_id"] = values["candidate_id"][::-1] + np.savez_compressed(archive, **values) + if defect != "checksum": + summary["result_sha256"] = file_sha256(archive) + summary_path.write_text(json.dumps(summary)) + with pytest.raises(ValueError, match=message): + executor.finalize_xfit_round( + round_dir, records, items=items, options=options + ) + assert not (round_dir / "scientific").exists() + + +def test_retained_worker_rejects_changed_source_archive( + tmp_path, monkeypatch +): + path = _input(tmp_path) + items, options = executor.plan_xfit_chunks( + path, chunk_size=2, fit_options={"backend": "numpy"} + ) + monkeypatch.setattr(executor, "collect_gpu_identity", lambda _: {}) + worker = executor.create_xfit_worker(options) + path.write_bytes(path.read_bytes() + b"changed") + with pytest.raises(ValueError, match="input archive changed"): + worker.run_item(items[0], tmp_path / "unused") + assert not (tmp_path / "unused").exists() + + +def test_stamp_chunks_retain_shared_basis_and_forced_finite_difference( + tmp_path, monkeypatch +): + y, x = np.mgrid[-4:5, -4:5] + basis = np.exp(-(x * x + y * y) / 3.0) + model = StampDipoleModel(basis, image_shape=(7, 9)) + truth = np.array( + [[-1.1, 0.3, 1.2, -0.2, 4.0], [-1.3, 0.2, 1.4, -0.3, 3.0]] + ) + images = model.evaluate(truth) + initial = truth + 0.01 + path = tmp_path / "stamp.npz" + np.savez_compressed( + path, + candidate_id=np.array([9, 2]), + images=images, + initial=initial, + stamp_basis=basis[None, :], + ) + items, options = executor.plan_xfit_chunks( + path, + chunk_size=1, + fit_options={ + "backend": "numpy", + "model": "stamp", + "max_evaluations": 8, + }, + ) + monkeypatch.setattr(executor, "collect_gpu_identity", lambda _: {}) + worker = executor.create_xfit_worker(options) + assert worker.config.use_finite_difference is True + ordinary = fit_dipoles( + images, + model=model, + initial=initial, + backend="numpy", + config=worker.config, + ) + round_dir = tmp_path / "round" + for item in items: + worker.run_item(item, round_dir / "items" / item.item_id) + executor.finalize_xfit_round( + round_dir, + [{"item_id": item.item_id, "status": "success"} for item in items], + items=items, + options=options, + ) + chunks = [ + executor._read_result(item, round_dir, options)[1] for item in items + ] + np.testing.assert_allclose( + np.concatenate([chunk["parameters"] for chunk in chunks]), + ordinary.parameters, + rtol=1e-10, + atol=1e-12, + ) + np.testing.assert_array_equal( + np.concatenate([chunk["evaluations"] for chunk in chunks]), + ordinary.evaluations, + ) diff --git a/tests/xscan/test_cli.py b/tests/xscan/test_cli.py index 29d7e94c..89753e9c 100644 --- a/tests/xscan/test_cli.py +++ b/tests/xscan/test_cli.py @@ -9390,3 +9390,52 @@ def fake_xpois_difference( assert "non-finite" in unstable["failure"]["message"] assert "triplet_xpois_nan_variant" not in payload["jobs"] assert "triplet_xpois_nan_variant" in payload["unstable_variants"] + + +@pytest.mark.parametrize("runtime", ["dragon", "mpi"]) +def test_distributed_inference_validates_output_name_before_launch( + tmp_path, monkeypatch, capsys, runtime +): + from cuphoton.core import executors + + monkeypatch.setattr( + executors, "run_workload", lambda **kwargs: pytest.fail("launched") + ) + rc = run_component( + "xscan", + [ + "infer-real-bogus", + "--run-dir", + str(tmp_path), + "--dataset-dir", + str(tmp_path), + "--executor", + runtime, + "--output-dir", + str(tmp_path / "invalid name"), + ], + ) + assert rc != 0 + assert "--output-dir basename" in capsys.readouterr().err + + +@pytest.mark.parametrize( + "option,value", [("--task-batches", "1"), ("--output-dir", "output")] +) +def test_local_inference_rejects_distributed_options( + tmp_path, capsys, option, value +): + rc = run_component( + "xscan", + [ + "infer-real-bogus", + "--run-dir", + str(tmp_path), + "--dataset-dir", + str(tmp_path), + option, + value, + ], + ) + assert rc != 0 + assert "require --executor dragon or mpi" in capsys.readouterr().err diff --git a/tests/xscan/test_executor.py b/tests/xscan/test_executor.py new file mode 100644 index 00000000..80bd04cd --- /dev/null +++ b/tests/xscan/test_executor.py @@ -0,0 +1,610 @@ +# 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 subprocess +import sys +import time +from functools import partial +from pathlib import Path + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +from cuphoton.core import execution # noqa: E402 +from cuphoton.core.artifacts import file_sha256 # noqa: E402 +from cuphoton.core.bulk import WorkItem # noqa: E402 +from cuphoton.xscan import executor, training, workflows # noqa: E402 +from cuphoton.xscan.model import build_model # noqa: E402 + + +def _mark_loader_ready(worker_id, *, directory): + time.sleep(0.1) + (directory / f"ready-{os.getpid()}").write_text(str(worker_id)) + + +def _inputs(tmp_path): + dataset_dir = tmp_path / "dataset" + run_dir = tmp_path / "model" + dataset_dir.mkdir() + run_dir.mkdir() + rng = np.random.default_rng(19) + for name in ("search", "template"): + np.save( + dataset_dir / f"{name}.npy", + rng.normal(size=(7, 17, 17)).astype(np.float32), + ) + np.save(dataset_dir / "labels.npy", np.array([0, 1, 0, 1, 1, 0, 1])) + np.save(dataset_dir / "split.npy", np.array([2, 0, 2, 2, 2, 0, 2])) + (dataset_dir / "metadata.jsonl").write_text( + "".join( + json.dumps({"candidate_id": f"candidate-{6 - index}"}) + "\n" + for index in range(7) + ) + ) + model_config = dict( + input_mode="pair", + image_size=17, + depths=[1, 1], + num_heads=[1, 1], + embed_dims=[8, 16], + decoder_embedding_dim=8, + pos_dim=4, + output_nc=2, + drop_rate=0.0, + attn_drop=0.0, + ) + torch.manual_seed(2) + model = build_model(**model_config) + torch.save( + { + "model_config": model_config, + "model_state": model.state_dict(), + "train_config": { + "performance": {"num_workers": 0, "amp_dtype": "off"} + }, + }, + run_dir / "checkpoint.pt", + ) + return run_dir, dataset_dir + + +def _cpu_worker(monkeypatch): + monkeypatch.setattr(executor, "_device", lambda: torch.device("cpu")) + monkeypatch.setattr( + executor, + "_gpu_identity", + lambda: {"backend": "torch", "device_index": 0, "uuid": "test"}, + ) + + +@pytest.mark.parametrize("runtime", ["dragon", "mpi"]) +def test_cli_dispatches_collective_preflight_and_preserves_batches( + tmp_path, monkeypatch, capsys, runtime +): + from types import SimpleNamespace + + from cuphoton.core import executors + from cuphoton.core.cli import run_component + + run_dir, dataset_dir = _inputs(tmp_path) + output_dir = tmp_path / "distributed" + calls = [] + + def run(**kwargs): + spec = kwargs["prepare_workload"](0) + calls.append((kwargs, spec)) + return SimpleNamespace( + status="success", to_dict=lambda: {"status": "success"} + ) + + monkeypatch.setattr(executors, "run_workload", run) + assert ( + run_component( + "xscan", + [ + "infer-real-bogus", + "--run-dir", + str(run_dir), + "--dataset-dir", + str(dataset_dir), + "--executor", + runtime, + "--output-dir", + str(output_dir), + "--batch-size", + "2", + "--task-batches", + "1", + "--num-workers", + "0", + "--warmup-rounds", + "1", + "--measure-rounds", + "2", + ], + ) + == 0 + ) + kwargs, spec = calls[0] + assert kwargs.get("prepare_on_root", False) == (runtime == "mpi") + assert kwargs["run_id"] == "distributed" + assert kwargs["output_root"] == tmp_path + assert kwargs["benchmark"].warmup_rounds == 1 + assert len(spec.items) == 3 + assert spec.options_payload["batch_size"] == 2 + assert spec.options_payload["num_workers"] == 0 + assert json.loads(capsys.readouterr().out)["status"] == "success" + + +def test_cli_requires_separate_distributed_output(tmp_path, capsys): + from cuphoton.core.cli import run_component + + assert ( + run_component( + "xscan", + [ + "infer-real-bogus", + "--run-dir", + str(tmp_path), + "--dataset-dir", + str(tmp_path), + "--executor", + "mpi", + ], + ) + != 0 + ) + assert "--output-dir is required" in capsys.readouterr().err + + +@pytest.mark.parametrize("num_workers", [0, 1]) +def test_persistent_workers_preserve_original_batches_predictions_and_order( + tmp_path, monkeypatch, num_workers +): + run_dir, dataset_dir = _inputs(tmp_path) + if num_workers: + # Spawned workers must be able to import the readiness test callback. + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2])) + checkpoint = torch.load(run_dir / "checkpoint.pt", weights_only=True) + checkpoint["train_config"]["performance"]["worker_start_method"] = ( + "spawn" + ) + torch.save(checkpoint, run_dir / "checkpoint.pt") + _cpu_worker(monkeypatch) + monkeypatch.setattr( + workflows, "resolve_device", lambda _: torch.device("cpu") + ) + ordinary = workflows.infer_workflow( + run_dir=run_dir, + dataset_dir=dataset_dir, + split="test", + batch_size=2, + num_workers=num_workers, + ) + expected = { + name: np.load(ordinary.run_dir / f"{name}.npy") + for name in ("logits", "labels", "probabilities") + } + checkpoint_bytes = (run_dir / "checkpoint.pt").read_bytes() + ordinary_summary = (ordinary.run_dir / "summary.json").read_bytes() + spec = executor.prepare_inference_workload( + run_dir=run_dir, + dataset_dir=dataset_dir, + split="test", + batch_size=2, + task_batches=1, + num_workers=num_workers, + ) + items, options = spec.items, spec.options_payload + assert [ + (item.payload["start"], item.payload["stop"]) for item in items + ] == [(0, 2), (2, 4), (4, 5)] + loads = [] + original_load = training.load_model_from_checkpoint + + def load(*args, **kwargs): + loads.append(True) + return original_load(*args, **kwargs) + + monkeypatch.setattr(training, "load_model_from_checkpoint", load) + from cuphoton.xscan import dataset as dataset_module + + metadata_reads = [] + loader_creations = [] + original_metadata = dataset_module.load_metadata_rows + original_loader = training.make_dataloader + + def metadata(path): + metadata_reads.append(path) + return original_metadata(path) + + def loader(*args, **kwargs): + loader_creations.append(True) + result = original_loader(*args, **kwargs) + result.worker_init_fn = partial( + _mark_loader_ready, directory=tmp_path + ) + return result + + monkeypatch.setattr(dataset_module, "load_metadata_rows", metadata) + monkeypatch.setattr(training, "make_dataloader", loader) + workers = [executor.create_xscan_worker(options) for _ in range(2)] + children = [ + process + for worker in workers + for process in getattr(worker.loader._iterator, "_workers", []) + ] + assert len(children) == 2 * num_workers + assert len(list(tmp_path.glob("ready-*"))) == 2 * num_workers + assert all(process.is_alive() for process in children) + from multiprocessing.process import BaseProcess + + def unexpected_process_start(_): + pytest.fail("loader processes must start before READY") + + monkeypatch.setattr(BaseProcess, "start", unexpected_process_start) + + def unexpected_metadata_read(_): + pytest.fail("metadata must be retained across tasks") + + monkeypatch.setattr( + training, "load_metadata_rows", unexpected_metadata_read + ) + for round_index in range(2): + round_dir = tmp_path / f"round-{round_index}" + round_id = f"round-{round_index}" + execution.prepare_run(round_dir, round_id, spec, "test") + shards = ((items[2], items[0]), (items[1],)) + results = [ + execution.execute_worker_round( + worker, + items=shards[index], + run_id=round_id, + run_dir=round_dir, + manifest_sha256=spec.manifest_sha256, + worker_id=index, + backend="torch", + provenance={ + "worker_id": index, + "pid": os.getpid(), + "hostname": "test-host", + "cuda_visible_devices": str(index), + "gpu": { + "backend": "torch", + "device_index": 0, + "uuid": f"GPU-{index}", + "identity_error": None, + }, + }, + ) + for index, worker in enumerate(workers) + ] + audited = execution.finalize_round( + round_dir, + round_id, + spec, + shards, + results, + artifact_timeout_sec=0.1, + ) + assert audited["status"] == "success", audited["errors"] + receipt = audited["result"] + assert receipt["sample_count"] == 5 + for name, values in expected.items(): + actual = np.load(round_dir / "scientific" / f"{name}.npy") + np.testing.assert_array_equal(actual, values) + np.testing.assert_array_equal( + np.load(round_dir / "scientific" / "sample_index.npy"), + [0, 2, 3, 4, 6], + ) + summary = json.loads( + (round_dir / "scientific" / "summary.json").read_text() + ) + for key in ( + "xfit_features", + "batch_size", + "dataset_dir", + "split", + ): + assert summary[key] == ordinary.summary[key] + assert summary["performance"] == { + **ordinary.summary["performance"], + "persistent_workers": num_workers > 0, + } + assert summary["saved"]["logits"] == "logits.npy" + assert len(loads) == 2 + assert len(metadata_reads) == 2 + assert len(loader_creations) == 2 + assert all(process.is_alive() for process in children) + assert children == [ + process + for worker in workers + for process in getattr(worker.loader._iterator, "_workers", []) + ] + assert (run_dir / "checkpoint.pt").read_bytes() == checkpoint_bytes + assert ( + ordinary.run_dir / "summary.json" + ).read_bytes() == ordinary_summary + for worker in workers: + worker.close() + assert all(not process.is_alive() for process in children) + + +def test_loader_defaults_and_prediction_timer(tmp_path, monkeypatch): + from types import SimpleNamespace + + run_dir, dataset_dir = _inputs(tmp_path) + checkpoint = torch.load(run_dir / "checkpoint.pt", weights_only=True) + checkpoint["train_config"]["performance"]["num_workers"] = 2 + torch.save(checkpoint, run_dir / "checkpoint.pt") + _cpu_worker(monkeypatch) + items, options = executor.plan_inference_chunks( + run_dir=run_dir, dataset_dir=dataset_dir, split="test", batch_size=2 + ) + worker = executor.create_xscan_worker(options) + assert worker.performance.num_workers == 0 + clock = [0.0] + original_check = executor._check_inputs + original_predict = training.predict_dataset + + def check(*args, **kwargs): + clock[0] += 100 + return original_check(*args, **kwargs) + + def predict(**kwargs): + clock[0] += 7 + return original_predict(**kwargs) + + monkeypatch.setattr( + executor, "time", SimpleNamespace(perf_counter=lambda: clock[0]) + ) + monkeypatch.setattr(executor, "_check_inputs", check) + monkeypatch.setattr(training, "predict_dataset", predict) + result = worker.run_item(items[0], tmp_path / "item") + assert result["timings_sec"]["predict_sec"] == 7 + assert result["wall_sec"]["item_runner"] == 207 + worker.close() + + +@pytest.mark.parametrize("ulps", [1, 4, 5]) +def test_probability_audit_tolerates_only_bounded_host_rounding( + tmp_path, monkeypatch, ulps +): + run_dir, dataset_dir = _inputs(tmp_path) + _cpu_worker(monkeypatch) + items, options = executor.plan_inference_chunks( + run_dir=run_dir, dataset_dir=dataset_dir, split="test" + ) + worker = executor.create_xscan_worker(options) + item = items[0] + round_dir = tmp_path / "round" + item_dir = round_dir / "items" / item.item_id + worker.run_item(item, item_dir) + path = item_dir / "probabilities.npy" + values = np.load(path) + for _ in range(ulps): + values = np.nextafter(values, np.inf) + np.save(path, values) + summary_path = item_dir / "summary.json" + summary = json.loads(summary_path.read_text()) + summary["artifact_sha256"]["probabilities"] = file_sha256(path) + summary_path.write_text(json.dumps(summary)) + if ulps > 4: + with pytest.raises(ValueError, match="more than 4 ULP"): + executor._read_result(item, round_dir, options) + else: + _, arrays = executor._read_result(item, round_dir, options) + np.testing.assert_array_equal(arrays["probabilities"], values) + worker.close() + + +@pytest.mark.parametrize( + "defect,message", + [ + ("missing", "one successful result"), + ("duplicate", "one successful result"), + ("failed", "one successful result"), + ("wrong_probability", "more than 4 ULP"), + ("wrong_identity", "identities or labels differ"), + ("checksum", "checksum mismatch"), + ("wrong_labels", "identities or labels differ"), + ("short_array", "row count mismatch"), + ("divergent_summary", "differ across workers"), + ("nonfinite", "finite float64"), + ("wrong_dtype", "finite float64"), + ], +) +def test_merge_rejects_incomplete_or_invalid_scientific_results( + tmp_path, monkeypatch, defect, message +): + run_dir, dataset_dir = _inputs(tmp_path) + _cpu_worker(monkeypatch) + items, options = executor.plan_inference_chunks( + run_dir=run_dir, + dataset_dir=dataset_dir, + split="test", + batch_size=2, + task_batches=1, + ) + worker = executor.create_xscan_worker(options) + round_dir = tmp_path / "round" + records = [] + for item in items: + worker.run_item(item, round_dir / "items" / item.item_id) + records.append({"item_id": item.item_id, "status": "success"}) + if defect == "missing": + records.pop() + elif defect == "duplicate": + records[1] = dict(records[0]) + elif defect == "failed": + records[0]["status"] = "failed" + else: + item_dir = round_dir / "items" / items[0].item_id + summary_path = item_dir / "summary.json" + summary = json.loads(summary_path.read_text()) + if defect == "wrong_identity": + summary["metadata_rows"][0]["candidate_id"] = "wrong" + elif defect == "divergent_summary": + summary["inference_summary"]["performance"]["num_workers"] = 10 + else: + name = "probabilities" + values = np.array([0.123, 0.456]) + if defect == "wrong_labels": + name = "labels" + values = 1 - np.load(item_dir / "labels.npy") + elif defect == "short_array": + name = "logits" + values = np.load(item_dir / "logits.npy")[:-1] + elif defect == "nonfinite": + values[0] = np.nan + elif defect == "wrong_dtype": + values = np.load(item_dir / "probabilities.npy").astype( + np.float32 + ) + path = item_dir / f"{name}.npy" + np.save(path, values) + if defect != "checksum": + summary["artifact_sha256"][name] = file_sha256(path) + summary_path.write_text(json.dumps(summary)) + with pytest.raises(ValueError, match=message): + executor.finalize_inference_round( + round_dir, records, items=items, options=options + ) + assert not (round_dir / "scientific").exists() + + +def test_worker_rejects_split_batches_and_changed_input( + tmp_path, monkeypatch +): + run_dir, dataset_dir = _inputs(tmp_path) + _cpu_worker(monkeypatch) + items, options = executor.plan_inference_chunks( + run_dir=run_dir, + dataset_dir=dataset_dir, + split="test", + batch_size=2, + task_batches=1, + ) + worker = executor.create_xscan_worker(options) + malformed = WorkItem(items[0].item_id, {**items[0].payload, "stop": 1}) + with pytest.raises(ValueError, match="minibatch boundary"): + worker.run_item(malformed, tmp_path / "unused") + np.save(dataset_dir / "labels.npy", np.zeros(7, dtype=np.int64)) + with pytest.raises(ValueError, match="input changed"): + worker.run_item(items[0], tmp_path / "unused") + + +def test_import_and_planning_do_not_import_torch(tmp_path): + run_dir, dataset_dir = _inputs(tmp_path) + command = ( + "import sys; from pathlib import Path; " + "from cuphoton.xscan.executor import plan_inference_chunks; " + "assert 'torch' not in sys.modules; " + f"plan_inference_chunks(run_dir=Path({str(run_dir)!r}), " + f"dataset_dir=Path({str(dataset_dir)!r}), split='test'); " + "assert 'torch' not in sys.modules" + ) + subprocess.run( + [sys.executable, "-c", command], check=True, env=os.environ.copy() + ) + + +def test_feature_matrix_is_reused_with_global_coverage(tmp_path, monkeypatch): + from cuphoton.xscan.xfit_features import XFitFeatureMatrix + + run_dir, dataset_dir = _inputs(tmp_path) + _cpu_worker(monkeypatch) + monkeypatch.setattr( + workflows, "resolve_device", lambda _: torch.device("cpu") + ) + checkpoint = torch.load(run_dir / "checkpoint.pt", weights_only=True) + names = ("fit_present", "separation_over_stamp") + checkpoint["model_config"]["xfit_feature_names"] = list(names) + model = build_model(**checkpoint["model_config"]) + torch.nn.init.ones_(model.xfit_head[-1].weight) + checkpoint["model_state"] = model.state_dict() + torch.save(checkpoint, run_dir / "checkpoint.pt") + feature_dir = tmp_path / "features" + feature_dir.mkdir() + values = np.column_stack( + (np.array([1, 0, 1, 0, 1, 0, 1]), np.arange(7) / 10) + ).astype(np.float32) + for name, array in { + "candidate-id.npy": np.arange(7), + "features.npy": values, + "input-image-sha256.npy": np.zeros(7, dtype="S64"), + }.items(): + np.save(feature_dir / name, array) + (feature_dir / "schema.json").write_text("{}") + matrix = XFitFeatureMatrix( + dataset_dir, + np.arange(7), + values, + names, + {}, + {"identity": "retained"}, + {}, + ) + loads = [] + + def features(**kwargs): + assert kwargs["use_xfit_features"] is True + assert kwargs["xfit_feature_dir"] == feature_dir + loads.append(True) + return matrix + + # Bundle validation already has dedicated tests; this exercises the + # adapter's reuse and original-index slicing with real fusion inference. + monkeypatch.setattr( + workflows, "_checkpoint_xfit_feature_matrix", features + ) + ordinary = workflows.infer_workflow( + run_dir=run_dir, + dataset_dir=dataset_dir, + split="test", + batch_size=2, + use_xfit_features=True, + xfit_feature_dir=feature_dir, + ) + items, options = executor.plan_inference_chunks( + run_dir=run_dir, + dataset_dir=dataset_dir, + split="test", + batch_size=2, + task_batches=1, + use_xfit_features=True, + xfit_feature_dir=feature_dir, + ) + worker = executor.create_xscan_worker(options) + assert worker.dataset.xfit_features is matrix.values + round_dir = tmp_path / "round" + for item in items: + worker.run_item(item, round_dir / "items" / item.item_id) + executor.finalize_inference_round( + round_dir, + [ + {"item_id": item.item_id, "status": "success"} + for item in reversed(items) + ], + items=items, + options=options, + ) + assert len(loads) == 2 # One ordinary call, one retained worker. + for name in ("logits", "probabilities", "labels"): + np.testing.assert_array_equal( + np.load(round_dir / "scientific" / f"{name}.npy"), + np.load(ordinary.run_dir / f"{name}.npy"), + ) + summary = json.loads( + (round_dir / "scientific" / "summary.json").read_text() + ) + assert summary["xfit_features"] == ordinary.summary["xfit_features"] + assert summary["xfit_features"]["fit_coverage"]["sample_count"] == 5 + assert summary["xfit_features"]["fit_coverage"]["fit_coverage"] == 0.8 diff --git a/tests/xscan/test_pipeline_executor.py b/tests/xscan/test_pipeline_executor.py new file mode 100644 index 00000000..b13291ca --- /dev/null +++ b/tests/xscan/test_pipeline_executor.py @@ -0,0 +1,375 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import numpy as np +import pytest + +from cuphoton.core.artifacts import file_sha256 +from cuphoton.core.bulk import WorkItem +from cuphoton.core.cli import run_component +from cuphoton.xscan import pipeline_executor as pipeline +from cuphoton.xscan.device_pipeline import ( + DevicePipelineCandidate, + DevicePipelineConfig, + DevicePipelineItem, + DeviceWorkerContext, + DeviceXPOISPipelineConfig, + NpyArrayDescriptor, +) + + +def _manifest(tmp_path): + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "checkpoint.pt").write_bytes(b"checkpoint") + schema = tmp_path / "schema.json" + schema.write_text("{}") + image = tmp_path / "image.npy" + array = np.ones((96, 96), dtype=np.float64) + np.save(image, array) + descriptor = NpyArrayDescriptor( + path=str(image), + sha256=file_sha256(image), + shape=array.shape, + dtype=array.dtype.str, + ) + config = DevicePipelineConfig( + device="cuda:0", + checkpoint_dir=str(checkpoint), + checkpoint_sha256=file_sha256(checkpoint / "checkpoint.pt"), + feature_schema_path=str(schema), + feature_schema_sha256=file_sha256(schema), + stamp_shape=(63, 63), + decision_threshold=0.5, + xpois=DeviceXPOISPipelineConfig( + kernel_shape=(3, 3), + basis_sigmas=(1.0,), + basis_degrees=(0,), + ), + ) + item = DevicePipelineItem( + item_id="pair", + reference=descriptor, + target=descriptor, + candidates=( + DevicePipelineCandidate( + candidate_id="candidate", + center_x=47, + center_y=47, + source_index=0, + ), + ), + ) + payload = { + "schema": pipeline.PIPELINE_MANIFEST_SCHEMA, + "configuration": config.to_payload(), + "items": [item.to_payload()], + } + path = tmp_path / "manifest.json" + path.write_text(json.dumps(payload)) + return path, config, item + + +def test_manifest_paths_are_relative_to_manifest_not_cwd( + tmp_path, monkeypatch +): + path, config, item = _manifest(tmp_path) + data = json.loads(path.read_text()) + data["configuration"]["checkpoint_dir"] = "checkpoint" + data["configuration"]["feature_schema_path"] = "schema.json" + for role in ("reference", "target"): + data["items"][0][role]["path"] = "image.npy" + path.write_text(json.dumps(data)) + monkeypatch.chdir(tmp_path.parent) + assert pipeline.load_pipeline_manifest(path) == (config, (item,)) + + +@pytest.mark.parametrize("directory_symlink", [False, True]) +def test_manifest_preserves_npy_symlink_and_validates_target( + tmp_path, directory_symlink +): + path, config, _ = _manifest(tmp_path) + image = tmp_path / "image.npy" + directory = tmp_path + input_path = "image.npy" + if directory_symlink: + directory = tmp_path / "store" + (directory / "nested").mkdir(parents=True) + (tmp_path / "inputs").symlink_to( + directory / "nested", target_is_directory=True + ) + input_path = "inputs/../image.npy" + target = directory / "content-addressed-input" + image.rename(target) + image = directory / "image.npy" + image.symlink_to(target.name) + data = json.loads(path.read_text()) + for role in ("reference", "target"): + data["items"][0][role]["path"] = input_path + path.write_text(json.dumps(data)) + + actual_config, items = pipeline.load_pipeline_manifest(path) + assert actual_config == config + assert items[0].reference.path == str(image) + assert items[0].target.path == str(image) + pipeline.prepare_pipeline_workload(items, config) + + target.write_bytes(target.read_bytes()[:-1] + b"x") + with pytest.raises(RuntimeError, match="changed before"): + pipeline.prepare_pipeline_workload(items, config) + + +@pytest.mark.parametrize( + "defect", ["duplicate", "unknown", "schema", "items"] +) +def test_manifest_rejects_ambiguous_or_invalid_descriptors(tmp_path, defect): + path, _, _ = _manifest(tmp_path) + data = json.loads(path.read_text()) + if defect == "unknown": + data["unexpected"] = 1 + elif defect == "schema": + data["schema"] = "wrong/v1" + elif defect == "items": + data["items"] = {} + path.write_text(json.dumps(data)) + if defect == "duplicate": + path.write_text( + path.read_text().replace('"items":', '"items": [], "items":') + ) + with pytest.raises(ValueError): + pipeline.load_pipeline_manifest(path) + + +@pytest.mark.parametrize("device", ["cpu", "cuda:1"]) +def test_manifest_rejects_device_before_item_preflight( + tmp_path, monkeypatch, device +): + path, _, _ = _manifest(tmp_path) + data = json.loads(path.read_text()) + data["configuration"]["device"] = device + path.write_text(json.dumps(data)) + monkeypatch.setattr( + DevicePipelineItem, + "from_payload", + lambda _: pytest.fail("items parsed before device validation"), + ) + with pytest.raises(ValueError, match="cuda:0"): + pipeline.load_pipeline_manifest(path) + + +def test_preflight_agrees_across_ranks_and_hashes_only_on_root( + tmp_path, monkeypatch +): + _, config, item = _manifest(tmp_path) + root = pipeline.prepare_pipeline_workload((item,), config) + + def unexpected_hash(path): + pytest.fail(f"non-root preflight hashed {path}") + + monkeypatch.setattr( + "cuphoton.xscan.dragon_pipeline.file_sha256", unexpected_hash + ) + peer = pipeline.prepare_pipeline_workload((item,), config, rank=1) + assert root.identity_payload() == peer.identity_payload() + assert peer.worker_factory is pipeline.create_pipeline_worker + + +def test_preflight_rejects_changed_input_on_root(tmp_path): + _, config, item = _manifest(tmp_path) + image = tmp_path / "image.npy" + image.write_bytes(image.read_bytes()[:-1] + b"x") + with pytest.raises(RuntimeError, match="changed before"): + pipeline.prepare_pipeline_workload((item,), config) + + +@pytest.mark.parametrize("executor", ["dragon", "mpi"]) +def test_entrypoint_prepares_inside_runtime_callback( + tmp_path, monkeypatch, executor +): + path, config, items = _manifest(tmp_path) + seen = {} + + def run(**kwargs): + seen.update(kwargs) + spec = kwargs["prepare_workload"](1) + assert spec.items[0].item_id == items.item_id + assert spec.options_payload["configuration"] == config.to_payload() + return "result" + + monkeypatch.setattr("cuphoton.core.executors.run_workload", run) + assert ( + pipeline.run_pipeline_manifest( + executor=executor, + manifest_path=path, + output_root=tmp_path / "runs", + run_id="example", + ) + == "result" + ) + assert seen["executor"] == executor + assert seen["run_id"] == "example" + + +def test_worker_reuses_context_and_cleans_both_streams(tmp_path, monkeypatch): + _, config, item = _manifest(tmp_path) + events = [] + context = SimpleNamespace( + producer_stream=SimpleNamespace( + synchronize=lambda: events.append("producer") + ), + consumer_stream=SimpleNamespace( + synchronize=lambda: events.append("consumer") + ), + ) + torch = SimpleNamespace( + cuda=SimpleNamespace(set_device=lambda n: events.append(("torch", n))) + ) + cupy = SimpleNamespace( + cuda=SimpleNamespace( + Device=lambda n: SimpleNamespace( + use=lambda: events.append(("cupy", n)) + ) + ) + ) + + def import_runtime(name): + events.append(name) + return {"torch": torch, "cupy": cupy}[name] + + monkeypatch.setattr(pipeline.importlib, "import_module", import_runtime) + monkeypatch.setattr( + pipeline, "_strict_worker_options", lambda options: config + ) + monkeypatch.setattr( + pipeline, "collect_gpu_identity", lambda backend: {"backend": backend} + ) + + def initialize(cls, seen): + assert seen == config + events.append("initialize") + return context + + monkeypatch.setattr( + DeviceWorkerContext, "initialize", classmethod(initialize) + ) + + def publish(**kwargs): + assert kwargs["context"] is context + events.append("item") + return {"status": "success"} + + monkeypatch.setattr(pipeline, "_publish_item_result", publish) + worker = pipeline.create_pipeline_worker({}) + work = WorkItem(item.item_id, item.to_payload()) + for index in range(3): + worker.run_item(work, tmp_path / str(index)) + worker.close() + assert events == [ + "torch", + "cupy", + ("torch", 0), + ("cupy", 0), + "initialize", + "item", + "item", + "item", + "producer", + "consumer", + ] + with pytest.raises(RuntimeError, match="closed"): + worker.run_item(work, tmp_path / "closed") + + +def test_close_attempts_second_stream_after_first_failure(): + calls = [] + + def fail(): + calls.append("producer") + raise RuntimeError("cannot synchronize") + + worker = object.__new__(pipeline._PipelineWorker) + worker.context = SimpleNamespace( + producer_stream=SimpleNamespace(synchronize=fail), + consumer_stream=SimpleNamespace( + synchronize=lambda: calls.append("consumer") + ), + ) + with pytest.raises(RuntimeError, match="cleanup"): + worker.close() + assert calls == ["producer", "consumer"] + assert worker.closed + + +@pytest.mark.parametrize("executor", ["dragon", "mpi"]) +def test_pipeline_command_forwards_execution_options( + tmp_path, monkeypatch, executor, capsys +): + seen = {} + + def run(**kwargs): + seen.update(kwargs) + return SimpleNamespace( + status="success", to_dict=lambda: {"executor": executor} + ) + + monkeypatch.setattr(pipeline, "run_pipeline_manifest", run) + assert ( + run_component( + "xscan", + [ + "run-pipeline", + "--executor", + executor, + "--manifest", + str(tmp_path / "m.json"), + "--output-dir", + str(tmp_path / "runs"), + "--name", + "qualified", + "--warmup-rounds", + "1", + "--measure-rounds", + "3", + ], + ) + == 0 + ) + assert json.loads(capsys.readouterr().out) == {"executor": executor} + assert seen["run_id"] == "qualified" + assert seen["benchmark"].to_payload() == { + "warmup_rounds": 1, + "measure_rounds": 3, + } + + +def test_pipeline_command_rejects_wrong_executor_limits_before_runtime( + tmp_path, monkeypatch +): + monkeypatch.setattr( + pipeline, + "run_pipeline_manifest", + lambda **kwargs: pytest.fail("invalid args reached runtime"), + ) + assert ( + run_component( + "xscan", + [ + "run-pipeline", + "--executor", + "mpi", + "--max-workers", + "2", + "--manifest", + str(tmp_path / "m.json"), + "--output-dir", + str(tmp_path), + ], + ) + != 0 + )