Worker pool with disk-based communication for distributed inference

cluster_example/README.md

@@ -0,0 +1,116 @@
+# Cluster distribution example
+
+This directory contains example configuration for distributing miss-alignment
+inference across a SLURM cluster. The same mechanism works for stack preparation
+(`--prepare-stacks`) and cross-correlation pre-alignment (`--preprocess`) phases too.
+
+## Files
+
+- **`cluster_config.json`** — describes how to submit, identify, and cancel cluster jobs.
+  Adapt the commands for your scheduler (SLURM shown; PBS/Torque and others work the
+  same way with different commands).
+- **`worker.sh`** — SLURM submission script template. The `{{command}}` placeholder
+  is filled in automatically with the `miss-alignment worker` invocation. Edit resource
+  requests (`--mem`, `--time`, `--gres`, `--partition`) to match your cluster.
+
+## Setup
+
+1. Point two environment variables at the files in this directory (or copies of them):
+
+   ```bash
+   export MISS_CLUSTER_CONFIG=/path/to/cluster_config.json
+   export MISS_CLUSTER_SCRIPT=/path/to/worker.sh
+   ```
+
+2. Set `MISS_CLUSTER_VAR_partition` to fill in the `{{partition}}` placeholder in
+   `worker.sh` (or hard-code your partition name directly in the script):
+
+   ```bash
+   export MISS_CLUSTER_VAR_partition=gpu
+   ```
+
+3. Add `--n-cluster-workers N` to your `miss-alignment train` or `miss-alignment infer`
+   command. `N` controls how many simultaneous cluster jobs are submitted. A good
+   starting point is one job per GPU node you want to use — each job claims and processes
+   tilt series from the shared queue until the queue is drained.
+
+   ```bash
+   miss-alignment train --config-file config.yaml --n-cluster-workers 8
+   ```
+
+## How it works
+
+The head node (where you run `miss-alignment train/infer`) writes one task JSON file
+per tilt series into `<training_dir>/tasks/pending/`. It then submits `N` cluster
+jobs, each running `miss-alignment worker --queue-dir <training_dir>/tasks --device 0`.
+Workers race to claim tasks by atomic file rename — no scheduler or lock files needed —
+and process series until the queue is empty. The head node blocks until all tasks are
+in `done/` or `failed/`, then continues to the next macro-iteration.
+
+The queue directory (`<training_dir>/tasks/`) must be on a shared filesystem visible
+to all worker nodes (Lustre, GPFS, NFS, etc.). Since the training directory already
+needs to hold the XML metadata and MRC stacks that workers read, this is satisfied
+automatically.
+
+## Adapting for other schedulers
+
+The four required fields in `cluster_config.json`:
+
+| Field | Purpose |
+|---|---|
+| `submit` | Command to submit `{{script_path}}` |
+| `submit_job_id_regex` | Regex (group 1) capturing the job ID from submit stdout |
+| `cancel` | Command to cancel `{{job_id}}` |
+| `status_list` | Command listing all your active jobs; `$USER` is expanded by the shell. Output must be one job per line in `id,STATUS` format. |
+
+The `status_list` command is called each scheduler tick (every 10s) to count alive
+(queued or running) jobs. Status tokens are auto-detected across SLURM, LSF, PBS, and
+SGE. Set `"scheduler": "slurm"` (or `"lsf"`, `"pbs"`, `"sge"`) to skip auto-detection.
+
+### Examples by scheduler
+
+**SLURM** (default):
+```json
+{
+  "submit": "sbatch {{script_path}}",
+  "submit_job_id_regex": "Submitted batch job (\\d+)",
+  "cancel": "scancel {{job_id}}",
+  "status_list": "squeue -u $USER -h -o \"%i,%T\""
+}
+```
+
+**LSF**:
+```json
+{
+  "submit": "bsub < {{script_path}}",
+  "submit_job_id_regex": "Job <(\\d+)> is submitted",
+  "cancel": "bkill {{job_id}}",
+  "status_list": "bjobs -u $USER -noheader -o 'jobid stat'",
+  "scheduler": "lsf"
+}
+```
+
+**PBS/Torque**:
+```json
+{
+  "submit": "qsub {{script_path}}",
+  "submit_job_id_regex": "(\\d+)\\.\\w+",
+  "cancel": "qdel {{job_id}}",
+  "status_list": "qstat -u $USER | awk 'NR>5 {print $1\",\"$10}'",
+  "scheduler": "pbs"
+}
+```
+
+**Custom scheduler** — provide your own alive status tokens:
+```json
+{
+  "submit": "...",
+  "submit_job_id_regex": "(\\d+)",
+  "cancel": "...",
+  "status_list": "...",
+  "scheduler": "custom",
+  "custom_alive_statuses": ["QUEUED", "ACTIVE"]
+}
+```
+
+Update `worker.sh` with the corresponding scheduler directives (`#PBS`, `#BSUB`, etc.).

cluster_example/cluster_config.json

@@ -0,0 +1,6 @@
+{
+  "submit": "sbatch {{script_path}}",
+  "submit_job_id_regex": "Submitted batch job (\\d+)",
+  "cancel": "scancel {{job_id}}",
+  "status_list": "squeue -u $USER -h -o \"%i,%T\""
+}

cluster_example/worker.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+#SBATCH --job-name=miss-alignment-worker
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:1
+#SBATCH --mem=32G
+#SBATCH --time=08:00:00
+#SBATCH --partition={{partition}}
+#SBATCH --output={{logs_dir}}/slurm-%j.out
+#SBATCH --error={{logs_dir}}/slurm-%j.err
+
+# Activate the miss-alignment conda environment.
+# Adjust the path to match your installation.
+source "$(conda info --base)/etc/profile.d/conda.sh"
+conda activate miss-alignment
+
+{{command}}

src/miss_alignment/__init__.py

@@ -14,8 +14,10 @@
     "cli",
     "train_miss_align",
     "infer_miss_align",
+    "worker_miss_align",
 ]
 
 from ._cli import cli
 from .train import train_miss_align
 from .infer import infer_miss_align
+from .distributed.worker import worker_miss_align

src/miss_alignment/__main__.py

@@ -0,0 +1,3 @@
+from miss_alignment import cli
+
+cli()

src/miss_alignment/_cli.py

@@ -11,3 +11,7 @@ def list_commands(self, ctx: Context):
 
 cli = typer.Typer(cls=OrderCommands, add_completion=False, no_args_is_help=True)
 OPTION_PROMPT_KWARGS = {"prompt": True, "prompt_required": True}
+
+from .distributed.worker import worker_miss_align  # noqa: E402
+
+cli.command(name="worker")(worker_miss_align)

src/miss_alignment/alignment/parallel.py

@@ -1,50 +1,6 @@
-import queue
-import torch
-from multiprocessing.managers import BaseProxy
 from pathlib import Path
 
-from .._parallel import run_device_pool
-from .tilt_series import evaluate_tilt_series
-
-
-def gpu_runner(
-    device: int,
-    task_queue: BaseProxy,
-    result_queue: BaseProxy,
-) -> None:
-    """Start a GPU runner, each runner should be initialized to a
-    multiprocessing.Process() and manage running jobs on a single GPU. Each runner will
-    grab jobs from the task_queue and assign jobs to the result_queue once they finish.
-    When the task_queue is empty the gpu_runner will stop.
-
-    Parameters
-    ----------
-    device: int
-        a GPU index to assign to the runner
-    task_queue: mp.managers.BaseProxy
-        shared queue from multiprocessing with jobs to run
-    result_queue: mp.manager.BaseProxy
-        shared queue from multiprocessing for finished jobs
-    """
-    torch.set_num_threads(1)
-    cuda_device = f"cuda:{device}"
-    while True:
-        try:
-            task_parameters = task_queue.get_nowait()
-            tilt_series_path, loss_values = evaluate_tilt_series(
-                **task_parameters,
-                device=cuda_device,
-            )
-            # place the name and final loss of the finished tilt_series
-            final_loss = float(loss_values[-1]) if loss_values else None
-            result_queue.put_nowait(
-                {
-                    "name": tilt_series_path.stem,
-                    "final_loss": final_loss,
-                }
-            )
-        except queue.Empty:
-            break
+from ..distributed.manager import run_distributed
 
 
 def run_alignment_parallel(
@@ -58,52 +14,30 @@ def run_alignment_parallel(
     apply_ctf: bool,
     downsample: int,
     devices_list: list[int],
+    n_cluster_workers: int | None = None,
 ) -> dict[str, float]:
-    """Run a job in parallel over a single or multiple GPUs. If no volume_splits are
-    given the search is parallelized by splitting the angular search. If volume_splits
-    are provided the job will first be split by volume, if there are still more GPUs
-    available, the subvolume jobs are still further split by angular search.
+    """Distribute per-tilt-series alignment across local GPUs or a cluster.
 
-    Parameters
-    ----------
-    model_checkpoint: Path
-    tilt_series_list: list[Path]
-    patches_per_dim: tuple[int, int, int]
-    patch_size: int
-    tomogram_shape: tuple[int, int, int]
-    output_directory: Path
-    devices_list: list[int]
-    ground_truth_list: list[Path]
+    Without --n-cluster-workers, one worker subprocess is spawned per GPU in
+    devices_list (local mode, unchanged behaviour). Set --n-cluster-workers N
+    to submit N cluster jobs instead; requires MISS_CLUSTER_CONFIG and
+    MISS_CLUSTER_SCRIPT to be set.
 
-    Returns
-    -------
-    dict[str, float]
-        Dictionary mapping tilt-series names to their final loss values.
+    Returns dict mapping tilt-series stem names to their final loss values.
     """
-    jobs = [
-        {
-            "model_checkpoint_path": model_checkpoint,
-            "tilt_series_path": tilt_series,
-            "output_directory": output_directory,
-            "setting": setting,
-            "patch_size": patch_size,
-            "patch_overlap": patch_overlap,
-            "batch_size": batch_size,
-            "apply_ctf": apply_ctf,
-            "downsample": downsample,
-        }
-        for tilt_series in tilt_series_list
-    ]
-
-    # one worker process per unique GPU, each pulling jobs from a shared queue
-    results = run_device_pool(
-        jobs=jobs,
-        runner=gpu_runner,
-        runner_args=(),
+    queue_root = output_directory / "tasks"
+
+    return run_distributed(
+        tilt_series_list=tilt_series_list,
+        model_checkpoint=model_checkpoint,
+        output_directory=output_directory,
+        setting=setting,
+        patch_size=patch_size,
+        patch_overlap=patch_overlap,
+        batch_size=batch_size,
+        apply_ctf=apply_ctf,
+        downsample=downsample,
         devices=devices_list,
-        desc="Tilt series alignment",
+        n_cluster_workers=n_cluster_workers,
+        queue_root=queue_root,
     )
-
-    # Convert results to dictionary of losses
-    losses = {result["name"]: result["final_loss"] for result in results}
-    return losses