Document Dragon GPU sharing, MPS and MPI execution

docs/README.md

@@ -44,10 +44,15 @@ requirements.
 ## Adapting and validating
 
 - [Architecture](architecture.md): package boundaries and execution model.
-- [Distributed execution](distributed-execution.md): prepare a pipeline
+- [Distributed architecture](distributed.md): work assignment, GPU ownership,
+  data movement and completion.
+- [Dragon architecture](dragon.md): process placement, GPU sharing with MPS,
+  queues and worker lifecycle.
+- [MPI architecture](mpi.md): rank binding, collective coordination and file mode.
+- [Distributed launch guide](distributed-execution.md): prepare a pipeline
   manifest and launch MPI or Dragon workers using Slurm or SSH host files.
-- [Dragon transport and coordination](dragon.md): execution models,
-  consumer-local queues, transport selection, and timing boundaries.
+- [Dragon transport and performance](dragon-performance.md): transport selection,
+  queue placement experiments and timing boundaries.
 - [Pipeline stage benchmark](components/pipeline-stage-benchmark.md): compare
   a persistent GPU pipeline with separate stages and inspect transfer costs.
 - [Adapting the workflows](adapting-workflows.md): a practical extension

docs/architecture.md

@@ -112,6 +112,8 @@ Core's public CLI facade.
 Core also supplies FITS metadata inspection and reader selection, plus the
 shared Dragon/MPI execution machinery. Components retain their scientific
 input validation and choose how to divide work among those executors.
+See the [distributed architecture](distributed.md) for process ownership,
+data movement and the Dragon/MPI worker lifecycles.
 
 xDataReader, xFit, xPois, and xRep expose a curated Python surface
 for embedding numerical operations. xScan and xRay are primarily

docs/components/xpois.md

@@ -748,7 +748,7 @@ selection remains a launcher option; `--executor dragon` cannot change the
 transport of an already running Dragon session. The worker wall-time limit
 covers the entire invocation, including warmups and all measured rounds.
 
-See [Dragon transport and coordination](../dragon.md) for native HSTA TCP
+See [Dragon transport and performance](../dragon-performance.md) for native HSTA TCP
 launch options, consumer-local queues in the persistent executors,
 and measurement boundaries.
 

docs/distributed-execution.md

@@ -3,7 +3,9 @@
 The xPois, xFit and xScan pipeline assigns complete image pairs to GPU
 workers. Each worker loads its model once and processes its assigned pairs
 through subtraction, dipole fitting and classification. MPI and Dragon
-provide alternative launchers for the same workload.
+provide alternative launchers for the same workload. The
+[distributed architecture](distributed.md) explains work assignment, process
+ownership and data movement.
 
 Three inputs determine a run:
 
@@ -16,8 +18,8 @@ Three inputs determine a run:
 `mpiexec -n 8` requests eight processes. The host file or scheduler allocation
 determines where they run. With Dragon, `--nodes 2` selects two nodes and
 cuPhoton's `--max-workers 8` caps the total GPU workers. The examples below
-use two nodes with four GPUs each, one worker per GPU, and at least eight
-image-pair items.
+use two nodes with four GPUs each and at least eight image-pair items.
+They use the default placement of one worker per GPU.
 
 ## Prepare the environment and workload
 
@@ -145,10 +147,16 @@ allocated nodes, and cuPhoton discovers their GPUs and places its workers:
   --max-workers 8 --warmup-rounds 1 --measure-rounds 3
 ```
 
-`--max-workers` is a cap: Dragon also limits workers by available GPUs and
-item count. This example selects TCP for both Dragon transports. Use the
+`--max-workers` caps the total process count. With the default
+`--workers-per-gpu 1`, Dragon also limits workers by available GPUs and item
+count. This example selects TCP for both Dragon transports. Use the
 site's routable interface selection when needed; see
-[Dragon transport and coordination](dragon.md).
+[Dragon transport and performance](dragon-performance.md).
+
+For multiple workers per GPU, see [Dragon GPU sharing and MPS](dragon.md#share-a-gpu-between-workers)
+and the [GPU sharing commands](components/xscan.md#share-a-gpu-between-image-pairs).
+An explicit MPS connection requires an existing service on each worker host.
+The launcher or administrator manages that service separately.
 
 ## Launch on named hosts through SSH
 
@@ -227,7 +235,8 @@ Each command creates `runs/<name>/summary.json` and retains warmup and
 measured outputs under `rounds/warmup-*` and `rounds/measure-*`. Inspect the
 summary and worker records for actual host/GPU placement, successful item
 completion and cleanup. MPI rejects more ranks than image-pair items;
-duplicate physical GPU assignments also fail validation.
+unexpected sharing or a physical GPU identity that disagrees with the
+requested placement also fails validation.
 
 Choose a new `--name` for another attempt. Readiness, batch execution,
 validation and cleanup have separate timing fields. Batch timing includes

docs/distributed.md

@@ -0,0 +1,172 @@
+# Distributed execution architecture
+
+cuPhoton uses Dragon or MPI to run independent work on several GPUs. Each
+GPU worker is a Python process bound to one GPU. The worker loads its inputs,
+runs the numerical code and writes its results. A coordinator assigns work
+and checks completion. Dragon defaults to one worker per GPU. Its shared
+executor can place several independent workers on the same GPU.
+
+The choice of Dragon or MPI controls process placement and coordination. The component
+and numerical backend control the calculation. For example, the combined
+imaging pipeline uses CuPy and PyTorch inside each worker with either runtime.
+
+Read [How cuPhoton uses Dragon](dragon.md) or [How cuPhoton uses MPI](mpi.md)
+for the process and message details. The [distributed launch guide](distributed-execution.md)
+contains environment setup, manifests and Slurm/SSH commands.
+
+## Units of work
+
+The optical workflows divide a batch into independent items:
+
+| Workflow | One item |
+| --- | --- |
+| xPois batch fitting | One aligned reference/target image pair |
+| xFit batch fitting | A range of candidate stamps |
+| xScan inference | A range of samples that preserves the planned minibatches |
+| Combined xPois, xFit and xScan pipeline | One image pair with its candidate coordinates |
+
+An imaging worker completes all stages for an assigned pair on its assigned GPU.
+The executors do not split that pair across GPUs or exchange intermediate
+images between workers. The distributed xRay detector workflow has a separate
+[tile and artifact contract](xray/DISTRIBUTED-DETECTOR-ARTIFACTS.md).
+
+A `WorkItem` contains an ID, a JSON description and a weight for assignment.
+The description refers to files, array properties or sample ranges. It does
+not contain a live GPU array.
+
+cuPhoton sorts items by descending weight and assigns each to the worker
+with the smallest total weight. Each worker receives a fixed list, called
+a *shard*. Weights usually estimate input bytes. xScan uses sample counts.
+They do not predict execution time exactly. Workers process their shards
+sequentially, with no work stealing. A round waits for its slowest worker.
+
+## Where the data goes
+
+This diagram shows the combined imaging pipeline inside one worker. The
+other workers run the same stages on their own assigned pairs.
+
+```mermaid
+flowchart TB
+    F["Shared files: images and model checkpoint"]
+    C["Coordinator"]
+    subgraph W["One Python worker process, one assigned GPU"]
+        R["Read and validate one image pair"]
+        P["xPois: match images and subtract"]
+        X["xFit: fit candidate stamps"]
+        S["xScan: classify candidates"]
+        H["Copy compact results to host memory"]
+        R --> P --> X --> S --> H
+    end
+    O["Shared files: results and completion records"]
+    F --> R
+    C -. "Launch reference and round control" .-> W
+    H --> O
+    W -. "Status, identity and timing" .-> C
+    O -->|"Read and validate"| C
+```
+
+Dragon passes a reference to a shared launch file. MPI ranks plan items
+locally or receive them in a broadcast.
+
+Workers read inputs directly from the shared filesystem. NPY inputs pass
+through host memory before a GPU upload. FITS inputs use the configured
+reader, which can return host or device arrays. See the
+[FITS reader guide](components/xdr.md#read-fits-images-in-a-workflow) for that choice.
+
+The combined pipeline keeps intermediate image and stamp arrays on the GPU.
+CuPy and PyTorch share compatible device buffers through DLPack, an array
+exchange interface. cuPhoton retains the buffer owners until the GPU work
+finishes. These transfers occur within one process. Dragon and MPI carry
+descriptions and completion metadata. They do not transport these GPU buffers.
+
+Workers write scientific outputs and completion records to shared storage.
+The coordinator reads those records and checks the expected item IDs,
+input identities and output files. The filesystem therefore carries both
+bulk input/output and part of the completion protocol.
+
+## Process ownership
+
+| Responsibility | Dragon | MPI |
+| --- | --- | --- |
+| Start the application | Dragon launcher starts a coordinator within its runtime | MPI launcher starts all ranks |
+| Start GPU workers | Coordinator creates a native Dragon `ProcessGroup` | Each rank is a GPU worker |
+| Coordinate the run | Separate coordinator process | Rank 0, which also computes its shard |
+| Select each GPU | Dragon host and GPU placement policy | Launcher binding and per-node rank mapping |
+| Exchange status | Dragon native queues | MPI collectives; the xPois file mode uses shared files |
+
+Each worker must see one assigned GPU before a numerical library initializes
+CUDA. Inside that process, the assigned GPU is device 0. cuPhoton records
+the host, process ID and physical GPU identity. A CUDA ordinal alone does
+not identify a physical GPU.
+
+MPI and the default Dragon configuration require distinct GPUs. With Dragon
+sharing enabled, workers assigned to the same GPU must report the same
+physical identity. Workers assigned to different GPUs must report different
+identities. See [Dragon GPU sharing and MPS](dragon.md#share-a-gpu-between-workers)
+for placement, connection checks and service ownership.
+
+Every node needs compatible software and access to the input, model and
+output paths at the same absolute locations. The launch guide describes
+these requirements and how to preserve scheduler GPU restrictions.
+
+## What workers retain
+
+In repeated runs, each GPU worker stays alive through all requested warmup
+and measured rounds. A round repeats the same shards. Warmup rounds run real work and retain their outputs.
+
+| Worker | State retained between items or rounds |
+| --- | --- |
+| xPois repeated batch | Process and CUDA context; inputs are read again each round |
+| xFit | Validated input, solver settings, model and CUDA context |
+| xScan | Dataset, metadata rows, model and data loader |
+| Combined pipeline | Model, CUDA streams and device context; each pair is read again |
+
+Each additional Dragon worker retains its own component state, including its
+model where applicable. More workers can increase host and device memory use.
+The pipeline also offers local process and thread executors; see the
+[GPU sharing guide](components/xscan.md#share-a-gpu-between-image-pairs).
+
+Keeping a process alive does not keep every input image on the GPU. Output
+writes, input checks and any required transfers still occur. Ordinary
+single-pass xPois uses its own executor path. The Dragon and MPI pages
+describe it.
+
+## Completion and failure
+
+The shared executors collect a completion record from each worker after a
+round. They then validate the item records and run any component-specific
+merge step. A record of success must agree with the files on disk.
+
+For a successful run, workers close their component resources before the
+coordinator publishes the final `summary.json`. Failed work, missing records
+or GPU identities that disagree with the requested placement make the run
+fail. Failed shutdown also makes the run fail. A successful round alone does not
+establish a successful complete run.
+
+These executors do not reschedule a failed item onto another worker.
+Completed files remain available for diagnosis. Validation confirms one terminal record per expected item. It does not
+retry or recover work.
+
+## Timing boundaries
+
+The reported `batch_wall_sec` covers round dispatch through worker completion
+collection. Worker completion includes its item processing and record writes.
+Coordinator validation and component merging follow this interval. The shared
+executor reports component merging as `finalization_sec`.
+
+Measure launcher-to-exit wall time separately to include runtime startup and
+process teardown. Keep warmup, readiness and shutdown costs visible when
+comparing Dragon and MPI. Subtracting the longest worker time from batch time
+does not isolate network latency. See
+[transport measurements and timing](dragon-performance.md#evidence-and-timing-boundaries)
+for the existing experiments and their limits.
+
+## Source map
+
+| Source | Responsibility |
+| --- | --- |
+| [`core/bulk.py`](../src/cuphoton/core/bulk.py) | Work items, shard assignment and completion-record checks |
+| [`core/execution.py`](../src/cuphoton/core/execution.py) | Component worker contract, item execution and round validation |
+| [`core/executors.py`](../src/cuphoton/core/executors.py) | Select Dragon or MPI |
+| [`xscan/pipeline_executor.py`](../src/cuphoton/xscan/pipeline_executor.py) | Adapt the imaging pipeline to the shared executors |
+| [`xscan/device_pipeline.py`](../src/cuphoton/xscan/device_pipeline.py) | Run the imaging stages and manage device-buffer lifetimes |