Run xFit, XScan and the device pipeline with Dragon or MPI

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 <command>` 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`,

docs/components/xfit.md

@@ -142,6 +142,50 @@ acceptable.
 See [Data and artifact contracts](../data-artifacts.md#xfit-dipole-batches)
 for the stable shapes and output fields.
 
+## Distributed fitting
+
+`fit-dipoles --executor dragon|mpi` distributes independent candidate chunks
+across GPUs. The default `--executor local` retains the original batch fit.
+Distributed fitting requires `--backend cupy` or `--backend cutile`, a shared
+filesystem for the input and output, and the same installed environment on
+every worker. `--chunk-size` sets candidates per task independently of worker
+count. Keep it fixed for matched comparisons; candidate IDs and input order
+are restored in the merged artifacts.
+The task count must be at least the MPI rank count. Dragon uses the smaller
+of the requested worker count and task count.
+
+Under an allocation with Dragon configured, run one warmup and two measured
+passes with workers retained across all three passes:
+
+```bash
+dragon .venv/bin/cuphoton xfit fit-dipoles \
+  --executor dragon --max-workers 8 \
+  --input /shared/dipoles.npz --model gaussian --backend cupy \
+  --chunk-size 256 --output-dir /shared/results/xfit-dragon \
+  --warmup-rounds 1 --measure-rounds 2
+```
+
+With Open MPI, use the installed rank wrapper to narrow GPU visibility before
+Python starts. The parent mask must list allocated GPUs in local-rank order:
+
+```bash
+: "${CUDA_VISIBLE_DEVICES:?must enumerate the allocated GPUs}"
+mpirun -n 8 --map-by slot --bind-to none -x CUDA_VISIBLE_DEVICES \
+  .venv/bin/cuphoton-openmpi-rank-exec -- \
+  .venv/bin/cuphoton xfit fit-dipoles \
+  --executor mpi --input /shared/dipoles.npz --model gaussian --backend cupy \
+  --chunk-size 256 --output-dir /shared/results/xfit-mpi \
+  --warmup-rounds 1 --measure-rounds 2
+```
+
+The output directory must be new. Its basename is the run ID: 1–128 ASCII
+letters, digits, dots, underscores or hyphens, starting with a letter or digit.
+Each pass retains normal xFit artifacts
+under `rounds/<round-id>/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

docs/components/xscan.md

@@ -236,6 +236,64 @@ The `*.blackwell.example.yaml` files demonstrate throughput-oriented settings
 for recent NVIDIA GPUs. Tune these starting points for your hardware and
 dataset.
 
+## Distributed inference
+
+`infer-real-bogus --executor dragon|mpi` scores a fixed split across GPUs.
+Each worker loads the checkpoint once and retains it across optional warmup
+and measured passes. Training and evaluation commands retain their existing
+local behavior. The default `--executor local` also preserves the original
+inference output location.
+
+Inputs, checkpoint, optional xFit feature bundle and output directory must
+reside on a filesystem shared by all workers. `--batch-size` retains its
+ordinary inference meaning. `--task-batches` groups whole minibatches into
+tasks; only the final task can contain the original final partial batch.
+Keep both values fixed when comparing worker counts. The merge restores
+selected-split order, original sample indices, labels and candidate metadata,
+and uses the same host probability calculation as local inference.
+The task count must be at least the MPI rank count. Dragon uses the smaller
+of the requested worker count and task count.
+
+Distributed inference defaults to `--num-workers 0`, independently of the
+checkpoint's loader setting. An explicit positive value creates persistent
+loader processes during worker setup; the loader and parsed metadata are
+reused across tasks and rounds. Workers retain inputs and model state, so
+their initial load is outside the timed rounds.
+
+Under a configured Dragon allocation:
+
+```bash
+dragon .venv/bin/cuphoton xscan infer-real-bogus \
+  --executor dragon --max-workers 8 \
+  --run-dir /shared/model --dataset-dir /shared/dataset --split test \
+  --batch-size 32 --task-batches 16 --num-workers 0 \
+  --output-dir /shared/results/inference-dragon \
+  --warmup-rounds 1 --measure-rounds 2
+```
+
+With Open MPI, the rank wrapper narrows GPU visibility before Python starts.
+The parent mask must list allocated GPUs in local-rank order:
+
+```bash
+: "${CUDA_VISIBLE_DEVICES:?must enumerate the allocated GPUs}"
+mpirun -n 8 --map-by slot --bind-to none -x CUDA_VISIBLE_DEVICES \
+  .venv/bin/cuphoton-openmpi-rank-exec -- \
+  .venv/bin/cuphoton xscan infer-real-bogus \
+  --executor mpi --run-dir /shared/model --dataset-dir /shared/dataset \
+  --split test --batch-size 32 --task-batches 16 --num-workers 0 \
+  --output-dir /shared/results/inference-mpi \
+  --warmup-rounds 1 --measure-rounds 2
+```
+
+Distributed inference requires a new `--output-dir`. Its basename is the run
+ID: 1–128 ASCII letters, digits, dots, underscores or hyphens, starting with
+a letter or digit. Each pass writes merged
+logits, labels, probabilities, sample indices and a summary under
+`rounds/<round-id>/scientific/`; warmup outputs are retained too. Without
+round flags, the single pass uses `scientific/`. The model directory remains
+unchanged. Execution receipts and timing are separate from these scientific
+outputs, and merging and validation occur after the timed worker phase.
+
 ## Persistent XPOIS, xFit and XScan pipeline
 
 The Python API in `cuphoton.xscan.device_pipeline` runs complete image pairs
@@ -356,6 +414,48 @@ scientific evidence; Dragon saves each result in
 `items/<item_id>/summary.json` and checks it again in the coordinator.
 The evidence decoder returns the 22 named arrays for comparison.
 
+The `run-pipeline` command accepts the same descriptors with either Dragon or
+MPI. Save the configuration and items above as a manifest:
+
+```python
+import json
+
+Path("pipeline.json").write_text(json.dumps({
+    "schema": "cuphoton.xscan.pipeline-manifest/v1",
+    "configuration": config.to_payload(),
+    "items": [item.to_payload() for item in items],
+}, indent=2) + "\n")
+```
+
+Paths may be absolute or relative to the manifest. All nodes must see the
+same source, inputs and output directory. Use your site's launcher settings;
+these examples run one warmup pass and two measured passes:
+
+```bash
+dragon cuphoton xscan run-pipeline --executor dragon \
+  --manifest pipeline.json --output-dir runs --name dragon-pipeline \
+  --max-workers 8 --warmup-rounds 1 --measure-rounds 2
+
+mpiexec -n 8 -x CUDA_VISIBLE_DEVICES cuphoton-openmpi-rank-exec -- \
+  cuphoton xscan run-pipeline --executor mpi \
+  --manifest pipeline.json --output-dir runs --name mpi-pipeline \
+  --warmup-rounds 1 --measure-rounds 2
+```
+
+The MPI example requires Open MPI and the allocation's visible GPU list on
+each node. Other MPI launchers must bind each rank to exactly one GPU before
+starting Python; do not use the Open MPI wrapper with MPICH. Each rank uses
+local `cuda:0`, and duplicate physical GPU assignments fail validation.
+
+Both executors initialize one worker context and reuse it across all items
+and rounds. Warmup outputs remain under `rounds/warmup-*`; measured outputs
+remain under `rounds/measure-*`. Each round retains ordinary item results and
+audit evidence. The parent `summary.json` separates readiness, batch time,
+artifact validation and cleanup. Batch time includes dispatch, input loading,
+numerical work and worker output publication; coordinator audits follow that
+timer. It is not kernel-only time. Failed rounds or cleanup invalidate the
+reported statistics. Omitting both round flags runs a single pass.
+
 The pipeline retains device owners through the blocking terminal copy and
 synchronizes failed work before reuse. Failed cleanup makes the context
 unusable. Transfer receipts count pipeline-owned uploads and the packed

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/<round-id>/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

src/cuphoton/xfit/commands.py

@@ -12,6 +12,7 @@
 
 import numpy as np
 
+from cuphoton.core.bulk import validate_identifier
 from cuphoton.core.cli import (
     BoolInvariant,
     CommandError,
@@ -22,6 +23,7 @@
     SetInvariant,
     StringInvariant,
 )
+from cuphoton.core.cli.executor import ExecutorOptions
 
 from ._types import (
     BACKEND_REQUESTS,
@@ -186,7 +188,7 @@ def run(self) -> None:
         self._emit_json(payload)
 
 
-class FitDipolesCommand(_ValidatedDatasetCommand):
+class FitDipolesCommand(ExecutorOptions, _ValidatedDatasetCommand):
     """Fit a batch of astronomical dipoles and persist safe artifacts."""
 
     output_dir = None
@@ -199,6 +201,13 @@ class FitDipolesCommand(_ValidatedDatasetCommand):
     g_tol = None
     max_evaluations = None
     use_finite_difference = None
+    chunk_size = None
+
+    class ChunkSizeArg(PositiveIntegerInvariant):
+        _arg = "--chunk-size"
+        _help = "Candidates per distributed task. [default: 256]"
+        _mandatory = False
+        _default = None
 
     class OutputDirArg(PathSpecInvariant):
         _arg = "--output-dir"
@@ -283,6 +292,59 @@ def _model_specification(
         )
 
     def run(self) -> None:
+        executor_options = self.executor_options()
+        if self.executor != "local":
+            from cuphoton.core.executors import run_workload
+
+            from .executor import prepare_xfit_workload
+
+            if self.executor == "mpi":
+                executor_options["prepare_on_root"] = True
+            output_dir = Path(self.output_dir).expanduser().resolve()
+            self._call(
+                validate_identifier,
+                output_dir.name,
+                field="--output-dir basename",
+            )
+            fit_options = {
+                name: getattr(self, name)
+                for name in (
+                    "model",
+                    "mode",
+                    "backend",
+                    "compute_dtype",
+                    "stamp_evaluation",
+                    "stamp_scale",
+                    "f_tol",
+                    "x_tol",
+                    "g_tol",
+                    "max_evaluations",
+                    "use_finite_difference",
+                )
+            }
+            result = self._call(
+                run_workload,
+                executor=self.executor,
+                prepare_workload=lambda rank: prepare_xfit_workload(
+                    input_path=self.input,
+                    chunk_size=self.chunk_size or 256,
+                    fit_options=fit_options,
+                    retain_input=rank == 0,
+                ),
+                output_root=output_dir.parent,
+                run_id=output_dir.name,
+                **executor_options,
+            )
+            if result is not None:
+                self._emit_json(result.to_dict())
+                if result.status != "success":
+                    raise CommandError("distributed xFit execution failed")
+            return
+        if self.chunk_size is not None:
+            raise CommandError(
+                "--chunk-size requires --executor dragon or mpi"
+            )
+
         from . import LMConfig, fit_dipoles
 
         dataset = self._load_dataset()