hq: phase 4

docs/spec.md

@@ -2048,6 +2048,12 @@ sampler.run_with_samples(
   the file may be much larger than RAM. Same key handling as the in-memory
   branch: missing keys raise, extra keys are ignored.
 
+Both branches return parameters in declaration order: the prior's nonlinear
+parameters, then the extensions', then any explicitly sampled linear ones. The
+order never depends on set iteration, so it is the same in every process
+whatever `PYTHONHASHSEED` is (results written by separate processes, as in
+harv-hq's shards, must agree on it).
+
 `randomize_prior_order` (HDF5 path only): when `True` (default), batch *order*
 is permuted via `jax.random.permutation(key, n_batches)`. Each
 batch is still a single contiguous h5py slice — no random seeks, no read

packages/hq/docs/implementation-plan.md

@@ -26,8 +26,8 @@ Updated as each phase lands. A phase is done when its PR is merged.
 | —     | hq scaffolding, spec, and this plan                    | Done (merged, PR #58)                             |
 | 1     | config, model file, IDs, Parquet I/O, provenance, CLI  | Done (merged, PR #59)                             |
 | 2     | prepare, `read_source`, `read_sources`                 | Done (PR #60)                                     |
-| 3     | prior cache, serial rejection, result parts, resume    | Implemented on branch `hq-phase3`, awaiting merge |
-| 4     | execution modes: shards, pool, MPI                     | Not started                                       |
+| 3     | prior cache, serial rejection, result parts, resume    | Done (PR #61)                                     |
+| 4     | execution modes: shards, pool, MPI                     | Implemented on branch `hq-phase4`, awaiting merge |
 | 5     | compaction, summarize, status                          | Not started                                       |
 | 6     | MCMC follow-up                                         | Not started                                       |
 | 7     | web viewer                                             | Not started                                       |
@@ -59,7 +59,6 @@ ______________________________________________________________________
 packages/hq/src/harv_hq/
 ├── __init__.py          # curated __all__ (tuple)
 ├── _version.py          # generated by hatch-vcs
-├── _threads.py          # pin_threads() -- called by pool workers before importing JAX
 ├── config.py            # Config (frozen dataclasses, Config.from_file), ConfigError
 ├── _model_file.py       # ModelFile: import prior.py, expose make_setup & hooks
 ├── ids.py               # stable_hash, source_key, prior_cache_key
@@ -70,7 +69,7 @@ packages/hq/src/harv_hq/
 ├── summary_stats.py     # per-source summary statistics + weighted resample
 ├── rejection.py         # process one source: rejection
 ├── mcmc.py              # process one source: MCMC
-├── execution.py         # serial / pool / MPI drivers around a per-source fn
+├── execution.py         # serial / pool drivers, mpi_comm, thread pinning, Progress
 ├── compact.py           # compact(): merge a stage's parts, drop superseded rows
 ├── summarize.py         # summarize(), status()
 ├── run.py               # Run, SourceResult
@@ -273,7 +272,7 @@ linear time (loose bound).
 
 ______________________________________________________________________
 
-## Phase 3: prior cache, serial rejection, result parts, resume — awaiting merge
+## Phase 3: prior cache, serial rejection, result parts, resume — done
 
 Implemented on branch `hq-phase3` (stacked on `hq-phase2`). Deviations from the
 text below: `Run.run_rejection` takes `shard` and `overwrite`
@@ -362,7 +361,30 @@ rows.
 
 ______________________________________________________________________
 
-## Phase 4: execution modes (shards, pool, MPI)
+## Phase 4: execution modes (shards, pool, MPI) — awaiting merge
+
+Implemented on branch `hq-phase4`. Deviations from the text below: there is
+no `_threads.py`. A worker cannot pin its own threads, because unpickling its
+initializer imports `harv_hq`, which loads NumPy and JAX first, so the parent
+sets the thread variables in the environment the workers are spawned with
+and restores its own afterwards (spec updated). The drivers take
+`(ids, data, setup, sink)`: `setup` is a picklable
+`functools.partial(rejection_setup, config, cache_path)` that builds the
+per-source function (each pool worker calls it once), and `sink` adds to the
+`PartWriter` and counts progress. There is no `run_mpi` driver: `--mpi` takes
+the slice from the rank and then runs the serial driver, or the pool driver
+with `--workers` (`--mpi --workers` is allowed, for one rank per node; spec
+updated). With `--overwrite`,
+only rank 0 supersedes, between two barriers (spec updated). Every mode logs
+progress about every 5% of its slice. The real-`mpirun` test is skipped
+unless `mpi4py` and `mpirun` are both present, and could not be run locally
+(Homebrew's Open MPI crashes even a one-rank job on this machine), so the MPI
+path is also tested with a stand-in communicator. Found along the way, and
+fixed in harv: `run_with_samples` ordered parameters by set iteration, so two
+processes could write differently ordered columns, and the structure check
+then refused the second one (harv spec updated). From the Copilot review of
+PR #63: a rank that raises calls `comm.Abort(1)`, so its peers cannot hang at
+a barrier (spec updated).
 
 **`execution.py`** wraps a per-source function in three drivers sharing one
 contract: given `(ids, data, process_fn, writer)`, process every ID and pass

packages/hq/docs/spec.md

@@ -601,22 +601,33 @@ falls as `N` grows.
 | Explicit shard           | `hq run --shard 3/16`               | `3/16`              | `0003-of-0016-<uuid>`     |
 | Local pool               | `hq run --workers 8 [--shard 3/16]` | the process's slice | as above, by the parent   |
 | MPI                      | `mpirun -n 64 hq run --mpi`         | `rank/size`         | `<rank>-of-<size>-<uuid>` |
+| MPI with a pool per rank | `hq run --mpi --workers 32`         | `rank/size`         | as above, by each rank    |
 
 - **Explicit shards** are for job arrays and task launchers (SLURM
   `--array=0-15` with `--shard $SLURM_ARRAY_TASK_ID/16`, disBatch). Shards are
   fully independent processes.
 - **Local pool** runs a `ProcessPoolExecutor` with the `spawn` start method.
-  Each worker is pinned to one CPU thread (BLAS and XLA environment variables
-  set before JAX is imported), so `W` workers use `W` cores. The worker
+  Each worker is pinned to one CPU thread, so `W` workers use `W` cores: the
+  parent sets the BLAS thread variables (`OMP_NUM_THREADS` and friends) and
+  appends single-thread flags to `XLA_FLAGS` in the environment the workers
+  are spawned with, and restores its own environment afterwards. The worker
   initializer imports the model file and builds the sampler once. The parent
-  loads the slice's data and sends each task its source's arrays; workers
-  return results as plain NumPy payloads, and only the parent buffers and
-  writes parts.
+  loads the slice's data and sends each task its source's data; at most
+  `2 * W` sources are in flight. Workers return plain payloads, and only the
+  parent buffers and writes parts, in completion order.
 - **MPI** uses `mpi4py` (the `harv-hq[mpi]` extra), imported lazily so a
   missing install raises `ImportError` naming `hq run --mpi`. Rank `r` of `n`
   processes slice `r/n` and writes its own parts, so ranks need no
-  communication beyond a final barrier. Rank 0 additionally logs aggregate
-  progress. `--mpi` is mutually exclusive with `--shard` and `--workers`.
+  communication beyond barriers: one after the start (with `--overwrite`,
+  only rank 0 moves the old results aside, and the others wait for it) and
+  one at the end. Rank 0 additionally logs progress extrapolated to all
+  ranks. A rank that raises (anything outside a single source's
+  processing, which only marks that source `failed`) logs the error and calls
+  `MPI.COMM_WORLD.Abort(1)`, so no rank is left waiting at a barrier.
+  `--mpi` is mutually exclusive with `--shard`. It combines with
+  `--workers`: each rank runs a local pool over its slice, so a cluster job
+  can use one rank per node (`srun --ntasks-per-node=1 hq run --mpi --workers 32`), and only the pool workers are pinned to one thread.
+- Every mode logs a progress line about every 5% of its slice.
 - GPU runs use the single-process or explicit-shard modes, one process per
   device.
 

packages/hq/pyproject.toml

@@ -39,6 +39,9 @@ dependencies = [
     "pyarrow>=17",
 ]
 
+[project.optional-dependencies]
+mpi = ["mpi4py>=4"]
+
 [project.scripts]
 hq = "harv_hq.cli:main"
 

packages/hq/src/harv_hq/cli.py

@@ -65,9 +65,11 @@ def _prior_cache(args: argparse.Namespace) -> int:
 
 
 def _run(args: argparse.Namespace) -> int:
-    if args.workers != 1 or args.mpi:
-        sys.exit("hq run: --workers and --mpi are not implemented yet; use --shard")
-    Run(args.run_dir).run_rejection(shard=args.shard, overwrite=args.overwrite)
+    if args.workers < 1:
+        sys.exit("hq run: --workers must be at least 1")
+    Run(args.run_dir).run_rejection(
+        shard=args.shard, workers=args.workers, mpi=args.mpi, overwrite=args.overwrite
+    )
     return 0
 
 

packages/hq/src/harv_hq/execution.py

@@ -0,0 +1,182 @@
+"""Drivers that run a per-source function over a slice of sources.
+
+See ``packages/hq/docs/spec.md``, "Execution modes". A driver gets the
+slice's IDs and data, a picklable ``setup`` that builds the per-source
+function (``setup() -> process(source_id, data) -> Payload``), and a ``sink``
+that receives every payload. Only the calling process writes results, so
+drivers know nothing about parts. The MPI mode is the serial driver run on
+slice ``rank/size`` (:func:`mpi_comm`).
+"""
+
+__all__ = ("Progress", "mpi_comm", "run_pool", "run_serial")
+
+import logging
+import multiprocessing
+import os
+from collections.abc import Callable, Mapping, Sequence
+from concurrent.futures import FIRST_COMPLETED, Future, ProcessPoolExecutor, wait
+from typing import Any, cast, final
+from unittest.mock import patch
+
+from harv_hq.results import Payload
+
+logger = logging.getLogger("harv_hq")
+
+ProcessFn = Callable[[Any, Any], Payload]
+Setup = Callable[[], ProcessFn]
+Sink = Callable[[Payload], None]
+
+# One CPU thread per pool worker, so W workers use W cores. Set in the
+# environment the workers are spawned with: by the time a worker's own code
+# runs, importing harv_hq has already loaded NumPy (BLAS) and JAX.
+_THREAD_ENV = {
+    "OMP_NUM_THREADS": "1",
+    "OPENBLAS_NUM_THREADS": "1",
+    "MKL_NUM_THREADS": "1",
+    "VECLIB_MAXIMUM_THREADS": "1",
+}
+_XLA_SINGLE_THREAD = "--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1"
+
+
+@final
+class Progress:
+    """Log a progress line about every 5% of a slice.
+
+    Parameters
+    ----------
+    label
+        The stage, for the log line.
+    total
+        Sources in this process's slice.
+    ranks
+        MPI ranks; above 1, the line also extrapolates this rank's count to
+        the whole run (rank 0 only logs this).
+    """
+
+    def __init__(self, label: str, total: int, *, ranks: int = 1) -> None:
+        self.label = label
+        self.total = total
+        self.ranks = ranks
+        self.done = 0
+        self._every = max(1, total // 20)
+
+    def tick(self) -> None:
+        """Count one finished source."""
+        self.done += 1
+        if self.done % self._every and self.done != self.total:
+            return
+        overall = (
+            f" (about {self.done * self.ranks}/{self.total * self.ranks} "
+            f"over {self.ranks} ranks)"
+            if self.ranks > 1
+            else ""
+        )
+        logger.info("%s: %d/%d sources%s", self.label, self.done, self.total, overall)
+
+
+def run_serial(
+    ids: Sequence[Any], data: Mapping[Any, Any], setup: Setup, sink: Sink
+) -> None:
+    """Process every source in this process, in order.
+
+    Parameters
+    ----------
+    ids
+        The sources to process, in order.
+    data
+        Source ID to its data.
+    setup
+        Builds the per-source function.
+    sink
+        Receives each payload.
+    """
+    process = setup()
+    for source_id in ids:
+        sink(process(source_id, data[source_id]))
+
+
+def run_pool(
+    ids: Sequence[Any],
+    data: Mapping[Any, Any],
+    setup: Setup,
+    sink: Sink,
+    *,
+    workers: int,
+) -> None:
+    """Process sources on ``workers`` spawned processes; payloads return here.
+
+    Each worker runs ``setup`` once (its initializer) and is pinned to one
+    CPU thread. At most ``2 * workers`` sources are in flight, so memory stays
+    flat however long ``ids`` is. Payloads reach ``sink`` in completion order.
+
+    Parameters
+    ----------
+    ids
+        The sources to process; submitted in this order.
+    data
+        Source ID to its data; each task carries one source's data.
+    setup
+        Builds the per-source function; must be picklable.
+    sink
+        Receives each payload, in this process.
+    workers
+        Number of worker processes.
+    """
+    changes = {**_THREAD_ENV, "XLA_FLAGS": _XLA_SINGLE_THREAD}
+    if flags := os.environ.get("XLA_FLAGS"):
+        changes["XLA_FLAGS"] = f"{flags} {_XLA_SINGLE_THREAD}"
+    # Set only while the workers are spawned; this process's environment is
+    # restored afterwards.
+    with patch.dict(os.environ, changes):
+        pool = ProcessPoolExecutor(
+            workers,
+            mp_context=multiprocessing.get_context("spawn"),
+            initializer=_init_worker,
+            initargs=(setup,),
+        )
+        try:
+            pending: set[Future[Payload]] = set()
+            for source_id in ids:
+                if len(pending) >= 2 * workers:
+                    finished, pending = wait(pending, return_when=FIRST_COMPLETED)
+                    for future in finished:
+                        sink(future.result())
+                pending.add(pool.submit(_process_in_worker, source_id, data[source_id]))
+            for future in wait(pending).done:
+                sink(future.result())
+        finally:
+            # On an interruption, drop queued work rather than wait for it.
+            pool.shutdown(wait=True, cancel_futures=True)
+
+
+def mpi_comm() -> Any:
+    """``MPI.COMM_WORLD``, importing ``mpi4py`` only when ``--mpi`` is used.
+
+    Returns
+    -------
+        The world communicator.
+
+    Raises
+    ------
+    ImportError
+        If ``mpi4py`` is not installed.
+    """
+    try:
+        from mpi4py import MPI  # noqa: PLC0415  # ty: ignore[unresolved-import]
+    except ImportError as err:
+        msg = "hq run --mpi needs mpi4py: pip install 'harv-hq[mpi]'"
+        raise ImportError(msg) from err
+    return MPI.COMM_WORLD
+
+
+# The per-source function, built once per worker by _init_worker.
+_worker_process: ProcessFn | None = None
+
+
+def _init_worker(setup: Setup) -> None:
+    global _worker_process  # noqa: PLW0603 -- one per worker process
+    _worker_process = setup()
+
+
+def _process_in_worker(source_id: Any, data: Any) -> Payload:
+    return cast("ProcessFn", _worker_process)(source_id, data)

packages/hq/src/harv_hq/rejection.py

@@ -5,11 +5,13 @@
 hand it across processes.
 """
 
-__all__ = ("process_rejection",)
+__all__ = ("process_rejection", "rejection_setup")
 
+import functools
 import time
 import traceback
 import warnings
+from collections.abc import Callable
 from datetime import UTC, datetime
 from pathlib import Path
 from typing import Any
@@ -18,12 +20,50 @@
 
 from harv.data import GaiaAstrometryData, RVData
 from harv.samplers import RejectionSampler
-from harv_hq.config import RejectionConfig
+from harv_hq._model_file import ModelFile
+from harv_hq.config import Config, RejectionConfig
 from harv_hq.ids import source_key, stable_hash
 from harv_hq.results import Payload, samples_table
 from harv_hq.summary_stats import summary_stats
 
 
+def rejection_setup(
+    config: Config, prior_cache: Path
+) -> Callable[[Any, RVData | GaiaAstrometryData], Payload]:
+    """Build a run's per-source rejection function, once per process.
+
+    Imports the model file, calls ``make_setup()``, and builds one
+    ``RejectionSampler``. Pass it with :func:`functools.partial` as an
+    execution driver's ``setup``, so that pool workers build their own.
+
+    Parameters
+    ----------
+    config
+        The run's configuration.
+    prior_cache
+        Path to the shared prior cache.
+
+    Returns
+    -------
+        ``process(source_id, data) -> Payload``.
+    """
+    prior, model = ModelFile.load(config.run.model_file).setup(config.run.kind)
+    rejection = config.rejection
+    sampler = RejectionSampler(
+        prior,
+        model,
+        batch_size=rejection.batch_size,
+        min_evidence_ess=rejection.min_evidence_ess,
+    )
+    return functools.partial(
+        process_rejection,
+        sampler=sampler,
+        prior_cache=prior_cache,
+        config=rejection,
+        seed=config.run.seed,
+    )
+
+
 def process_rejection(
     source_id: Any,
     data: RVData | GaiaAstrometryData,