From 30fc28ccd02ba5bef1241f18ec558ff5200b21e7 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 23 Sep 2026 21:31:18 -0700 Subject: [PATCH 01/10] Add candidate-sharded xFit execution adapter Signed-off-by: Trent Nelson --- src/cuphoton/xfit/executor.py | 509 ++++++++++++++++++++++++++++++++++ tests/xfit/test_executor.py | 251 +++++++++++++++++ 2 files changed, 760 insertions(+) create mode 100644 src/cuphoton/xfit/executor.py create mode 100644 tests/xfit/test_executor.py diff --git a/src/cuphoton/xfit/executor.py b/src/cuphoton/xfit/executor.py new file mode 100644 index 00000000..18a1d712 --- /dev/null +++ b/src/cuphoton/xfit/executor.py @@ -0,0 +1,509 @@ +# 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.""" + + 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, + "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),) + ) + return items, options + + +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], +) -> 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" + ) + dataset = _load_input(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] +): + """Preflight one standalone fit for either shared GPU executor.""" + from cuphoton.core.execution import WorkloadSpec + + items, options = plan_xfit_chunks( + input_path, chunk_size=chunk_size, fit_options=fit_options + ) + backend = options["fit_options"]["backend"] + if backend not in {"cupy", "cutile"}: + raise ValueError("distributed xFit requires backend cupy or cutile") + 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 + ), + ) diff --git a/tests/xfit/test_executor.py b/tests/xfit/test_executor.py new file mode 100644 index 00000000..12c3a1ff --- /dev/null +++ b/tests/xfit/test_executor.py @@ -0,0 +1,251 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from dataclasses import fields + +import numpy as np +import pyarrow.parquet as pq +import pytest + +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( + "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", ["missing", "duplicate", "failed", "corrupt"] +) +def test_merge_rejects_incomplete_or_changed_scientific_results( + tmp_path, monkeypatch, defect +): + 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" + else: + (round_dir / "items" / items[0].item_id / "result.npz").write_bytes( + b"changed" + ) + with pytest.raises(ValueError): + 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, + ) From f82ea8cfcadacc7847680af849f06ec7dced32fc Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 23 Sep 2026 21:31:23 -0700 Subject: [PATCH 02/10] Add persistent xScan inference execution adapter Signed-off-by: Trent Nelson --- src/cuphoton/xscan/executor.py | 519 +++++++++++++++++++++++++++++++++ tests/xscan/test_executor.py | 363 +++++++++++++++++++++++ 2 files changed, 882 insertions(+) create mode 100644 src/cuphoton/xscan/executor.py create mode 100644 tests/xscan/test_executor.py diff --git a/src/cuphoton/xscan/executor.py b/src/cuphoton/xscan/executor.py new file mode 100644 index 00000000..83452a0a --- /dev/null +++ b/src/cuphoton/xscan/executor.py @@ -0,0 +1,519 @@ +# 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": 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 .training import ( + load_model_from_checkpoint, + 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 + ) + if options["num_workers"] is not None: + self.performance = normalize_performance_config( + replace(self.performance, num_workers=options["num_workers"]), + 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.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) + + 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" + ) + payload = predict_dataset( + model=self.model, + dataset=dataset, + batch_size=self.options["batch_size"], + device=self.device, + performance=self.performance, + ) + _check_inputs(self.options) + prediction_seconds = time.perf_counter() - started + 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.""" + self.model = None + self.dataset = 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") + if not np.all(np.isfinite(arrays["logits"])) or not np.array_equal( + arrays["probabilities"], sigmoid(arrays["logits"]) + ): + raise ValueError( + "XScan chunk probabilities differ from the original host sigmoid" + ) + 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().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/tests/xscan/test_executor.py b/tests/xscan/test_executor.py new file mode 100644 index 00000000..e3e7b175 --- /dev/null +++ b/tests/xscan/test_executor.py @@ -0,0 +1,363 @@ +# 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 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 _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"}, + ) + + +def test_persistent_workers_preserve_original_batches_predictions_and_order( + tmp_path, monkeypatch +): + run_dir, dataset_dir = _inputs(tmp_path) + _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=0, + ) + 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=0, + ) + 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) + workers = [executor.create_xscan_worker(options) for _ in range(2)] + 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 ( + "performance", + "xfit_features", + "batch_size", + "dataset_dir", + "split", + ): + assert summary[key] == ordinary.summary[key] + assert summary["saved"]["logits"] == "logits.npy" + assert len(loads) == 2 + 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() + + +@pytest.mark.parametrize( + "defect", + ["missing", "duplicate", "failed", "wrong_probability", "wrong_identity"], +) +def test_merge_rejects_incomplete_or_invalid_scientific_results( + tmp_path, monkeypatch, defect +): + 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_probability": + np.save(item_dir / "probabilities.npy", np.array([0.123, 0.456])) + summary["artifact_sha256"]["probabilities"] = file_sha256( + item_dir / "probabilities.npy" + ) + else: + summary["metadata_rows"][0]["candidate_id"] = "wrong" + summary_path.write_text(json.dumps(summary)) + with pytest.raises(ValueError): + 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 From 77dfa21fcd46358697c0d857ee80405167d9cc62 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 23 Sep 2026 21:35:48 -0700 Subject: [PATCH 03/10] Expose the device pipeline through Dragon and MPI commands Signed-off-by: Trent Nelson --- docs/components/xscan.md | 42 +++ src/cuphoton/xscan/commands.py | 49 ++++ src/cuphoton/xscan/dragon_pipeline.py | 22 +- src/cuphoton/xscan/pipeline_executor.py | 215 ++++++++++++++++ tests/core/test_cli_contract.py | 10 +- tests/xscan/test_pipeline_executor.py | 323 ++++++++++++++++++++++++ 6 files changed, 652 insertions(+), 9 deletions(-) create mode 100644 src/cuphoton/xscan/pipeline_executor.py create mode 100644 tests/xscan/test_pipeline_executor.py diff --git a/docs/components/xscan.md b/docs/components/xscan.md index 54493928..1d1e06ef 100644 --- a/docs/components/xscan.md +++ b/docs/components/xscan.md @@ -356,6 +356,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 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/src/cuphoton/xscan/commands.py b/src/cuphoton/xscan/commands.py index 51cadc8e..f8b79c91 100644 --- a/src/cuphoton/xscan/commands.py +++ b/src/cuphoton/xscan/commands.py @@ -20,6 +20,7 @@ SetInvariant, StringInvariant, ) +from cuphoton.core.cli.executor import ExecutorOptions T = TypeVar("T") _DEFAULT_REVIEW_DIR = os.environ.get("CUPHOTON_XSCAN_REVIEW_DIR") @@ -888,6 +889,54 @@ def run(self) -> None: self._emit_json({"run_dir": str(result.run_dir), **result.summary}) +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, started by the 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(XScanCommand): """Run inference for a trained real-bogus model on one split.""" 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/pipeline_executor.py b/src/cuphoton/xscan/pipeline_executor.py new file mode 100644 index 00000000..d9196c85 --- /dev/null +++ b/src/cuphoton/xscan/pipeline_executor.py @@ -0,0 +1,215 @@ +# 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() + return str((path.parent / value).resolve()) + + 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 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/tests/core/test_cli_contract.py b/tests/core/test_cli_contract.py index 4be2b23d..90f94d28 100644 --- a/tests/core/test_cli_contract.py +++ b/tests/core/test_cli_contract.py @@ -73,15 +73,15 @@ def test_public_command_surface_counts_are_exact() -> None: ("xdr", 1, 0, 14, 0), ("xfit", 3, 3, 17, 1), ("xpois", 7, 7, 146, 1), - ("xscan", 42, 42, 171, 1), + ("xscan", 43, 43, 181, 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) == 851 def test_public_registry_order_and_component_derivations() -> None: diff --git a/tests/xscan/test_pipeline_executor.py b/tests/xscan/test_pipeline_executor.py new file mode 100644 index 00000000..d8c796db --- /dev/null +++ b/tests/xscan/test_pipeline_executor.py @@ -0,0 +1,323 @@ +# 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( + "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) + + +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 + ) From ebad58f8301ede64c0e9abc88ff709b4995f3644 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 23 Sep 2026 21:36:44 -0700 Subject: [PATCH 04/10] Expose distributed fitting and inference commands Signed-off-by: Trent Nelson --- docs/components/xfit.md | 40 +++++++++++++++++ docs/components/xscan.md | 48 +++++++++++++++++++++ src/cuphoton/xfit/commands.py | 55 ++++++++++++++++++++++- src/cuphoton/xscan/commands.py | 56 +++++++++++++++++++++++- tests/xfit/test_executor.py | 58 +++++++++++++++++++++++++ tests/xscan/test_executor.py | 79 ++++++++++++++++++++++++++++++++++ 6 files changed, 334 insertions(+), 2 deletions(-) diff --git a/docs/components/xfit.md b/docs/components/xfit.md index e1468252..13fcfdb6 100644 --- a/docs/components/xfit.md +++ b/docs/components/xfit.md @@ -142,6 +142,46 @@ 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. + +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. 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 1d1e06ef..3b434f41 100644 --- a/docs/components/xscan.md +++ b/docs/components/xscan.md @@ -236,6 +236,54 @@ 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. + +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`. 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 diff --git a/src/cuphoton/xfit/commands.py b/src/cuphoton/xfit/commands.py index ad79edb3..f75397a9 100644 --- a/src/cuphoton/xfit/commands.py +++ b/src/cuphoton/xfit/commands.py @@ -22,6 +22,7 @@ SetInvariant, StringInvariant, ) +from cuphoton.core.cli.executor import ExecutorOptions from ._types import ( BACKEND_REQUESTS, @@ -186,7 +187,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 +200,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 +291,51 @@ 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 + + output_dir = Path(self.output_dir).expanduser().resolve() + 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, + ), + 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/xscan/commands.py b/src/cuphoton/xscan/commands.py index f8b79c91..2bfac36d 100644 --- a/src/cuphoton/xscan/commands.py +++ b/src/cuphoton/xscan/commands.py @@ -937,7 +937,7 @@ def run(self) -> None: ) -class InferRealBogusCommand(XScanCommand): +class InferRealBogusCommand(ExecutorOptions, XScanCommand): """Run inference for a trained real-bogus model on one split.""" run_dir = None @@ -947,6 +947,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" @@ -992,6 +1008,44 @@ 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" + ) + output_dir = Path(self.output_dir).expanduser().resolve() + 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/tests/xfit/test_executor.py b/tests/xfit/test_executor.py index 12c3a1ff..70d66003 100644 --- a/tests/xfit/test_executor.py +++ b/tests/xfit/test_executor.py @@ -54,6 +54,64 @@ def _input(tmp_path, *, mode="difference", auxiliary="candidate"): return path +@pytest.mark.parametrize("runtime", ["dragon", "mpi"]) +def test_cli_dispatches_collective_preflight_and_preserves_options( + tmp_path, monkeypatch, capsys, runtime +): + 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 = [] + + 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( + "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["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", [ diff --git a/tests/xscan/test_executor.py b/tests/xscan/test_executor.py index e3e7b175..5b86bc5a 100644 --- a/tests/xscan/test_executor.py +++ b/tests/xscan/test_executor.py @@ -76,6 +76,85 @@ def _cpu_worker(monkeypatch): ) +@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["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 + + def test_persistent_workers_preserve_original_batches_predictions_and_order( tmp_path, monkeypatch ): From c7fe4577f7d8a688a8ef5ec46172b121278491e6 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 23 Sep 2026 21:39:44 -0700 Subject: [PATCH 05/10] Update the distributed command surface and launch examples Signed-off-by: Trent Nelson --- docs/components/xscan.md | 2 +- tests/core/test_cli_contract.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/components/xscan.md b/docs/components/xscan.md index 3b434f41..213873d9 100644 --- a/docs/components/xscan.md +++ b/docs/components/xscan.md @@ -426,7 +426,7 @@ 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 cuphoton-openmpi-rank-exec \ +mpiexec -n 8 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 diff --git a/tests/core/test_cli_contract.py b/tests/core/test_cli_contract.py index 90f94d28..74d369dd 100644 --- a/tests/core/test_cli_contract.py +++ b/tests/core/test_cli_contract.py @@ -71,9 +71,9 @@ 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", 43, 43, 181, 1), + ("xscan", 43, 43, 190, 1), ("xrep", 6, 6, 103, 1), ("xray", 33, 31, 390, 1), ] @@ -81,7 +81,7 @@ def test_public_command_surface_counts_are_exact() -> None: 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) == 851 + assert sum(item[3] for item in per_group) == 868 def test_public_registry_order_and_component_derivations() -> None: From 9f65f751c13f84aa127b6ef4c590af592bb9861c Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 23 Sep 2026 21:48:53 -0700 Subject: [PATCH 06/10] Reject unsupported cuTile fitting settings before launch Signed-off-by: Trent Nelson --- src/cuphoton/xfit/executor.py | 9 +++++++++ tests/xfit/test_executor.py | 36 +++++++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/src/cuphoton/xfit/executor.py b/src/cuphoton/xfit/executor.py index 18a1d712..9e1cc0ed 100644 --- a/src/cuphoton/xfit/executor.py +++ b/src/cuphoton/xfit/executor.py @@ -477,6 +477,15 @@ def prepare_xfit_workload( backend = options["fit_options"]["backend"] if backend not in {"cupy", "cutile"}: raise ValueError("distributed xFit requires backend cupy or cutile") + settings = options["fit_options"] + 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" + ) by_id = {item.item_id: item for item in items} def validate(record, round_dir): diff --git a/tests/xfit/test_executor.py b/tests/xfit/test_executor.py index 70d66003..d7d63be3 100644 --- a/tests/xfit/test_executor.py +++ b/tests/xfit/test_executor.py @@ -54,6 +54,42 @@ def _input(tmp_path, *, mode="difference", auxiliary="candidate"): 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, model, finite_difference, message +): + path = _input(tmp_path) + if model == "stamp": + np.savez_compressed( + path, + candidate_id=np.array(["stamp"]), + images=np.zeros((1, 7, 9)), + stamp_basis=np.ones((7, 9)), + ) + with pytest.raises(ValueError, match=message): + executor.prepare_xfit_workload( + input_path=path, + fit_options={ + "backend": "cutile", + "model": model, + "use_finite_difference": finite_difference, + }, + ) + + +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", ["dragon", "mpi"]) def test_cli_dispatches_collective_preflight_and_preserves_options( tmp_path, monkeypatch, capsys, runtime From acc18c0bd9bcb2576aa04cc33090290b5a3c94f0 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Thu, 24 Sep 2026 14:15:20 -0700 Subject: [PATCH 07/10] Retain inference loaders and tolerate host sigmoid rounding Signed-off-by: Trent Nelson --- docs/cli.md | 9 ++ docs/components/xfit.md | 6 +- docs/components/xscan.md | 14 ++- docs/data-artifacts.md | 6 + src/cuphoton/xfit/commands.py | 6 + src/cuphoton/xfit/executor.py | 38 +++++-- src/cuphoton/xscan/commands.py | 15 ++- src/cuphoton/xscan/executor.py | 70 ++++++++++-- src/cuphoton/xscan/training.py | 26 +++-- tests/core/test_mpi_adapters.py | 172 ++++++++++++++++++++++++++++ tests/xfit/test_commands.py | 52 +++++++++ tests/xfit/test_executor.py | 119 ++++++++++++++++--- tests/xscan/test_cli.py | 49 ++++++++ tests/xscan/test_executor.py | 195 +++++++++++++++++++++++++++++--- 14 files changed, 716 insertions(+), 61 deletions(-) create mode 100644 tests/core/test_mpi_adapters.py 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 13fcfdb6..1be478d6 100644 --- a/docs/components/xfit.md +++ b/docs/components/xfit.md @@ -151,6 +151,8 @@ 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: @@ -176,7 +178,9 @@ mpirun -n 8 --map-by slot --bind-to none -x CUDA_VISIBLE_DEVICES \ --warmup-rounds 1 --measure-rounds 2 ``` -The output directory must be new. Each pass retains normal xFit artifacts +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 diff --git a/docs/components/xscan.md b/docs/components/xscan.md index 213873d9..4f0c90e5 100644 --- a/docs/components/xscan.md +++ b/docs/components/xscan.md @@ -251,6 +251,14 @@ 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: @@ -277,7 +285,9 @@ mpirun -n 8 --map-by slot --bind-to none -x CUDA_VISIBLE_DEVICES \ --warmup-rounds 1 --measure-rounds 2 ``` -Distributed inference requires a new `--output-dir`. Each pass writes merged +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 @@ -426,7 +436,7 @@ 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 cuphoton-openmpi-rank-exec -- \ +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 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 f75397a9..c0d872ae 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, @@ -298,6 +299,11 @@ def run(self) -> None: from .executor import prepare_xfit_workload 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 ( diff --git a/src/cuphoton/xfit/executor.py b/src/cuphoton/xfit/executor.py index 9e1cc0ed..c43d3372 100644 --- a/src/cuphoton/xfit/executor.py +++ b/src/cuphoton/xfit/executor.py @@ -77,6 +77,16 @@ def plan_xfit_chunks( ) -> 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) @@ -99,6 +109,7 @@ def plan_xfit_chunks( "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" @@ -128,7 +139,10 @@ def plan_xfit_chunks( for start in range(0, dataset.batch_size, chunk_size) for stop in (min(start + chunk_size, dataset.batch_size),) ) - return items, options + 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: @@ -407,6 +421,7 @@ def finalize_xfit_round( *, 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} @@ -419,7 +434,14 @@ def finalize_xfit_round( raise ValueError( "xFit merge requires one successful result per chunk" ) - dataset = _load_input(options) + 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 @@ -471,13 +493,10 @@ def prepare_xfit_workload( """Preflight one standalone fit for either shared GPU executor.""" from cuphoton.core.execution import WorkloadSpec - items, options = plan_xfit_chunks( - input_path, chunk_size=chunk_size, fit_options=fit_options - ) - backend = options["fit_options"]["backend"] + settings = {**_DEFAULTS, **fit_options} + backend = settings["backend"] if backend not in {"cupy", "cutile"}: raise ValueError("distributed xFit requires backend cupy or cutile") - settings = options["fit_options"] if backend == "cutile" and settings["model"] == "stamp": raise ValueError( "backend='cutile' currently supports only the Gaussian model" @@ -486,6 +505,9 @@ def prepare_xfit_workload( 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 + ) by_id = {item.item_id: item for item in items} def validate(record, round_dir): @@ -513,6 +535,6 @@ def validate(record, round_dir): 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 + round_dir, records, items=items, options=options, _dataset=dataset ), ) diff --git a/src/cuphoton/xscan/commands.py b/src/cuphoton/xscan/commands.py index 2bfac36d..2a44d411 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, @@ -900,7 +901,9 @@ class ExecutorArg(ExecutorOptions.ExecutorArg): _set = {"dragon", "mpi"} _default = None _mandatory = True - _help = "Distributed runtime, started by the matching launcher." + _help = ( + "Distributed runtime: dragon or mpi, with its matching launcher." + ) class ManifestArg(PathSpecInvariant): _arg = "--manifest" @@ -988,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 @@ -1019,6 +1025,11 @@ def run(self) -> None: "--output-dir is required for distributed inference" ) 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, diff --git a/src/cuphoton/xscan/executor.py b/src/cuphoton/xscan/executor.py index 83452a0a..5e802d22 100644 --- a/src/cuphoton/xscan/executor.py +++ b/src/cuphoton/xscan/executor.py @@ -123,7 +123,7 @@ def plan_inference_chunks( "split": split, "batch_size": batch_size, "task_batches": task_batches, - "num_workers": num_workers, + "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 @@ -202,8 +202,10 @@ 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, ) @@ -219,11 +221,14 @@ def __init__(self, options: Mapping[str, Any]) -> None: self.model, checkpoint, self.performance = load_model_from_checkpoint( Path(options["run_dir"]), device=self.device ) - if options["num_workers"] is not None: - self.performance = normalize_performance_config( - replace(self.performance, num_workers=options["num_workers"]), - 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 @@ -248,6 +253,7 @@ def __init__(self, options: Mapping[str, Any]) -> None: 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"], @@ -269,6 +275,24 @@ def __init__(self, options: Mapping[str, Any]) -> 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 @@ -285,15 +309,19 @@ def run_item(self, item: WorkItem, item_dir: Path) -> Mapping[str, Any]: 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) - prediction_seconds = time.perf_counter() - started item_dir.mkdir(parents=True, exist_ok=False) arrays = { name: payload[name] @@ -331,8 +359,13 @@ def run_item(self, item: WorkItem, item_dir: Path) -> Mapping[str, Any]: 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: @@ -381,12 +414,25 @@ def _read_result(item: WorkItem, round_dir: Path, options: Mapping[str, Any]): != item.payload["sample_indices_sha256"] ): raise ValueError("XScan chunk sample identity mismatch") - if not np.all(np.isfinite(arrays["logits"])) or not np.array_equal( - arrays["probabilities"], sigmoid(arrays["logits"]) + 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 probabilities differ from the original host sigmoid" + 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 @@ -412,7 +458,7 @@ def finalize_inference_round( metadata_path = Path(options["dataset_dir"]) / "metadata.jsonl" metadata = [ json.loads(line) - for line in metadata_path.read_text().splitlines() + for line in metadata_path.read_text(encoding="utf-8").splitlines() if line.strip() ] labels = np.load( 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_mpi_adapters.py b/tests/core/test_mpi_adapters.py new file mode 100644 index 00000000..cf1cfad8 --- /dev/null +++ b/tests/core/test_mpi_adapters.py @@ -0,0 +1,172 @@ +# 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] + + 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) + + 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): + return executor.prepare_xfit_workload( + input_path=path, + chunk_size=2, + fit_options={"backend": "cupy", "max_evaluations": 2}, + ) + + 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): + 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, + 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 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 index d7d63be3..1c7bbbeb 100644 --- a/tests/xfit/test_executor.py +++ b/tests/xfit/test_executor.py @@ -5,12 +5,14 @@ from __future__ import annotations import json -from dataclasses import fields +import os +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, @@ -62,19 +64,15 @@ def _input(tmp_path, *, mode="difference", auxiliary="candidate"): ], ) def test_preflight_rejects_unsupported_cutile_settings( - tmp_path, model, finite_difference, message + tmp_path, monkeypatch, model, finite_difference, message ): - path = _input(tmp_path) - if model == "stamp": - np.savez_compressed( - path, - candidate_id=np.array(["stamp"]), - images=np.zeros((1, 7, 9)), - stamp_basis=np.ones((7, 9)), - ) + 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=path, + input_path=tmp_path / "missing.npz", fit_options={ "backend": "cutile", "model": model, @@ -83,6 +81,67 @@ def test_preflight_rejects_unsupported_cutile_settings( ) +@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"} @@ -233,10 +292,20 @@ def test_chunks_preserve_numerics_broadcasting_and_order( @pytest.mark.parametrize( - "defect", ["missing", "duplicate", "failed", "corrupt"] + "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 + tmp_path, monkeypatch, defect, message ): path = _input(tmp_path) items, options = executor.plan_xfit_chunks( @@ -257,11 +326,31 @@ def test_merge_rejects_incomplete_or_changed_scientific_results( records[1] = dict(records[0]) elif defect == "failed": records[1]["status"] = "failed" - else: + elif defect == "corrupt": (round_dir / "items" / items[0].item_id / "result.npz").write_bytes( b"changed" ) - with pytest.raises(ValueError): + 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 ) 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 index 5b86bc5a..02403436 100644 --- a/tests/xscan/test_executor.py +++ b/tests/xscan/test_executor.py @@ -8,6 +8,9 @@ import os import subprocess import sys +import time +from functools import partial +from pathlib import Path import numpy as np import pytest @@ -21,6 +24,11 @@ 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" @@ -155,10 +163,19 @@ def test_cli_requires_separate_distributed_output(tmp_path, capsys): 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 + 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") @@ -168,7 +185,7 @@ def test_persistent_workers_preserve_original_batches_predictions_and_order( dataset_dir=dataset_dir, split="test", batch_size=2, - num_workers=0, + num_workers=num_workers, ) expected = { name: np.load(ordinary.run_dir / f"{name}.npy") @@ -182,7 +199,7 @@ def test_persistent_workers_preserve_original_batches_predictions_and_order( split="test", batch_size=2, task_batches=1, - num_workers=0, + num_workers=num_workers, ) items, options = spec.items, spec.options_payload assert [ @@ -196,7 +213,49 @@ def load(*args, **kwargs): 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}" @@ -248,29 +307,121 @@ def load(*args, **kwargs): (round_dir / "scientific" / "summary.json").read_text() ) for key in ( - "performance", "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", - ["missing", "duplicate", "failed", "wrong_probability", "wrong_identity"], + "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 + tmp_path, monkeypatch, defect, message ): run_dir, dataset_dir = _inputs(tmp_path) _cpu_worker(monkeypatch) @@ -297,15 +448,31 @@ def test_merge_rejects_incomplete_or_invalid_scientific_results( 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_probability": - np.save(item_dir / "probabilities.npy", np.array([0.123, 0.456])) - summary["artifact_sha256"]["probabilities"] = file_sha256( - item_dir / "probabilities.npy" - ) - else: + 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): + with pytest.raises(ValueError, match=message): executor.finalize_inference_round( round_dir, records, items=items, options=options ) From 7f27a94fde9afe5b7c007208c1ecc3838667ab1a Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Thu, 24 Sep 2026 14:42:22 -0700 Subject: [PATCH 08/10] Release unused xFit planning input on MPI workers Retain the planning dataset only on rank zero, where scientific result finalization reuses it. Other MPI ranks release it before loading their worker dataset. Signed-off-by: Trent Nelson --- src/cuphoton/xfit/commands.py | 1 + src/cuphoton/xfit/executor.py | 8 +++++++- tests/core/test_mpi_adapters.py | 1 + tests/xfit/test_executor.py | 21 ++++++++++++++++++--- 4 files changed, 27 insertions(+), 4 deletions(-) diff --git a/src/cuphoton/xfit/commands.py b/src/cuphoton/xfit/commands.py index c0d872ae..84fcdcd3 100644 --- a/src/cuphoton/xfit/commands.py +++ b/src/cuphoton/xfit/commands.py @@ -327,6 +327,7 @@ def run(self) -> None: 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, diff --git a/src/cuphoton/xfit/executor.py b/src/cuphoton/xfit/executor.py index c43d3372..73a90779 100644 --- a/src/cuphoton/xfit/executor.py +++ b/src/cuphoton/xfit/executor.py @@ -488,7 +488,11 @@ def finalize_xfit_round( def prepare_xfit_workload( - *, input_path: Path, chunk_size: int = 256, fit_options: Mapping[str, Any] + *, + 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 @@ -508,6 +512,8 @@ def prepare_xfit_workload( 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): diff --git a/tests/core/test_mpi_adapters.py b/tests/core/test_mpi_adapters.py index cf1cfad8..8a26c2b0 100644 --- a/tests/core/test_mpi_adapters.py +++ b/tests/core/test_mpi_adapters.py @@ -109,6 +109,7 @@ def prepare(rank): input_path=path, chunk_size=2, fit_options={"backend": "cupy", "max_evaluations": 2}, + retain_input=rank == 0, ) else: diff --git a/tests/xfit/test_executor.py b/tests/xfit/test_executor.py index 1c7bbbeb..d48bb576 100644 --- a/tests/xfit/test_executor.py +++ b/tests/xfit/test_executor.py @@ -4,8 +4,10 @@ from __future__ import annotations +import gc import json import os +import weakref from dataclasses import fields, replace import numpy as np @@ -149,9 +151,11 @@ def test_preflight_accepts_analytic_gaussian_cutile(tmp_path): assert spec.backend == "cutile" -@pytest.mark.parametrize("runtime", ["dragon", "mpi"]) +@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 + tmp_path, monkeypatch, capsys, runtime, rank ): from types import SimpleNamespace @@ -161,15 +165,26 @@ def test_cli_dispatches_collective_preflight_and_preserves_options( 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"](0) + 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", From fd81c49fbdfad0e630805df535d844143feae98d Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Thu, 24 Sep 2026 17:45:51 -0700 Subject: [PATCH 09/10] Plan MPI component workloads once before broadcasting identity Signed-off-by: Trent Nelson --- src/cuphoton/xfit/commands.py | 2 ++ src/cuphoton/xscan/commands.py | 2 ++ src/cuphoton/xscan/pipeline_executor.py | 2 ++ tests/core/test_mpi_adapters.py | 16 ++++++++++++++++ tests/xfit/test_executor.py | 1 + tests/xscan/test_executor.py | 1 + tests/xscan/test_pipeline_executor.py | 17 +++++++++++++++++ 7 files changed, 41 insertions(+) diff --git a/src/cuphoton/xfit/commands.py b/src/cuphoton/xfit/commands.py index 84fcdcd3..28c829e3 100644 --- a/src/cuphoton/xfit/commands.py +++ b/src/cuphoton/xfit/commands.py @@ -298,6 +298,8 @@ def run(self) -> None: 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, diff --git a/src/cuphoton/xscan/commands.py b/src/cuphoton/xscan/commands.py index 2a44d411..22669ebc 100644 --- a/src/cuphoton/xscan/commands.py +++ b/src/cuphoton/xscan/commands.py @@ -1024,6 +1024,8 @@ def run(self) -> 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, diff --git a/src/cuphoton/xscan/pipeline_executor.py b/src/cuphoton/xscan/pipeline_executor.py index d9196c85..579ce2ba 100644 --- a/src/cuphoton/xscan/pipeline_executor.py +++ b/src/cuphoton/xscan/pipeline_executor.py @@ -72,6 +72,8 @@ def resolve(raw: Any) -> str: 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 = [] diff --git a/tests/core/test_mpi_adapters.py b/tests/core/test_mpi_adapters.py index 8a26c2b0..4b91450a 100644 --- a/tests/core/test_mpi_adapters.py +++ b/tests/core/test_mpi_adapters.py @@ -43,6 +43,8 @@ def test_two_rank_adapters_publish_scientific_results( state = local() barrier = Barrier(2, timeout=10) messages = [None, None] + planning_ranks = [] + input_load_ranks = [] class Comm: def Get_rank(self): @@ -96,6 +98,13 @@ def gpu_identity(backend): 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"}) @@ -105,6 +114,7 @@ def cpu_fit(*args, **kwargs): 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, @@ -125,6 +135,7 @@ def prepare(rank): ) def prepare(rank): + planning_ranks.append(rank) return executor.prepare_inference_workload( run_dir=run_dir, dataset_dir=dataset_dir, @@ -137,6 +148,7 @@ 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, @@ -147,6 +159,10 @@ def run(rank): 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" diff --git a/tests/xfit/test_executor.py b/tests/xfit/test_executor.py index d48bb576..4c00db8d 100644 --- a/tests/xfit/test_executor.py +++ b/tests/xfit/test_executor.py @@ -213,6 +213,7 @@ def run(**kwargs): == 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 diff --git a/tests/xscan/test_executor.py b/tests/xscan/test_executor.py index 02403436..80bd04cd 100644 --- a/tests/xscan/test_executor.py +++ b/tests/xscan/test_executor.py @@ -133,6 +133,7 @@ def run(**kwargs): == 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 diff --git a/tests/xscan/test_pipeline_executor.py b/tests/xscan/test_pipeline_executor.py index d8c796db..56b6cebf 100644 --- a/tests/xscan/test_pipeline_executor.py +++ b/tests/xscan/test_pipeline_executor.py @@ -111,6 +111,23 @@ def test_manifest_rejects_ambiguous_or_invalid_descriptors(tmp_path, defect): 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 ): From 1279f39aeadc0ed0df8e3a28034a19f9ff751da3 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Thu, 24 Sep 2026 18:17:37 -0700 Subject: [PATCH 10/10] Preserve NPY symlink names in pipeline manifests Signed-off-by: Trent Nelson --- src/cuphoton/xscan/pipeline_executor.py | 3 ++- tests/xscan/test_pipeline_executor.py | 35 +++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/src/cuphoton/xscan/pipeline_executor.py b/src/cuphoton/xscan/pipeline_executor.py index 579ce2ba..2caf0a75 100644 --- a/src/cuphoton/xscan/pipeline_executor.py +++ b/src/cuphoton/xscan/pipeline_executor.py @@ -63,7 +63,8 @@ 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() - return str((path.parent / value).resolve()) + nested = path.parent / value + return str(nested.parent.resolve() / nested.name) configuration = json_mapping( payload["configuration"], field="configuration" diff --git a/tests/xscan/test_pipeline_executor.py b/tests/xscan/test_pipeline_executor.py index 56b6cebf..b13291ca 100644 --- a/tests/xscan/test_pipeline_executor.py +++ b/tests/xscan/test_pipeline_executor.py @@ -90,6 +90,41 @@ def test_manifest_paths_are_relative_to_manifest_not_cwd( 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"] )