diff --git a/cluster_example/README.md b/cluster_example/README.md new file mode 100644 index 00000000..db3cb18f --- /dev/null +++ b/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 `/tasks/pending/`. It then submits `N` cluster +jobs, each running `miss-alignment worker --queue-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 (`/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.). diff --git a/cluster_example/cluster_config.json b/cluster_example/cluster_config.json new file mode 100644 index 00000000..15fe8c7e --- /dev/null +++ b/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\"" +} diff --git a/cluster_example/worker.sh b/cluster_example/worker.sh new file mode 100644 index 00000000..bf8a8b7d --- /dev/null +++ b/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}} diff --git a/src/miss_alignment/__init__.py b/src/miss_alignment/__init__.py index 8745d7f6..dface4e6 100644 --- a/src/miss_alignment/__init__.py +++ b/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 diff --git a/src/miss_alignment/__main__.py b/src/miss_alignment/__main__.py new file mode 100644 index 00000000..7f9e2d32 --- /dev/null +++ b/src/miss_alignment/__main__.py @@ -0,0 +1,3 @@ +from miss_alignment import cli + +cli() diff --git a/src/miss_alignment/_cli.py b/src/miss_alignment/_cli.py index d89206d8..a304641f 100644 --- a/src/miss_alignment/_cli.py +++ b/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) diff --git a/src/miss_alignment/_parallel.py b/src/miss_alignment/_parallel.py deleted file mode 100644 index e756d43d..00000000 --- a/src/miss_alignment/_parallel.py +++ /dev/null @@ -1,83 +0,0 @@ -"""Shared one-process-per-GPU work-queue helper. - -Mirrors the scheme used by ``alignment.run_alignment_parallel``: start exactly -one worker process per unique device, feed all jobs through a shared queue, and -let each worker pull jobs until the queue is empty. This binds a device to a -*process* (deterministic, one job per GPU at a time) instead of to a task index, -and uses every available device rather than a fixed process count. -""" - -import multiprocessing as mp -import sys -import time -from collections.abc import Callable -from typing import Any - -import tqdm - - -def run_device_pool( - jobs: list[Any], - runner: Callable, - runner_args: tuple, - devices: list[int] | None, - desc: str, -) -> list[Any]: - """Run ``jobs`` across one worker process per unique device. - - Each worker runs ``runner(device, task_queue, result_queue, *runner_args)``, - pulling jobs off ``task_queue`` until empty and putting one result on - ``result_queue`` per finished job. One process is started per unique entry in - ``devices`` (or a single default-device process when ``devices`` is falsy). - - Returns the list of results collected from the workers. Raises ``RuntimeError`` - if any worker exits with a non-zero code (all workers are then terminated). - """ - ctx = mp.get_context("spawn") - device_slots = sorted(set(devices)) if devices else [None] - - with ctx.Manager() as manager: - task_queue = manager.Queue() - result_queue = manager.Queue() - for job in jobs: - task_queue.put_nowait(job) - - procs = [ - ctx.Process( - target=runner, - args=(device, task_queue, result_queue, *runner_args), - ) - for device in device_slots - ] - [p.start() for p in procs] - - results: list[Any] = [] - pbar = tqdm.tqdm( - total=len(jobs), - desc=desc, - file=sys.stdout, - ) - while len(results) < len(jobs): - while not result_queue.empty(): - results.append(result_queue.get_nowait()) - pbar.update(1) - - for p in procs: - # a worker that died with a non-zero exit code means a job failed; - # tear everything down rather than hang waiting for its result - if not p.is_alive() and p.exitcode != 0: - for x in procs: - x.terminate() - for x in procs: - x.join(timeout=5.0) - pbar.close() - raise RuntimeError( - f"A worker process for '{desc}' stopped unexpectedly." - ) - - time.sleep(0.1) - - pbar.close() - [p.join() for p in procs] - - return results diff --git a/src/miss_alignment/alignment/parallel.py b/src/miss_alignment/alignment/parallel.py index a2161325..ba691ca9 100644 --- a/src/miss_alignment/alignment/parallel.py +++ b/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 diff --git a/src/miss_alignment/alignment/tilt_series.py b/src/miss_alignment/alignment/tilt_series.py index d55267c5..dcdec76f 100644 --- a/src/miss_alignment/alignment/tilt_series.py +++ b/src/miss_alignment/alignment/tilt_series.py @@ -128,6 +128,7 @@ def evaluate_tilt_series( device: str = "cpu", initial_reliable_fraction: float = 1 / 2, n_control_points: int = 7, + model: MissAlignment | None = None, ) -> tuple[Path, list[float]]: """Evaluate and optimize tilt series alignment using trained model. @@ -168,21 +169,23 @@ def evaluate_tilt_series( tuple[Path, list[float]] Path to output JSON and list of loss values. """ - # load the best model and run alignment optimization - model = MissAlignment.load_from_checkpoint( - model_checkpoint_path, - map_location="cpu", - ) - # load_from_checkpoint calls configure_model(), which compiles self.net via - # torch.compile. Unwrap it here so alignment inference runs in eager mode. - # torch.compile is fundamentally incompatible with spawned worker processes: - # the inductor spawns its own subprocess pool for async compilation and races - # on shared cache files when multiple workers compile simultaneously, causing - # stochastic FileNotFoundError crashes. See pytorch/pytorch#134384. - # The reconstruction step dominates alignment runtime so the eager-mode - # overhead on the model forward pass is negligible. - if hasattr(model.net, "_orig_mod"): - model.net = model.net._orig_mod + # Load model from checkpoint when not provided by the caller. Workers pass a + # pre-loaded model to amortize checkpoint loading across many series per run. + if model is None: + model = MissAlignment.load_from_checkpoint( + model_checkpoint_path, + map_location="cpu", + ) + # load_from_checkpoint calls configure_model(), which compiles self.net via + # torch.compile. Unwrap it here so alignment inference runs in eager mode. + # torch.compile is fundamentally incompatible with spawned worker processes: + # the inductor spawns its own subprocess pool for async compilation and races + # on shared cache files when multiple workers compile simultaneously, causing + # stochastic FileNotFoundError crashes. See pytorch/pytorch#134384. + # The reconstruction step dominates alignment runtime so the eager-mode + # overhead on the model forward pass is negligible. + if hasattr(model.net, "_orig_mod"): + model.net = model.net._orig_mod # load tilt_series and set its name for output tilt_series_data = TiltSeriesData(xml_metadata_path=tilt_series_path) diff --git a/src/miss_alignment/data/training_datamodule.py b/src/miss_alignment/data/training_datamodule.py index 7a36f762..f20c3dba 100644 --- a/src/miss_alignment/data/training_datamodule.py +++ b/src/miss_alignment/data/training_datamodule.py @@ -215,7 +215,7 @@ def setup(self, stage: str = "fit"): self.pool_dir.mkdir(parents=True, exist_ok=True) print(f"Created pool directory: {self.pool_dir}") print( - f"Pool configuration: {self.n_workers} workers -> " + f"Pool configuration: {self.n_workers} reconstruction workers, " f"{self.n_partitions} partitions ({self.partition_size} files each)" ) diff --git a/src/miss_alignment/distributed/__init__.py b/src/miss_alignment/distributed/__init__.py new file mode 100644 index 00000000..1b05278b --- /dev/null +++ b/src/miss_alignment/distributed/__init__.py @@ -0,0 +1,41 @@ +"""Disk-based distributed task queue for miss-alignment inference.""" + +from .config import ClusterConfig, load_cluster_config +from .manager import run_distributed +from .provisioner import ( + ClusterProvisioner, + CompositeProvisioner, + LocalProvisioner, + WorkerProvisioner, +) +from .queue import ( + QueueLayout, + TaskSpec, + claim_one, + clear_queue, + compute_fingerprint, + mark_done, + mark_failed, + write_pending, +) +from .worker import run_worker_loop, worker_miss_align + +__all__ = [ + "ClusterConfig", + "load_cluster_config", + "run_distributed", + "ClusterProvisioner", + "CompositeProvisioner", + "LocalProvisioner", + "WorkerProvisioner", + "QueueLayout", + "TaskSpec", + "claim_one", + "clear_queue", + "compute_fingerprint", + "mark_done", + "mark_failed", + "write_pending", + "run_worker_loop", + "worker_miss_align", +] diff --git a/src/miss_alignment/distributed/config.py b/src/miss_alignment/distributed/config.py new file mode 100644 index 00000000..fe483410 --- /dev/null +++ b/src/miss_alignment/distributed/config.py @@ -0,0 +1,67 @@ +"""Read cluster configuration from environment variables. + +load_cluster_config() is called only when --n-cluster-workers is set. +It raises RuntimeError immediately if either required env var is absent, +so the user gets a clear error rather than a silent fallback to local mode. +""" + +from __future__ import annotations + +import json +import os +from dataclasses import dataclass, field +from pathlib import Path + + +@dataclass +class ClusterConfig: + submit: str + submit_job_id_regex: str + cancel: str + status_list: str + script_path: Path + # Optional: 'slurm', 'lsf', 'pbs', 'sge', or 'custom'. + # When absent the status output is auto-detected by trying all parsers. + scheduler: str = "auto" + # Custom parser fields — only used when scheduler='custom'. + custom_alive_statuses: list[str] = field(default_factory=list) + custom_status_regex: str = "" # captures (job_id, status) per line + + +def load_cluster_config() -> ClusterConfig: + """Return ClusterConfig. Raises RuntimeError if env vars are missing.""" + config_path_str = os.environ.get("MISS_CLUSTER_CONFIG") + if not config_path_str: + raise RuntimeError( + "MISS_CLUSTER_CONFIG environment variable is required when " + "--n-cluster-workers is set. Point it to a JSON file with " + "'submit', 'submit_job_id_regex', 'cancel', and 'status_list' keys." + ) + + script_path_str = os.environ.get("MISS_CLUSTER_SCRIPT") + if not script_path_str: + raise RuntimeError( + "MISS_CLUSTER_SCRIPT environment variable is required when " + "--n-cluster-workers is set. Point it to a shell script template " + "containing a {{command}} placeholder." + ) + + config_path = Path(config_path_str) + script_path = Path(script_path_str) + + if not config_path.exists(): + raise FileNotFoundError(f"MISS_CLUSTER_CONFIG not found: {config_path}") + if not script_path.exists(): + raise FileNotFoundError(f"MISS_CLUSTER_SCRIPT not found: {script_path}") + + data = json.loads(config_path.read_text()) + return ClusterConfig( + submit=data["submit"], + submit_job_id_regex=data["submit_job_id_regex"], + cancel=data["cancel"], + status_list=data["status_list"], + script_path=script_path, + scheduler=data.get("scheduler", "auto"), + custom_alive_statuses=data.get("custom_alive_statuses", []), + custom_status_regex=data.get("custom_status_regex", ""), + ) diff --git a/src/miss_alignment/distributed/manager.py b/src/miss_alignment/distributed/manager.py new file mode 100644 index 00000000..497463dc --- /dev/null +++ b/src/miss_alignment/distributed/manager.py @@ -0,0 +1,288 @@ +"""Head-node coordinator: writes tasks, starts provisioner, blocks until done.""" + +from __future__ import annotations + +import json +import os +import shutil +import sys +import threading +import time +from pathlib import Path + +import tqdm + +from .config import load_cluster_config +from .provisioner import ( + ClusterProvisioner, + CompositeProvisioner, + LocalProvisioner, + WorkerProvisioner, +) +from .queue import ( + QueueLayout, + TaskSpec, + clear_queue, + compute_fingerprint, + write_pending, +) + +_POLL_INTERVAL_S = 0.5 +_SCHEDULER_INTERVAL_S = 10.0 +_MANAGER_HB_INTERVAL_S = 5.0 +_WORKER_STALL_TIMEOUT_S = 120.0 + + +def _format_task_id(index: int, tilt_series_path: Path) -> str: + return f"{index:07d}-{tilt_series_path.stem}" + + +def _write_manager_hb(layout: QueueLayout, seq: int) -> None: + new_hb = layout.manager_hb / f"hb-{seq}" + new_hb.write_text("") + if seq > 0: + (layout.manager_hb / f"hb-{seq - 1}").unlink(missing_ok=True) + + +def _sweep_stalled_workers(layout: QueueLayout) -> None: + for worker_dir in layout.running.iterdir(): + if not worker_dir.is_dir(): + continue + ticks = list(worker_dir.glob("hb-*")) + latest_mtime = None + for tick in ticks: + try: + mtime = tick.stat().st_mtime + if latest_mtime is None or mtime > latest_mtime: + latest_mtime = mtime + except FileNotFoundError: + pass # heartbeat thread replaced this tick; ignore + if latest_mtime is not None: + age = time.time() - latest_mtime + else: + age = time.time() - worker_dir.stat().st_mtime + + if age <= _WORKER_STALL_TIMEOUT_S: + continue + + for task_file in worker_dir.glob("*.json"): + # Skip tasks that already completed while the worker was being swept. + if (layout.done / task_file.name).exists(): + task_file.unlink(missing_ok=True) + continue + if (layout.failed / task_file.name).exists(): + task_file.unlink(missing_ok=True) + continue + dest = layout.pending / task_file.name + try: + os.rename(task_file, dest) + except FileNotFoundError: + pass + for hb in worker_dir.glob("hb-*"): + hb.unlink(missing_ok=True) + try: + worker_dir.rmdir() + except OSError: + pass + + +def _format_worker_counts(counts: dict[str, int]) -> str: + return " ".join(f"{label}={n}" for label, n in sorted(counts.items())) + + +def _scheduler_thread( + layout: QueueLayout, + provisioner: WorkerProvisioner, + n_workers: int, + stop_event: threading.Event, + error_box: list, + worker_counts: dict, +) -> None: + hb_seq = 1 # seq 0 written before thread starts + last_hb = time.time() + # Start from now so the first sweep is delayed by _SCHEDULER_INTERVAL_S; + # run_distributed calls ensure_workers() explicitly at startup instead. + last_sweep = time.time() + + try: + while not stop_event.is_set(): + now = time.time() + + if now - last_hb >= _MANAGER_HB_INTERVAL_S: + _write_manager_hb(layout, hb_seq) + hb_seq += 1 + last_hb = now + + if now - last_sweep >= _SCHEDULER_INTERVAL_S: + _sweep_stalled_workers(layout) + n_pending = len(list(layout.pending.glob("*.json"))) + if n_pending > 0: + # Cap at n_pending: no point submitting more workers than tasks. + # This also prevents resubmission when workers exit cleanly on + # an empty queue — only re-fill if the sweep re-pended orphans. + provisioner.ensure_workers(min(n_workers, n_pending)) + worker_counts.clear() + worker_counts.update(provisioner.worker_counts_by_type()) + last_sweep = now + + stop_event.wait(timeout=1.0) + except Exception as exc: + error_box.append(exc) + stop_event.set() # wake the poll loop so it notices immediately + + +def run_distributed( + tilt_series_list: list[Path], + model_checkpoint: Path, + output_directory: Path, + setting: str | tuple, + patch_size: int, + patch_overlap: float, + batch_size: int, + apply_ctf: bool, + downsample: int, + devices: list[int], + n_cluster_workers: int | None, + queue_root: Path, + task_type: str = "alignment", + desired_pixel_size: float | None = None, + lowpass_cutoff: float | None = None, +) -> dict[str, float]: + """Write tasks, provision workers, block until all tasks are terminal. + + Returns dict[series_name → final_loss]. Raises RuntimeError listing all + failed series if any task ends in failed/. Deletes queue_root on exit. + + task_type selects the worker dispatch path: + "alignment" — evaluate_tilt_series (default) + "prepare_stacks" — _prepare_single_tilt_series + "cross_correlation" — _run_cross_correlation_single + """ + layout = QueueLayout(queue_root) + layout.ensure_directories() + clear_queue(layout) + + # Fingerprint only meaningful for alignment tasks (amortizes model load). + if task_type == "alignment": + fingerprint = compute_fingerprint( + model_checkpoint_path=str(model_checkpoint), + setting=setting if isinstance(setting, str) else list(setting), + patch_size=patch_size, + patch_overlap=patch_overlap, + batch_size=batch_size, + apply_ctf=apply_ctf, + downsample=downsample, + ) + else: + fingerprint = "" + + task_ids = [] + for i, ts_path in enumerate(tilt_series_list): + task_id = _format_task_id(i, ts_path) + task_ids.append(task_id) + spec = TaskSpec( + task_id=task_id, + model_checkpoint_path=( + str(model_checkpoint) if task_type == "alignment" else "" + ), + tilt_series_path=str(ts_path), + output_directory=str(output_directory), + setting=setting if isinstance(setting, str) else list(setting), + patch_size=patch_size, + patch_overlap=patch_overlap, + batch_size=batch_size, + apply_ctf=apply_ctf, + downsample=downsample, + init_fingerprint=fingerprint, + task_type=task_type, + desired_pixel_size=desired_pixel_size, + lowpass_cutoff=lowpass_cutoff, + ) + write_pending(layout, spec) + + # Write first heartbeat before starting workers so workers never see a + # missing heartbeat on startup (would cause immediate exit). + _write_manager_hb(layout, seq=0) + + if n_cluster_workers is not None: + cluster_config = load_cluster_config() + cluster = ClusterProvisioner(queue_dir=queue_root, config=cluster_config) + if devices: + # Also use local GPUs — the queue is shared so local and cluster + # workers pull from the same task list with no extra coordination. + local = LocalProvisioner(queue_dir=queue_root, devices=devices) + provisioner: WorkerProvisioner = CompositeProvisioner([cluster, local]) + else: + provisioner = cluster + n_workers = n_cluster_workers + else: + provisioner = LocalProvisioner(queue_dir=queue_root, devices=devices) + n_workers = len(devices) + + stop_event = threading.Event() + scheduler_errors: list = [] + worker_counts: dict = {} + scheduler = threading.Thread( + target=_scheduler_thread, + args=( + layout, provisioner, n_workers, stop_event, scheduler_errors, worker_counts + ), + daemon=True, + ) + scheduler.start() + provisioner.ensure_workers(n_workers) + + pending_ids = set(task_ids) + losses: dict[str, float] = {} + failed_series: list[str] = [] + + _desc = { + "alignment": "Tilt series alignment", + "prepare_stacks": "Preparing stacks", + "cross_correlation": "Cross-correlation alignment", + }.get(task_type, task_type) + pbar = tqdm.tqdm(total=len(task_ids), desc=_desc, file=sys.stdout) + try: + while pending_ids: + time.sleep(_POLL_INTERVAL_S) + + if scheduler_errors: + raise RuntimeError("Scheduler thread crashed") from scheduler_errors[0] + + pbar.set_postfix_str(_format_worker_counts(worker_counts)) + + for done_file in layout.done.glob("*.json"): + data = json.loads(done_file.read_text()) + tid = data["task_id"] + if tid in pending_ids: + ts_name = Path(data["tilt_series_path"]).stem + losses[ts_name] = data.get("final_loss", float("nan")) + pending_ids.discard(tid) + pbar.update(1) + + for fail_file in layout.failed.glob("*.json"): + data = json.loads(fail_file.read_text()) + tid = data["task_id"] + if tid in pending_ids: + ts_name = Path(data["tilt_series_path"]).stem + failed_series.append(ts_name) + pending_ids.discard(tid) + pbar.update(1) + print( + f"[manager] FAILED {ts_name}: {data.get('error', '')}", + file=sys.stderr, + ) + finally: + pbar.close() + stop_event.set() + scheduler.join(timeout=5.0) + provisioner.shutdown() + shutil.rmtree(queue_root, ignore_errors=True) + + if failed_series: + raise RuntimeError( + f"Alignment failed for {len(failed_series)} tilt series: " + + ", ".join(failed_series) + ) + + return losses diff --git a/src/miss_alignment/distributed/provisioner.py b/src/miss_alignment/distributed/provisioner.py new file mode 100644 index 00000000..93ae8fac --- /dev/null +++ b/src/miss_alignment/distributed/provisioner.py @@ -0,0 +1,392 @@ +"""Worker provisioners: spawn local child processes or submit cluster jobs.""" + +from __future__ import annotations + +import os +import re +import subprocess +import sys +import threading +import time +from abc import ABC, abstractmethod +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path + +from .config import ClusterConfig + +# A worker dir with no heartbeat file yet is assumed alive for this long +# after the dir was created (worker is still starting up). +_STARTUP_GRACE_S = 60.0 +# A worker dir whose newest heartbeat is older than this is considered dead. +_STALL_TIMEOUT_S = 120.0 + + +def _live_worker_dirs(running_dir: Path) -> int: + """Count running// subdirs whose heartbeat is fresh enough to be alive.""" + if not running_dir.exists(): + return 0 + count = 0 + now = time.time() + for wdir in running_dir.iterdir(): + if not wdir.is_dir(): + continue + ticks = list(wdir.glob("hb-*")) + if ticks: + newest_mtime = None + for t in ticks: + try: + mtime = t.stat().st_mtime + if newest_mtime is None or mtime > newest_mtime: + newest_mtime = mtime + except FileNotFoundError: + pass # heartbeat thread rotated this tick; ignore + if newest_mtime is not None and now - newest_mtime < _STALL_TIMEOUT_S: + count += 1 + else: + # No heartbeat yet — consider alive if dir was created recently. + try: + if now - wdir.stat().st_mtime < _STARTUP_GRACE_S: + count += 1 + except FileNotFoundError: + pass # dir vanished between iterdir and stat; skip + return count + + +class WorkerProvisioner(ABC): + @abstractmethod + def ensure_workers(self, n_workers: int) -> None: + """Ensure workers are running. + + Called once at startup and on each scheduler tick to respawn dead workers. + """ + + @abstractmethod + def shutdown(self) -> None: + """Terminate all managed workers.""" + + @abstractmethod + def live_worker_count(self) -> int: + """Return the number of workers currently considered alive.""" + + def worker_counts_by_type(self) -> dict[str, int]: + """Return a dict of label → live count for display purposes.""" + return {"workers": self.live_worker_count()} + + +class LocalProvisioner(WorkerProvisioner): + """Spawns one miss-alignment worker subprocess per GPU device.""" + + def __init__(self, queue_dir: Path, devices: list[int]) -> None: + self._queue_dir = queue_dir + self._devices = devices + self._procs: dict[int, subprocess.Popen] = {} + self._lock = threading.Lock() + + def live_worker_count(self) -> int: + with self._lock: + return sum(1 for p in self._procs.values() if p.poll() is None) + + def worker_counts_by_type(self) -> dict[str, int]: + return {"local": self.live_worker_count()} + + def ensure_workers(self, n_workers: int) -> None: + with self._lock: + for device in self._devices: + proc = self._procs.get(device) + if proc is not None and proc.poll() is None: + continue # still running + new_proc = subprocess.Popen( + [ + sys.executable, + "-m", + "miss_alignment", + "worker", + "--queue-dir", + str(self._queue_dir), + "--device", + str(device), + ], + stdout=subprocess.DEVNULL, + stderr=sys.stderr, + ) + self._procs[device] = new_proc + + def shutdown(self) -> None: + for proc in self._procs.values(): + if proc.poll() is None: + proc.terminate() + for proc in self._procs.values(): + try: + proc.wait(timeout=10.0) + except subprocess.TimeoutExpired: + proc.kill() + self._procs.clear() + + +# --------------------------------------------------------------------------- +# Per-scheduler status parsers. +# Each entry maps a comma-separated "id,STATUS" token to alive/dead. +# The status_list command is expected to emit one "id,STATUS" pair per line +# (e.g. SLURM: squeue -u $USER -h -o "%i,%T"). +# --------------------------------------------------------------------------- + +# Statuses where the job is actively executing on a node. +_RUNNING_STATUSES: dict[str, set[str]] = { + "slurm": {"RUNNING", "R", "COMPLETING", "CG", "RESIZING"}, + "lsf": {"RUN"}, + "pbs": {"R", "E"}, + "sge": {"r", "t", "Rr"}, +} + +# Statuses where the job is alive but not yet on a node. +_PENDING_STATUSES: dict[str, set[str]] = { + "slurm": {"PENDING", "PD", "SUSPENDED", "S"}, + "lsf": {"PEND", "SSUSP", "USUSP", "PSUSP"}, + "pbs": {"Q", "H", "W", "T", "S"}, + "sge": {"qw", "Rq", "hqw", "hRwq"}, +} + + +# Jobs absent from squeue for longer than this are considered gone. +# Must be long enough to cover scheduler registration delay (usually <30s) +# plus a full squeue poll cycle. +_JOB_GRACE_S = 120.0 + + +def _parse_status_output( + output: str, + our_job_ids: set[str], + scheduler: str, + custom_alive_statuses: list[str], + custom_status_regex: str, +) -> dict[str, str]: + """Parse status_list output; return dict[job_id → 'running'|'pending']. + + Only jobs from our_job_ids that appear in the output are returned. + Jobs absent from the output are not included — the caller decides + whether to prune them based on how long they have been absent. + """ + result: dict[str, str] = {} + + if scheduler == "auto": + schedulers_to_try = ["slurm", "lsf", "pbs", "sge"] + elif scheduler == "custom": + schedulers_to_try = [] + else: + schedulers_to_try = [scheduler] + + for line in output.splitlines(): + line = line.strip() + if not line: + continue + + if "," in line: + job_id, _, status_token = line.partition(",") + job_id = job_id.strip() + status_token = status_token.strip() + else: + job_id = line + status_token = "" + + if job_id not in our_job_ids: + continue + + if scheduler == "custom": + if status_token in custom_alive_statuses or not status_token: + result[job_id] = "pending" + # Otherwise not recognised as alive — leave absent from result. + continue + + if not status_token: + result[job_id] = "pending" + continue + + for sched in schedulers_to_try: + if status_token in _RUNNING_STATUSES[sched]: + result[job_id] = "running" + break + if status_token in _PENDING_STATUSES[sched]: + result[job_id] = "pending" + break + # Jobs with unrecognised status are simply absent from the result. + + return result + + +class ClusterProvisioner(WorkerProvisioner): + """Submits cluster jobs; tracks liveness by querying the batch scheduler. + + The status_list command (configured in cluster_config.json) is called each + scheduler tick to get the set of alive jobs. Supports SLURM, LSF, PBS, SGE, + and custom schedulers. Scheduler type is auto-detected from status output + unless 'scheduler' is set explicitly in the config. + + Jobs absent from squeue output are kept for _JOB_GRACE_S seconds before + being pruned — this covers the window between submission and scheduler + registration, which can be tens of seconds on busy clusters. + + Recommended status_list format (one job per line, id,STATUS): + SLURM: squeue -u $USER -h -o "%i,%T" + """ + + def __init__(self, queue_dir: Path, config: ClusterConfig) -> None: + self._queue_dir = queue_dir + self._config = config + # job_id → submission timestamp + self._job_submit_time: dict[str, float] = {} + # Monotonically increasing index for unique script names so indices + # never repeat even after preemption-triggered resubmissions. + self._next_index: int = 0 + self._scripts_dir = queue_dir / "cluster" + self._scripts_dir.mkdir(parents=True, exist_ok=True) + + @property + def _job_ids(self) -> list[str]: + return list(self._job_submit_time.keys()) + + def _query_job_states(self) -> dict[str, str]: + """Query the scheduler; return dict[job_id → 'running'|'pending']. + + Jobs absent from squeue are pruned only after _JOB_GRACE_S seconds, + allowing for scheduler registration delay. On error, returns all + tracked jobs as 'pending'. + """ + if not self._job_submit_time: + return {} + our_ids = set(self._job_submit_time.keys()) + status_cmd = self._config.status_list.replace( + "{{user}}", os.environ.get("USER", os.environ.get("USERNAME", "")) + ) + try: + result = subprocess.run( + status_cmd, + shell=True, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + timeout=30, + ) + states = _parse_status_output( + output=result.stdout, + our_job_ids=our_ids, + scheduler=self._config.scheduler, + custom_alive_statuses=self._config.custom_alive_statuses, + custom_status_regex=self._config.custom_status_regex, + ) + now = time.time() + # Prune jobs absent from squeue only after the grace period. + for jid in list(self._job_submit_time): + if jid not in states: + age = now - self._job_submit_time[jid] + if age > _JOB_GRACE_S: + del self._job_submit_time[jid] + return states + except Exception: + return {jid: "pending" for jid in self._job_submit_time} + + def live_worker_count(self) -> int: + return len(self._job_submit_time) + + def worker_counts_by_type(self) -> dict[str, int]: + states = self._query_job_states() + running = sum(1 for s in states.values() if s == "running") + # pending = squeue-visible pending + not-yet-visible (within grace period) + visible = set(states.keys()) + pending = sum(1 for s in states.values() if s == "pending") + pending += sum(1 for jid in self._job_submit_time if jid not in visible) + counts = {} + if running: + counts["cluster-running"] = running + if pending: + counts["cluster-pending"] = pending + return counts + + def ensure_workers(self, n_workers: int) -> None: + """Submit new jobs until tracked job count reaches n_workers. + + Prunes grace-expired absent jobs first, then submits the deficit. + """ + self._query_job_states() + deficit = n_workers - len(self._job_submit_time) + for _ in range(deficit): + self._submit_one() + + def _render_script(self, index: int) -> Path: + template_text = self._config.script_path.read_text() + # $(hostname) and $$ are expanded by the compute node's shell at runtime. + command = ( + f"miss-alignment worker" + f" --queue-dir {self._queue_dir}" + f" --device 0" + f' --worker-id "$(hostname)-$$-{index}"' + ) + logs_dir = self._queue_dir / "logs" + logs_dir.mkdir(parents=True, exist_ok=True) + rendered = template_text.replace("{{command}}", command) + rendered = rendered.replace("{{tasks_dir}}", str(self._queue_dir)) + rendered = rendered.replace("{{logs_dir}}", str(logs_dir)) + for key, value in os.environ.items(): + if key.startswith("MISS_CLUSTER_VAR_"): + var_name = key[len("MISS_CLUSTER_VAR_") :].lower() + rendered = rendered.replace(f"{{{{{var_name}}}}}", value) + + script_path = self._scripts_dir / f"worker-{index}.sh" + script_path.write_text(rendered) + return script_path + + def _submit_one(self) -> None: + index = self._next_index + self._next_index += 1 + script_path = self._render_script(index) + submit_cmd = self._config.submit.replace( + "{{script_path}}", str(script_path) + ) + result = subprocess.run( + submit_cmd, + shell=True, + stdout=subprocess.PIPE, + stderr=sys.stderr, + text=True, + ) + match = re.search(self._config.submit_job_id_regex, result.stdout) + if match: + self._job_submit_time[match.group(1)] = time.time() + + def shutdown(self) -> None: + def _cancel(job_id: str) -> None: + cancel_cmd = self._config.cancel.replace("{{job_id}}", job_id) + subprocess.run(cancel_cmd, shell=True, stderr=subprocess.DEVNULL) + + job_ids = list(self._job_submit_time.keys()) + with ThreadPoolExecutor(max_workers=min(32, len(job_ids) or 1)) as pool: + list(as_completed([pool.submit(_cancel, jid) for jid in job_ids])) + self._job_submit_time.clear() + + +class CompositeProvisioner(WorkerProvisioner): + """Delegates to multiple provisioners simultaneously. + + Used to run local GPU workers alongside cluster workers so local GPUs + are not left idle during cluster-distributed alignment phases. + """ + + def __init__(self, provisioners: list[WorkerProvisioner]) -> None: + self._provisioners = provisioners + + def live_worker_count(self) -> int: + return sum(p.live_worker_count() for p in self._provisioners) + + def worker_counts_by_type(self) -> dict[str, int]: + counts: dict[str, int] = {} + for p in self._provisioners: + for label, n in p.worker_counts_by_type().items(): + counts[label] = counts.get(label, 0) + n + return counts + + def ensure_workers(self, n_workers: int) -> None: + for p in self._provisioners: + p.ensure_workers(n_workers) + + def shutdown(self) -> None: + for p in self._provisioners: + p.shutdown() diff --git a/src/miss_alignment/distributed/queue.py b/src/miss_alignment/distributed/queue.py new file mode 100644 index 00000000..5238e408 --- /dev/null +++ b/src/miss_alignment/distributed/queue.py @@ -0,0 +1,208 @@ +"""Filesystem queue: task JSON files + atomic-rename claim protocol. + +Directory layout under /: + pending/ one JSON per queued task + running// claimed task + heartbeat ticks + done/ completed task JSONs (result fields appended) + failed/ failed task JSONs (error field appended) + manager/ manager heartbeat ticks + cluster/ rendered cluster submission scripts +""" + +from __future__ import annotations + +import hashlib +import json +import os +import random +from dataclasses import asdict, dataclass +from pathlib import Path + + +@dataclass +class QueueLayout: + root: Path + + @property + def pending(self) -> Path: + return self.root / "pending" + + @property + def running(self) -> Path: + return self.root / "running" + + @property + def done(self) -> Path: + return self.root / "done" + + @property + def failed(self) -> Path: + return self.root / "failed" + + @property + def manager_hb(self) -> Path: + return self.root / "manager" + + @property + def cluster(self) -> Path: + return self.root / "cluster" + + @property + def logs(self) -> Path: + return self.root / "logs" + + def ensure_directories(self) -> None: + for d in ( + self.pending, + self.running, + self.done, + self.failed, + self.manager_hb, + self.cluster, + self.logs, + ): + d.mkdir(parents=True, exist_ok=True) + + def worker_dir(self, worker_id: str) -> Path: + return self.running / worker_id + + +@dataclass +class TaskSpec: + task_id: str + model_checkpoint_path: str + tilt_series_path: str + output_directory: str + setting: str | list + patch_size: int + patch_overlap: float + batch_size: int + apply_ctf: bool + downsample: int + init_fingerprint: str + # Task type and optional parameters for non-alignment tasks. + task_type: str = "alignment" + desired_pixel_size: float | None = None + lowpass_cutoff: float | None = None + + +def _atomic_write(path: Path, data: dict) -> None: + """Write JSON atomically via a temp file + rename.""" + tmp = path.with_suffix(f".tmp.{os.getpid()}") + tmp.write_text(json.dumps(data, indent=2)) + os.replace(tmp, path) + + +def _read_spec(path: Path) -> TaskSpec: + data = json.loads(path.read_text()) + return TaskSpec( + **{k: v for k, v in data.items() if k in TaskSpec.__dataclass_fields__} + ) + + +def compute_fingerprint( + model_checkpoint_path: str, + setting: str | list, + patch_size: int, + patch_overlap: float, + batch_size: int, + apply_ctf: bool, + downsample: int, +) -> str: + """SHA-256 over the fields that require reloading the model/settings.""" + payload = json.dumps( + { + "model_checkpoint_path": model_checkpoint_path, + "setting": setting, + "patch_size": patch_size, + "patch_overlap": patch_overlap, + "batch_size": batch_size, + "apply_ctf": apply_ctf, + "downsample": downsample, + }, + sort_keys=True, + ).encode() + return hashlib.sha256(payload).hexdigest() + + +def write_pending(layout: QueueLayout, spec: TaskSpec) -> None: + _atomic_write(layout.pending / f"{spec.task_id}.json", asdict(spec)) + + +def claim_one(layout: QueueLayout, worker_id: str) -> TaskSpec | None: + """Attempt to claim a pending task via atomic rename. + + Returns the claimed TaskSpec, or None if the queue is empty. + """ + worker_dir = layout.worker_dir(worker_id) + worker_dir.mkdir(parents=True, exist_ok=True) + + candidates = list(layout.pending.glob("*.json")) + random.shuffle(candidates) + + for candidate in candidates: + dest = worker_dir / candidate.name + try: + os.rename(candidate, dest) + return _read_spec(dest) + except FileNotFoundError: + # another worker claimed it first + continue + + return None + + +def mark_done( + layout: QueueLayout, + worker_id: str, + spec: TaskSpec, + final_loss: float, + device: str, +) -> None: + """Write result to done/, then remove from running/ (publish-before-delete).""" + data = asdict(spec) + data["final_loss"] = final_loss + data["device"] = device + _atomic_write(layout.done / f"{spec.task_id}.json", data) + (layout.worker_dir(worker_id) / f"{spec.task_id}.json").unlink(missing_ok=True) + + +def mark_failed( + layout: QueueLayout, + worker_id: str, + spec: TaskSpec, + error: str, +) -> None: + """Write error to failed/, then remove from running/ (publish-before-delete).""" + data = asdict(spec) + data["error"] = error + data["worker_id"] = worker_id + _atomic_write(layout.failed / f"{spec.task_id}.json", data) + (layout.worker_dir(worker_id) / f"{spec.task_id}.json").unlink(missing_ok=True) + + +def clear_queue(layout: QueueLayout) -> None: + """Delete stale queue state from a prior run; recover running orphans to pending. + + Order matters: wipe pending/done/failed first, THEN recover orphans into + the now-empty pending/ so they are not immediately re-deleted. + """ + for directory in (layout.pending, layout.done, layout.failed): + for f in directory.glob("*.json"): + f.unlink(missing_ok=True) + + for worker_dir in layout.running.iterdir(): + if not worker_dir.is_dir(): + continue + for task_file in worker_dir.glob("*.json"): + dest = layout.pending / task_file.name + try: + os.rename(task_file, dest) + except FileNotFoundError: + pass + for hb in worker_dir.glob("hb-*"): + hb.unlink(missing_ok=True) + try: + worker_dir.rmdir() + except OSError: + pass # not empty yet; scheduler will sweep it diff --git a/src/miss_alignment/distributed/worker.py b/src/miss_alignment/distributed/worker.py new file mode 100644 index 00000000..f5ba87ae --- /dev/null +++ b/src/miss_alignment/distributed/worker.py @@ -0,0 +1,230 @@ +"""Worker subcommand: claims tasks from the queue and runs evaluate_tilt_series. + +Each worker processes many series per run. The model checkpoint is loaded once +when the first task is claimed, then reused for all subsequent tasks that share +the same init_fingerprint (all tasks in one alignment phase do). + +Usage (launched by provisioner): + miss-alignment worker --queue-dir --device [--worker-id ] +""" + +from __future__ import annotations + +import os +import sys +import threading +import time +import traceback +from pathlib import Path +from typing import Optional + +import torch +import typer + +from ..alignment.tilt_series import evaluate_tilt_series +from ..models.models import MissAlignment +from ..prepare_stacks import _prepare_single_tilt_series +from ..preprocessing import _run_cross_correlation_single +from .queue import ( + QueueLayout, + TaskSpec, + claim_one, + mark_done, + mark_failed, +) + +_MANAGER_HB_TIMEOUT_S = 120.0 +_HB_INTERVAL_S = 5.0 + + +def _write_worker_hb(worker_dir: Path, seq: int) -> None: + new_hb = worker_dir / f"hb-{seq}" + try: + new_hb.write_text("") + except FileNotFoundError: + # Worker dir was swept by the manager stall check; recreate it. + worker_dir.mkdir(parents=True, exist_ok=True) + new_hb.write_text("") + if seq > 0: + (worker_dir / f"hb-{seq - 1}").unlink(missing_ok=True) + + +def _start_heartbeat_thread( + worker_dir: Path, stop_event: threading.Event +) -> threading.Thread: + """Write a heartbeat tick every _HB_INTERVAL_S in a background daemon thread. + + Runs independently of task execution so long-running series don't make + the worker appear stale to the manager's sweep. + """ + + def _loop() -> None: + seq = 0 + while not stop_event.wait(timeout=_HB_INTERVAL_S): + _write_worker_hb(worker_dir, seq) + seq += 1 + + t = threading.Thread(target=_loop, daemon=True) + t.start() + return t + + +def _manager_hb_age_s(layout: QueueLayout) -> float: + """Seconds since the manager's most recent heartbeat tick, or infinity.""" + ticks = list(layout.manager_hb.glob("hb-*")) + if not ticks: + return float("inf") + latest = max(ticks, key=lambda p: p.stat().st_mtime) + return time.time() - latest.stat().st_mtime + + +def _load_model(checkpoint_path: str) -> MissAlignment: + model = MissAlignment.load_from_checkpoint(checkpoint_path, map_location="cpu") + # Unwrap torch.compile: incompatible with spawned processes (see tilt_series.py). + if hasattr(model.net, "_orig_mod"): + model.net = model.net._orig_mod + return model + + +def _device_int(device: str) -> int | None: + """Convert 'cuda:0' → 0, 'cpu' → None.""" + if device.startswith("cuda:"): + return int(device.split(":")[1]) + return None + + +def _execute_task( + spec: TaskSpec, + device: str, + cached_model: MissAlignment | None, +) -> float: + """Run the work described by spec and return a scalar result (final loss or 0.0).""" + if spec.task_type == "alignment": + # Convert setting back to tuple if serialised as list. + setting = ( + tuple(spec.setting) if isinstance(spec.setting, list) else spec.setting + ) + _, loss_values = evaluate_tilt_series( + model_checkpoint_path=Path(spec.model_checkpoint_path), + tilt_series_path=Path(spec.tilt_series_path), + output_directory=Path(spec.output_directory), + setting=setting, + patch_size=spec.patch_size, + patch_overlap=spec.patch_overlap, + batch_size=spec.batch_size, + apply_ctf=spec.apply_ctf, + downsample=spec.downsample, + device=device, + model=cached_model, + ) + return float(loss_values[-1]) if loss_values else float("nan") + + elif spec.task_type == "prepare_stacks": + _prepare_single_tilt_series( + xml_path=Path(spec.tilt_series_path), + desired_pixel_size=spec.desired_pixel_size, + device=_device_int(device), + ) + return 0.0 + + elif spec.task_type == "cross_correlation": + _run_cross_correlation_single( + xml_file=Path(spec.tilt_series_path), + device=_device_int(device), + lowpass_cutoff=spec.lowpass_cutoff or 0.25, + ) + return 0.0 + + else: + raise ValueError(f"Unknown task_type: {spec.task_type!r}") + + +def run_worker_loop( + layout: QueueLayout, + worker_id: str, + device: str, + manager_hb_timeout_s: float = _MANAGER_HB_TIMEOUT_S, +) -> str: + """Main worker loop: claim → evaluate → write result. Repeat until queue empty. + + Returns an exit-reason string suitable for writing to logs/.exit. + """ + worker_dir = layout.worker_dir(worker_id) + worker_dir.mkdir(parents=True, exist_ok=True) + + stop_hb = threading.Event() + _start_heartbeat_thread(worker_dir, stop_hb) + + last_fingerprint: str | None = None + cached_model: MissAlignment | None = None + tasks_done = 0 + tasks_failed = 0 + + try: + while True: + age = _manager_hb_age_s(layout) + if age > manager_hb_timeout_s: + return ( + f"manager heartbeat stale ({age:.0f}s > " + f"{manager_hb_timeout_s:.0f}s); " + f"done={tasks_done} failed={tasks_failed}" + ) + + spec = claim_one(layout, worker_id) + if spec is None: + return f"queue empty; done={tasks_done} failed={tasks_failed}" + + # For alignment tasks, (re)load model when fingerprint changes. + if ( + spec.task_type == "alignment" + and spec.init_fingerprint != last_fingerprint + ): + cached_model = _load_model(spec.model_checkpoint_path) + last_fingerprint = spec.init_fingerprint + + try: + final_loss = _execute_task(spec, device, cached_model) + mark_done(layout, worker_id, spec, final_loss=final_loss, device=device) + tasks_done += 1 + except Exception: + error = traceback.format_exc() + mark_failed(layout, worker_id, spec, error=error) + tasks_failed += 1 + print( + f"[{worker_id}] Failed {spec.task_id}:\n{error}", + file=sys.stderr, + ) + finally: + stop_hb.set() + + # Unreachable — loop only exits via return statements above. + return f"done={tasks_done} failed={tasks_failed}" + + +def worker_miss_align( + queue_dir: Path = typer.Option(..., help="Path to the tasks/ queue directory."), + device: int = typer.Option(0, help="GPU device index to use."), + worker_id: Optional[str] = typer.Option( + None, help="Unique worker ID. Defaults to local--gpu." + ), +) -> None: + """Claim and run inference tasks from the distributed queue.""" + if worker_id is None: + worker_id = f"local-{os.getpid()}-gpu{device}" + + layout = QueueLayout(queue_dir) + layout.ensure_directories() + + cuda_device = f"cuda:{device}" if torch.cuda.is_available() else "cpu" + exit_reason = "unknown (unhandled exception)" + try: + exit_reason = run_worker_loop(layout, worker_id, cuda_device) + except Exception: + exit_reason = f"unhandled exception:\n{traceback.format_exc()}" + raise + finally: + exit_file = layout.logs / f"{worker_id}.exit" + try: + exit_file.write_text(exit_reason) + except Exception: + pass # don't mask the original error diff --git a/src/miss_alignment/infer.py b/src/miss_alignment/infer.py index d884bd7a..d18cd91a 100644 --- a/src/miss_alignment/infer.py +++ b/src/miss_alignment/infer.py @@ -39,6 +39,13 @@ def infer_miss_align( help="Run cross-correlation based alignment before the inference " "iterations. This performs coarse alignment.", ), + n_cluster_workers: Optional[int] = typer.Option( + None, + help="Number of cluster jobs to submit for the alignment phase. " + "When set, activates cluster mode; requires MISS_CLUSTER_CONFIG " + "and MISS_CLUSTER_SCRIPT environment variables to be set. " + "When absent, local multi-GPU mode is used.", + ), ) -> None: """Align a dataset by applying models from a previous training run. @@ -93,6 +100,7 @@ def infer_miss_align( training_directory=data_directory, desired_pixel_size=prepare_stacks, devices=devices_alignment, + n_cluster_workers=n_cluster_workers, ) # Run preprocessing if requested @@ -112,6 +120,7 @@ def infer_miss_align( run_cross_correlation_alignment_parallel( training_directory=data_directory, devices=devices_alignment, + n_cluster_workers=n_cluster_workers, ) start_iter = start_at_iteration @@ -159,6 +168,7 @@ def infer_miss_align( apply_ctf=general_config["apply_ctf"], downsample=iteration_settings["downsample"], devices_list=devices_alignment, + n_cluster_workers=n_cluster_workers, ) # make copies of the xml files after alignment diff --git a/src/miss_alignment/prepare_stacks.py b/src/miss_alignment/prepare_stacks.py index 8bc0b6a7..f4ce53db 100644 --- a/src/miss_alignment/prepare_stacks.py +++ b/src/miss_alignment/prepare_stacks.py @@ -4,14 +4,13 @@ preprocessed tilt stacks ready for training. """ -import queue from pathlib import Path import mrcfile from warpylib import TiltSeries from warpylib.movie import Movie -from ._parallel import run_device_pool +from .distributed.manager import run_distributed def _get_original_pixel_size(tilt_series: TiltSeries) -> float: @@ -89,7 +88,6 @@ def _prepare_single_tilt_series( ts = TiltSeries(xml_path) original_pixel_size = _get_original_pixel_size(ts) - # print(f"{xml_path.stem}: original pixel size = {original_pixel_size:.4f} Å") images, _, _ = ts.load_images( original_pixel_size=original_pixel_size, @@ -106,29 +104,11 @@ def _prepare_single_tilt_series( ) -def _prepare_stacks_runner( - device: int | None, - task_queue, - result_queue, - desired_pixel_size: float, -) -> None: - """Pull tilt-series off the queue and prepare them on a single device.""" - import torch - - torch.set_num_threads(1) - while True: - try: - xml_path = task_queue.get_nowait() - except queue.Empty: - break - _prepare_single_tilt_series(xml_path, desired_pixel_size, device) - result_queue.put_nowait(xml_path.stem) - - def prepare_stacks_parallel( training_directory: Path, desired_pixel_size: float, devices: list[int] | None = None, + n_cluster_workers: int | None = None, ) -> None: """Prepare tilt stacks for all tilt series in the training directory. @@ -144,13 +124,16 @@ def prepare_stacks_parallel( devices : list[int] | None CUDA device indices to distribute work across (one worker process per unique device). If None, a single default-device worker is used. + n_cluster_workers : int | None + Number of cluster jobs to submit. When set, activates cluster mode; + requires MISS_CLUSTER_CONFIG and MISS_CLUSTER_SCRIPT to be set. Raises ------ FileNotFoundError If no XML files are found in the training directory. RuntimeError - If any tilt series fails to process (terminates on first error). + If any tilt series fails to process. """ xml_files = list(training_directory.glob("*.xml")) if not xml_files: @@ -159,15 +142,26 @@ def prepare_stacks_parallel( ) print( - f"Preparing stacks for {len(xml_files)} tilt series at {desired_pixel_size} Å" + f"Preparing stacks for {len(xml_files)} tilt series " + f"at {desired_pixel_size} Å" ) - run_device_pool( - jobs=xml_files, - runner=_prepare_stacks_runner, - runner_args=(desired_pixel_size,), - devices=devices, - desc="Preparing stacks", + queue_root = training_directory / "tasks" + run_distributed( + tilt_series_list=xml_files, + model_checkpoint=Path(""), + output_directory=training_directory, + setting="", + patch_size=0, + patch_overlap=0.0, + batch_size=0, + apply_ctf=False, + downsample=1, + devices=devices or [], + n_cluster_workers=n_cluster_workers, + queue_root=queue_root, + task_type="prepare_stacks", + desired_pixel_size=desired_pixel_size, ) print(f"Successfully prepared stacks for {len(xml_files)} tilt series") diff --git a/src/miss_alignment/preprocessing.py b/src/miss_alignment/preprocessing.py index 4a59a936..156648c2 100644 --- a/src/miss_alignment/preprocessing.py +++ b/src/miss_alignment/preprocessing.py @@ -1,11 +1,10 @@ """Preprocessing utilities for tilt-series alignment.""" -import queue import torch from pathlib import Path -from ._parallel import run_device_pool from .data.io import TiltSeriesData +from .distributed.manager import run_distributed def _run_cross_correlation_single( @@ -60,29 +59,11 @@ def _run_cross_correlation_single( ts_data.save_metadata_to_xml(ts) -def _cross_correlation_runner( - device: int | None, task_queue, result_queue, lowpass_cutoff: float -) -> None: - """Pull tilt-series off the queue and align them on a single device.""" - - torch.set_num_threads(1) - while True: - try: - xml_file = task_queue.get_nowait() - except queue.Empty: - break - _run_cross_correlation_single( - xml_file, - device, - lowpass_cutoff, - ) - result_queue.put_nowait(xml_file.stem) - - def run_cross_correlation_alignment_parallel( training_directory: Path, devices: list[int] | None = None, lowpass_cutoff: float = 0.25, + n_cluster_workers: int | None = None, ) -> None: """ Run cross-correlation based alignment in parallel. @@ -100,6 +81,9 @@ def run_cross_correlation_alignment_parallel( unique device). If None, a single default-device worker is used. lowpass_cutoff : float, optional Low-pass filter cutoff frequency (default: 0.25). + n_cluster_workers : int | None + Number of cluster jobs to submit. When set, activates cluster mode; + requires MISS_CLUSTER_CONFIG and MISS_CLUSTER_SCRIPT to be set. """ # Get list of all XML files to process xml_files = list(training_directory.glob("*.xml")) @@ -114,12 +98,22 @@ def run_cross_correlation_alignment_parallel( else: print(" Using default device assignment\n") - run_device_pool( - jobs=xml_files, - runner=_cross_correlation_runner, - runner_args=(lowpass_cutoff,), - devices=devices, - desc="Cross-correlation alignment", + queue_root = training_directory / "tasks" + run_distributed( + tilt_series_list=xml_files, + model_checkpoint=Path(""), + output_directory=training_directory, + setting="", + patch_size=0, + patch_overlap=0.0, + batch_size=0, + apply_ctf=False, + downsample=1, + devices=devices or [], + n_cluster_workers=n_cluster_workers, + queue_root=queue_root, + task_type="cross_correlation", + lowpass_cutoff=lowpass_cutoff, ) print("\nCross-correlation alignment complete!\n") diff --git a/src/miss_alignment/train.py b/src/miss_alignment/train.py index 08538c1a..8cc06c8e 100644 --- a/src/miss_alignment/train.py +++ b/src/miss_alignment/train.py @@ -310,6 +310,13 @@ def train_miss_align( help="Run cross-correlation based alignment before training iterations. " "This performs coarse alignment.", ), + n_cluster_workers: Optional[int] = typer.Option( + None, + help="Number of cluster jobs to submit for the alignment phase. " + "When set, activates cluster mode; requires MISS_CLUSTER_CONFIG " + "and MISS_CLUSTER_SCRIPT environment variables to be set. " + "When absent, local multi-GPU mode is used.", + ), ) -> None: """Iteratively train and realign a dataset over a series of coarse-to-fine macro-iterations. @@ -387,6 +394,7 @@ def train_miss_align( training_directory=training_directory, desired_pixel_size=prepare_stacks, devices=devices_alignment, + n_cluster_workers=n_cluster_workers, ) # Run preprocessing if requested @@ -410,6 +418,7 @@ def train_miss_align( run_cross_correlation_alignment_parallel( training_directory=training_directory, devices=devices_alignment, + n_cluster_workers=n_cluster_workers, ) start_iter = start_at_iteration @@ -485,6 +494,7 @@ def train_miss_align( apply_ctf=general_config["apply_ctf"], downsample=iteration_settings["downsample"], devices_list=devices_alignment, + n_cluster_workers=n_cluster_workers, ) # make copies of the xml files and model after alignment diff --git a/tests/distributed/__init__.py b/tests/distributed/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/distributed/test_config.py b/tests/distributed/test_config.py new file mode 100644 index 00000000..1fa005aa --- /dev/null +++ b/tests/distributed/test_config.py @@ -0,0 +1,82 @@ +import json +import pytest +from pathlib import Path +from miss_alignment.distributed.config import ClusterConfig, load_cluster_config + + +@pytest.fixture() +def cluster_json(tmp_path): + cfg = { + "submit": "sbatch {{script_path}}", + "submit_job_id_regex": r"Submitted batch job (\d+)", + "cancel": "scancel {{job_id}}", + "status_list": "squeue -u $USER -h -o '%i,%T'", + } + p = tmp_path / "cluster.json" + p.write_text(json.dumps(cfg)) + return p + + +@pytest.fixture() +def cluster_script(tmp_path): + p = tmp_path / "worker.sh" + p.write_text("#!/bin/bash\n{{command}}\n") + return p + + +def test_load_cluster_config_raises_when_config_unset(monkeypatch): + monkeypatch.delenv("MISS_CLUSTER_CONFIG", raising=False) + monkeypatch.delenv("MISS_CLUSTER_SCRIPT", raising=False) + with pytest.raises(RuntimeError, match="MISS_CLUSTER_CONFIG"): + load_cluster_config() + + +def test_load_cluster_config_raises_when_script_unset(monkeypatch, cluster_json): + monkeypatch.setenv("MISS_CLUSTER_CONFIG", str(cluster_json)) + monkeypatch.delenv("MISS_CLUSTER_SCRIPT", raising=False) + with pytest.raises(RuntimeError, match="MISS_CLUSTER_SCRIPT"): + load_cluster_config() + + +def test_load_cluster_config_returns_config(monkeypatch, cluster_json, cluster_script): + monkeypatch.setenv("MISS_CLUSTER_CONFIG", str(cluster_json)) + monkeypatch.setenv("MISS_CLUSTER_SCRIPT", str(cluster_script)) + cfg = load_cluster_config() + assert isinstance(cfg, ClusterConfig) + assert "sbatch" in cfg.submit + assert cfg.script_path == cluster_script + assert r"(\d+)" in cfg.submit_job_id_regex + assert "squeue" in cfg.status_list + assert cfg.scheduler == "auto" + + +def test_load_cluster_config_explicit_scheduler(monkeypatch, tmp_path, cluster_script): + cfg_data = { + "submit": "sbatch {{script_path}}", + "submit_job_id_regex": r"Submitted batch job (\d+)", + "cancel": "scancel {{job_id}}", + "status_list": "squeue -u $USER -h -o '%i,%T'", + "scheduler": "slurm", + } + p = tmp_path / "cluster.json" + p.write_text(json.dumps(cfg_data)) + monkeypatch.setenv("MISS_CLUSTER_CONFIG", str(p)) + monkeypatch.setenv("MISS_CLUSTER_SCRIPT", str(cluster_script)) + cfg = load_cluster_config() + assert cfg.scheduler == "slurm" + + +def test_load_cluster_config_raises_on_missing_file(monkeypatch, tmp_path, cluster_script): + monkeypatch.setenv("MISS_CLUSTER_CONFIG", str(tmp_path / "nonexistent.json")) + monkeypatch.setenv("MISS_CLUSTER_SCRIPT", str(cluster_script)) + with pytest.raises(FileNotFoundError): + load_cluster_config() + + +def test_load_cluster_config_raises_on_missing_key(monkeypatch, tmp_path, cluster_script): + bad = tmp_path / "bad.json" + bad.write_text('{"submit": "sbatch {{script_path}}"}') + monkeypatch.setenv("MISS_CLUSTER_CONFIG", str(bad)) + monkeypatch.setenv("MISS_CLUSTER_SCRIPT", str(cluster_script)) + with pytest.raises(KeyError): + load_cluster_config() diff --git a/tests/distributed/test_manager.py b/tests/distributed/test_manager.py new file mode 100644 index 00000000..50a70c0c --- /dev/null +++ b/tests/distributed/test_manager.py @@ -0,0 +1,214 @@ +"""Integration tests for the manager coordinator. + +A fake worker thread simulates cluster workers by polling pending/ and +writing done/ or failed/ files. +""" +import json +import os +import threading +import time +from pathlib import Path +from unittest.mock import patch + +import pytest + +from miss_alignment.distributed.manager import run_distributed +from miss_alignment.distributed.queue import QueueLayout + + +def _make_xml(tmp_path, name): + p = tmp_path / f"{name}.xml" + p.write_text(f"{name}") + return p + + +def _fake_worker_thread(queue_root, n_tasks, fail=False): + """Simulates a worker: claims pending tasks, writes done or failed.""" + + def _run(): + layout = QueueLayout(queue_root) + done_count = 0 + deadline = time.time() + 15 + while done_count < n_tasks and time.time() < deadline: + for f in list(layout.pending.glob("*.json")): + data = json.loads(f.read_text()) + task_id = data["task_id"] + running_dir = layout.running / "fake-worker" + running_dir.mkdir(parents=True, exist_ok=True) + try: + os.rename(f, running_dir / f.name) + except FileNotFoundError: + continue + if fail: + fail_data = {**data, "error": "boom", "worker_id": "fake-worker"} + (layout.failed / f"{task_id}.json").write_text( + json.dumps(fail_data) + ) + else: + done_data = {**data, "final_loss": 0.01, "device": "cpu"} + (layout.done / f"{task_id}.json").write_text( + json.dumps(done_data) + ) + (running_dir / f"{task_id}.json").unlink(missing_ok=True) + done_count += 1 + time.sleep(0.05) + + return threading.Thread(target=_run, daemon=True) + + +class _NoOpProvisioner: + def ensure_workers(self, n_workers): + pass + + def shutdown(self): + pass + + def live_worker_count(self): + return 0 + + def worker_counts_by_type(self): + return {"workers": 0} + + +def test_run_distributed_returns_losses(tmp_path): + xml1 = _make_xml(tmp_path, "ts01") + xml2 = _make_xml(tmp_path, "ts02") + ckpt = tmp_path / "model.ckpt" + ckpt.write_text("") + queue_root = tmp_path / "tasks" + + worker = _fake_worker_thread(queue_root, n_tasks=2) + worker.start() + + with patch( + "miss_alignment.distributed.manager.LocalProvisioner", + return_value=_NoOpProvisioner(), + ): + losses = run_distributed( + tilt_series_list=[xml1, xml2], + model_checkpoint=ckpt, + output_directory=tmp_path, + setting="anchoring", + patch_size=96, + patch_overlap=0.1, + batch_size=32, + apply_ctf=False, + downsample=2, + devices=[0], + n_cluster_workers=None, + queue_root=queue_root, + ) + + assert set(losses.keys()) == {"ts01", "ts02"} + assert all(v == pytest.approx(0.01) for v in losses.values()) + + +def test_run_distributed_raises_on_any_failure(tmp_path): + xml1 = _make_xml(tmp_path, "ts01") + ckpt = tmp_path / "model.ckpt" + ckpt.write_text("") + queue_root = tmp_path / "tasks" + + worker = _fake_worker_thread(queue_root, n_tasks=1, fail=True) + worker.start() + + with patch( + "miss_alignment.distributed.manager.LocalProvisioner", + return_value=_NoOpProvisioner(), + ): + with pytest.raises(RuntimeError, match="ts01"): + run_distributed( + tilt_series_list=[xml1], + model_checkpoint=ckpt, + output_directory=tmp_path, + setting="anchoring", + patch_size=96, + patch_overlap=0.1, + batch_size=32, + apply_ctf=False, + downsample=2, + devices=[0], + n_cluster_workers=None, + queue_root=queue_root, + ) + + +def test_run_distributed_cleans_up_tasks_dir(tmp_path): + xml1 = _make_xml(tmp_path, "ts01") + ckpt = tmp_path / "model.ckpt" + ckpt.write_text("") + queue_root = tmp_path / "tasks" + + worker = _fake_worker_thread(queue_root, n_tasks=1) + worker.start() + + with patch( + "miss_alignment.distributed.manager.LocalProvisioner", + return_value=_NoOpProvisioner(), + ): + run_distributed( + tilt_series_list=[xml1], + model_checkpoint=ckpt, + output_directory=tmp_path, + setting="anchoring", + patch_size=96, + patch_overlap=0.1, + batch_size=32, + apply_ctf=False, + downsample=2, + devices=[0], + n_cluster_workers=None, + queue_root=queue_root, + ) + + assert not queue_root.exists() + + +def test_run_distributed_prepare_stacks_task_type(tmp_path): + """Manager creates prepare_stacks tasks with the correct task_type field.""" + import json + + xml1 = _make_xml(tmp_path, "ts01") + queue_root = tmp_path / "tasks" + + # Capture the task JSON before the worker thread consumes it + written_tasks = [] + + worker = _fake_worker_thread(queue_root, n_tasks=1) + + with patch( + "miss_alignment.distributed.manager.LocalProvisioner", + return_value=_NoOpProvisioner(), + ): + # Run in a thread so we can inspect tasks/ briefly — but just check + # the return value here since the fake worker resolves immediately. + losses = None + + def run(): + nonlocal losses + losses = run_distributed( + tilt_series_list=[xml1], + model_checkpoint=Path(""), + output_directory=tmp_path, + setting="", + patch_size=0, + patch_overlap=0.0, + batch_size=0, + apply_ctf=False, + downsample=1, + devices=[0], + n_cluster_workers=None, + queue_root=queue_root, + task_type="prepare_stacks", + desired_pixel_size=10.0, + ) + + import threading + + worker.start() + t = threading.Thread(target=run) + t.start() + t.join(timeout=15) + + assert losses is not None + assert "ts01" in losses diff --git a/tests/distributed/test_provisioner.py b/tests/distributed/test_provisioner.py new file mode 100644 index 00000000..c0eddbca --- /dev/null +++ b/tests/distributed/test_provisioner.py @@ -0,0 +1,340 @@ +import sys +import time +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from miss_alignment.distributed.config import ClusterConfig +from miss_alignment.distributed.provisioner import ( + ClusterProvisioner, + CompositeProvisioner, + LocalProvisioner, + _parse_status_output, +) + + +def _cluster_cfg(tmp_path) -> ClusterConfig: + script = tmp_path / "worker.sh" + script.write_text("#!/bin/bash\n{{command}}\n") + return ClusterConfig( + submit="sbatch {{script_path}}", + submit_job_id_regex=r"Submitted batch job (\d+)", + cancel="scancel {{job_id}}", + status_list="squeue -u $USER -h -o '%i,%T'", + script_path=script, + ) + + +# --------------------------------------------------------------------------- +# Status parser unit tests +# --------------------------------------------------------------------------- + +def test_parse_status_output_slurm_running_and_pending(): + output = "12345,PENDING\n12346,RUNNING\n12347,COMPLETING\n" + our_ids = {"12345", "12346", "12347", "99999"} + states = _parse_status_output(output, our_ids, "slurm", [], "") + assert states == {"12345": "pending", "12346": "running", "12347": "running"} + + +def test_parse_status_output_slurm_terminal_absent(): + """Terminal jobs (COMPLETED/FAILED) are absent from the result, not included.""" + output = "12345,COMPLETED\n12346,FAILED\n12347,CANCELLED\n" + our_ids = {"12345", "12346", "12347"} + states = _parse_status_output(output, our_ids, "slurm", [], "") + assert states == {} + + +def test_parse_status_output_lsf(): + output = "12345,PEND\n12346,RUN\n" + our_ids = {"12345", "12346"} + states = _parse_status_output(output, our_ids, "lsf", [], "") + assert states == {"12345": "pending", "12346": "running"} + + +def test_parse_status_output_pbs(): + output = "12345,Q\n12346,R\n12347,C\n" + our_ids = {"12345", "12346", "12347"} + states = _parse_status_output(output, our_ids, "pbs", [], "") + assert states == {"12345": "pending", "12346": "running"} + + +def test_parse_status_output_sge(): + output = "12345,qw\n12346,r\n" + our_ids = {"12345", "12346"} + states = _parse_status_output(output, our_ids, "sge", [], "") + assert states == {"12345": "pending", "12346": "running"} + + +def test_parse_status_output_auto_detects_slurm(): + output = "12345,PENDING\n12346,RUNNING\n" + our_ids = {"12345", "12346"} + states = _parse_status_output(output, our_ids, "auto", [], "") + assert states == {"12345": "pending", "12346": "running"} + + +def test_parse_status_output_ignores_unknown_ids(): + output = "99999,RUNNING\n" + our_ids = {"12345"} + states = _parse_status_output(output, our_ids, "slurm", [], "") + assert states == {} + + +def test_parse_status_output_no_status_token_treated_as_pending(): + """A job present in output with no comma is assumed pending.""" + output = "12345\n" + our_ids = {"12345"} + states = _parse_status_output(output, our_ids, "slurm", [], "") + assert states == {"12345": "pending"} + + +def test_parse_status_output_custom(): + """Custom: alive statuses appear as 'pending'; others are absent.""" + output = "12345,INPROGRESS\n12346,DONE\n" + our_ids = {"12345", "12346"} + states = _parse_status_output(output, our_ids, "custom", ["INPROGRESS"], "") + assert states == {"12345": "pending"} + + +# --------------------------------------------------------------------------- +# LocalProvisioner tests +# --------------------------------------------------------------------------- + +def test_local_provisioner_spawns_one_process_per_device(tmp_path): + with patch("miss_alignment.distributed.provisioner.subprocess.Popen") as mock_popen: + mock_proc = MagicMock() + mock_proc.poll.return_value = None + mock_popen.return_value = mock_proc + + p = LocalProvisioner(queue_dir=tmp_path, devices=[0, 1, 2]) + p.ensure_workers(n_workers=10) + + assert mock_popen.call_count == 3 + all_args = [str(c) for c in mock_popen.call_args_list] + assert any("0" in s for s in all_args) + + +def test_local_provisioner_does_not_respawn_running(tmp_path): + with patch("miss_alignment.distributed.provisioner.subprocess.Popen") as mock_popen: + mock_proc = MagicMock() + mock_proc.poll.return_value = None + mock_popen.return_value = mock_proc + + p = LocalProvisioner(queue_dir=tmp_path, devices=[0]) + p.ensure_workers(n_workers=5) + p.ensure_workers(n_workers=5) + + assert mock_popen.call_count == 1 + + +def test_local_provisioner_respawns_dead_worker(tmp_path): + with patch("miss_alignment.distributed.provisioner.subprocess.Popen") as mock_popen: + mock_proc = MagicMock() + mock_proc.poll.return_value = 0 # exited + mock_popen.return_value = mock_proc + + p = LocalProvisioner(queue_dir=tmp_path, devices=[0]) + p.ensure_workers(n_workers=5) + p.ensure_workers(n_workers=5) + + assert mock_popen.call_count == 2 + + +def test_local_provisioner_shutdown_terminates(tmp_path): + with patch("miss_alignment.distributed.provisioner.subprocess.Popen") as mock_popen: + mock_proc = MagicMock() + mock_proc.poll.return_value = None + mock_popen.return_value = mock_proc + + p = LocalProvisioner(queue_dir=tmp_path, devices=[0]) + p.ensure_workers(n_workers=5) + p.shutdown() + + mock_proc.terminate.assert_called() + + +# --------------------------------------------------------------------------- +# ClusterProvisioner tests +# --------------------------------------------------------------------------- + +def test_cluster_provisioner_submits_n_workers_jobs(tmp_path): + """ensure_workers submits until _job_submit_time reaches target.""" + cfg = _cluster_cfg(tmp_path) + call_num = [0] + + def fake_run(cmd, **kwargs): + call_num[0] += 1 + result = MagicMock() + result.stdout = ( + "" if "squeue" in cmd + else f"Submitted batch job {call_num[0]}\n" + ) + return result + + with patch( + "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run + ): + p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) + p.ensure_workers(n_workers=4) + + assert len(p._job_submit_time) == 4 + + +def test_cluster_provisioner_does_not_resubmit_within_grace(tmp_path): + """Jobs recently submitted but absent from squeue are not double-counted.""" + cfg = _cluster_cfg(tmp_path) + + def fake_run(cmd, **kwargs): + result = MagicMock() + result.stdout = "" # squeue returns nothing (jobs not yet registered) + return result + + with patch( + "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run + ): + p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) + now = time.time() + p._job_submit_time = {"1": now, "2": now, "3": now, "4": now} + p.ensure_workers(n_workers=4) + + assert len(p._job_submit_time) == 4 # no new submissions + + +def test_cluster_provisioner_prunes_grace_expired_absent_jobs(tmp_path): + """Jobs absent from squeue past the grace period are pruned.""" + cfg = _cluster_cfg(tmp_path) + + def fake_run(cmd, **kwargs): + result = MagicMock() + result.stdout = "3,RUNNING\n" # only job 3 visible + return result + + with patch( + "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run + ): + p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) + old = time.time() - 200 # well past grace period + now = time.time() + p._job_submit_time = {"1": old, "2": old, "3": now} + states = p._query_job_states() + + # Old absent jobs pruned; recent absent job kept; visible job in states + assert "1" not in p._job_submit_time + assert "2" not in p._job_submit_time + assert "3" in p._job_submit_time + assert states == {"3": "running"} + + +def test_cluster_provisioner_absent_within_grace_not_pruned(tmp_path): + """Jobs absent from squeue within the grace period are kept.""" + cfg = _cluster_cfg(tmp_path) + + def fake_run(cmd, **kwargs): + result = MagicMock() + result.stdout = "11111,RUNNING\n" + return result + + with patch( + "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run + ): + p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) + now = time.time() + p._job_submit_time = {"11111": now, "22222": now, "33333": now} + states = p._query_job_states() + + assert states == {"11111": "running"} + assert set(p._job_submit_time.keys()) == {"11111", "22222", "33333"} + + +def test_cluster_provisioner_replenishes_grace_expired_jobs(tmp_path): + """ensure_workers resubmits for jobs that have aged out.""" + cfg = _cluster_cfg(tmp_path) + submitted = [] + + def fake_run(cmd, **kwargs): + result = MagicMock() + if "squeue" in cmd: + result.stdout = "3,RUNNING\n4,RUNNING\n" + else: + submitted.append(cmd) + result.stdout = f"Submitted batch job {len(submitted) + 10}\n" + return result + + with patch( + "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run + ): + p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) + old = time.time() - 200 + now = time.time() + p._job_submit_time = {"1": old, "2": old, "3": now, "4": now} + p.ensure_workers(n_workers=4) + + assert len(submitted) == 2 # replaced the 2 expired-absent jobs + + +def test_cluster_provisioner_worker_counts_split_running_pending(tmp_path): + cfg = _cluster_cfg(tmp_path) + + def fake_run(cmd, **kwargs): + result = MagicMock() + result.stdout = "1,RUNNING\n2,RUNNING\n3,PENDING\n" + return result + + with patch( + "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run + ): + p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) + now = time.time() + # 5 tracked: 3 visible (2 running, 1 pending), 2 not yet visible + p._job_submit_time = {"1": now, "2": now, "3": now, "4": now, "5": now} + counts = p.worker_counts_by_type() + + assert counts["cluster-running"] == 2 + assert counts["cluster-pending"] == 3 # 1 explicit + 2 not-yet-visible + + +def test_cluster_provisioner_cancels_on_shutdown(tmp_path): + cfg = _cluster_cfg(tmp_path) + cancel_calls = [] + submit_count = [0] + + def fake_run(cmd, **kwargs): + if "sbatch" in cmd: + submit_count[0] += 1 + result = MagicMock() + result.stdout = f"Submitted batch job {submit_count[0]}\n" + return result + if "squeue" in cmd: + result = MagicMock() + result.stdout = "" + return result + cancel_calls.append(cmd) + return MagicMock() + + with patch( + "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run + ): + p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) + p.ensure_workers(n_workers=3) + p.shutdown() + + assert len(cancel_calls) == 3 + assert all("scancel" in c for c in cancel_calls) + + +# --------------------------------------------------------------------------- +# CompositeProvisioner tests +# --------------------------------------------------------------------------- + +def test_composite_provisioner_delegates_to_all(tmp_path): + a = MagicMock(spec=LocalProvisioner) + b = MagicMock(spec=LocalProvisioner) + p = CompositeProvisioner([a, b]) + + p.ensure_workers(5) + a.ensure_workers.assert_called_once_with(5) + b.ensure_workers.assert_called_once_with(5) + + p.shutdown() + a.shutdown.assert_called_once() + b.shutdown.assert_called_once() diff --git a/tests/distributed/test_queue.py b/tests/distributed/test_queue.py new file mode 100644 index 00000000..e86451e6 --- /dev/null +++ b/tests/distributed/test_queue.py @@ -0,0 +1,127 @@ +import json +import os +from pathlib import Path +import pytest +from miss_alignment.distributed.queue import ( + QueueLayout, + TaskSpec, + clear_queue, + claim_one, + compute_fingerprint, + mark_done, + mark_failed, + write_pending, +) + + +@pytest.fixture() +def layout(tmp_path): + layout = QueueLayout(tmp_path / "tasks") + layout.ensure_directories() + return layout + + +def _spec(task_id="0000001-ts01"): + return TaskSpec( + task_id=task_id, + model_checkpoint_path="/data/model.ckpt", + tilt_series_path="/data/ts01.xml", + output_directory="/data/out", + setting="anchoring", + patch_size=96, + patch_overlap=0.1, + batch_size=32, + apply_ctf=False, + downsample=2, + init_fingerprint="abc123", + ) + + +def test_write_pending_creates_json(layout): + write_pending(layout, _spec()) + assert (layout.pending / "0000001-ts01.json").exists() + + +def test_claim_one_returns_spec_and_moves_file(layout): + write_pending(layout, _spec()) + result = claim_one(layout, "worker-0") + assert result is not None + assert result.task_id == "0000001-ts01" + assert not (layout.pending / "0000001-ts01.json").exists() + assert (layout.running / "worker-0" / "0000001-ts01.json").exists() + + +def test_claim_one_returns_none_when_empty(layout): + assert claim_one(layout, "worker-0") is None + + +def test_claim_one_exclusive(layout): + """Two sequential claimers: exactly one wins.""" + write_pending(layout, _spec()) + r0 = claim_one(layout, "worker-0") + r1 = claim_one(layout, "worker-1") + claimed = [r for r in (r0, r1) if r is not None] + assert len(claimed) == 1 + + +def test_mark_done_writes_done_and_removes_running(layout): + spec = _spec() + write_pending(layout, spec) + claim_one(layout, "worker-0") + mark_done(layout, "worker-0", spec, final_loss=0.042, device="cuda:0") + done_path = layout.done / "0000001-ts01.json" + assert done_path.exists() + data = json.loads(done_path.read_text()) + assert data["final_loss"] == pytest.approx(0.042) + assert data["device"] == "cuda:0" + assert not (layout.running / "worker-0" / "0000001-ts01.json").exists() + + +def test_mark_failed_writes_failed_and_removes_running(layout): + spec = _spec() + write_pending(layout, spec) + claim_one(layout, "worker-0") + mark_failed(layout, "worker-0", spec, error="CUDA OOM") + failed_path = layout.failed / "0000001-ts01.json" + assert failed_path.exists() + data = json.loads(failed_path.read_text()) + assert data["error"] == "CUDA OOM" + assert not (layout.running / "worker-0" / "0000001-ts01.json").exists() + + +def test_clear_queue_recovers_orphans(layout): + spec = _spec() + write_pending(layout, spec) + claim_one(layout, "worker-0") + # simulate crash: running file remains; clear should put it back in pending + clear_queue(layout) + assert (layout.pending / "0000001-ts01.json").exists() + + +def test_clear_queue_wipes_done_and_failed(layout): + spec = _spec() + write_pending(layout, spec) + claim_one(layout, "worker-0") + mark_done(layout, "worker-0", spec, final_loss=0.1, device="cpu") + # write a second spec directly to failed + spec2 = _spec("0000002-ts02") + write_pending(layout, spec2) + claim_one(layout, "worker-0") + mark_failed(layout, "worker-0", spec2, error="boom") + + clear_queue(layout) + assert list(layout.done.glob("*.json")) == [] + assert list(layout.failed.glob("*.json")) == [] + + +def test_compute_fingerprint_is_deterministic(): + fp1 = compute_fingerprint("/ckpt", "anchoring", 96, 0.1, 32, False, 2) + fp2 = compute_fingerprint("/ckpt", "anchoring", 96, 0.1, 32, False, 2) + assert fp1 == fp2 + assert len(fp1) == 64 # SHA-256 hex + + +def test_compute_fingerprint_differs_on_change(): + fp1 = compute_fingerprint("/ckpt", "anchoring", 96, 0.1, 32, False, 2) + fp2 = compute_fingerprint("/other.ckpt", "anchoring", 96, 0.1, 32, False, 2) + assert fp1 != fp2 diff --git a/tests/distributed/test_worker.py b/tests/distributed/test_worker.py new file mode 100644 index 00000000..5d95a879 --- /dev/null +++ b/tests/distributed/test_worker.py @@ -0,0 +1,247 @@ +"""Unit tests for the worker claim loop. + +evaluate_tilt_series is mocked throughout; these tests validate the claim, +model-reuse, heartbeat-exit, and result-write logic without CUDA. +""" + +import json +import os +import time +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from miss_alignment.distributed.queue import QueueLayout, TaskSpec, write_pending +from miss_alignment.distributed.worker import run_worker_loop + + +@pytest.fixture() +def layout(tmp_path): + layout = QueueLayout(tmp_path / "tasks") + layout.ensure_directories() + return layout + + +def _write_manager_hb(layout, seq=0): + for old in layout.manager_hb.glob("hb-*"): + old.unlink(missing_ok=True) + (layout.manager_hb / f"hb-{seq}").write_text("") + + +def _spec(task_id="0000001-ts01", fingerprint="abc123"): + return TaskSpec( + task_id=task_id, + model_checkpoint_path="/data/model.ckpt", + tilt_series_path=f"/data/{task_id}.xml", + output_directory="/data/out", + setting="anchoring", + patch_size=96, + patch_overlap=0.1, + batch_size=32, + apply_ctf=False, + downsample=2, + init_fingerprint=fingerprint, + ) + + +def test_worker_processes_task_and_writes_done(layout): + _write_manager_hb(layout) + write_pending(layout, _spec()) + + with ( + patch( + "miss_alignment.distributed.worker.evaluate_tilt_series", + return_value=(Path("/data/ts01.xml"), [0.5, 0.3, 0.1]), + ), + patch( + "miss_alignment.distributed.worker.MissAlignment.load_from_checkpoint", + return_value=MagicMock(), + ), + ): + run_worker_loop(layout, "worker-0", "cpu", manager_hb_timeout_s=30.0) + + done = layout.done / "0000001-ts01.json" + assert done.exists() + data = json.loads(done.read_text()) + assert data["final_loss"] == pytest.approx(0.1) + + +def test_worker_writes_failed_on_exception(layout): + _write_manager_hb(layout) + write_pending(layout, _spec()) + + with ( + patch( + "miss_alignment.distributed.worker.evaluate_tilt_series", + side_effect=RuntimeError("CUDA OOM"), + ), + patch( + "miss_alignment.distributed.worker.MissAlignment.load_from_checkpoint", + return_value=MagicMock(), + ), + ): + run_worker_loop(layout, "worker-0", "cpu", manager_hb_timeout_s=30.0) + + failed = layout.failed / "0000001-ts01.json" + assert failed.exists() + data = json.loads(failed.read_text()) + assert "CUDA OOM" in data["error"] + + +def test_worker_exits_when_manager_hb_stale(layout): + # Write a heartbeat file that is 200 seconds old + hb_file = layout.manager_hb / "hb-0" + hb_file.write_text("") + old_time = time.time() - 200 + os.utime(hb_file, (old_time, old_time)) + + write_pending(layout, _spec()) + + called = [] + with ( + patch( + "miss_alignment.distributed.worker.evaluate_tilt_series", + side_effect=lambda **kw: called.append(True), + ), + patch( + "miss_alignment.distributed.worker.MissAlignment.load_from_checkpoint", + return_value=MagicMock(), + ), + ): + run_worker_loop(layout, "worker-0", "cpu", manager_hb_timeout_s=120.0) + + assert called == [] + + +def test_worker_reuses_model_when_fingerprint_matches(layout): + """Model is loaded once when two tasks share the same init_fingerprint.""" + _write_manager_hb(layout) + write_pending(layout, _spec("0000001-ts01", fingerprint="same")) + write_pending(layout, _spec("0000002-ts02", fingerprint="same")) + + load_calls = [] + + def fake_evaluate(**kwargs): + return (Path(kwargs["tilt_series_path"]), [0.1]) + + def fake_load(path, map_location=None): + load_calls.append(path) + return MagicMock() + + with ( + patch( + "miss_alignment.distributed.worker.evaluate_tilt_series", + side_effect=fake_evaluate, + ), + patch( + "miss_alignment.distributed.worker.MissAlignment.load_from_checkpoint", + side_effect=fake_load, + ), + ): + run_worker_loop(layout, "worker-0", "cpu", manager_hb_timeout_s=30.0) + + assert len(load_calls) == 1 # loaded once despite two tasks + + +def test_worker_reloads_model_when_fingerprint_changes(layout): + """Model is reloaded when fingerprint differs between tasks.""" + _write_manager_hb(layout) + write_pending(layout, _spec("0000001-ts01", fingerprint="fp-a")) + write_pending(layout, _spec("0000002-ts02", fingerprint="fp-b")) + + load_calls = [] + + def fake_evaluate(**kwargs): + return (Path(kwargs["tilt_series_path"]), [0.1]) + + def fake_load(path, map_location=None): + load_calls.append(path) + return MagicMock() + + with ( + patch( + "miss_alignment.distributed.worker.evaluate_tilt_series", + side_effect=fake_evaluate, + ), + patch( + "miss_alignment.distributed.worker.MissAlignment.load_from_checkpoint", + side_effect=fake_load, + ), + ): + run_worker_loop(layout, "worker-0", "cpu", manager_hb_timeout_s=30.0) + + assert len(load_calls) == 2 + + +def test_worker_dispatches_prepare_stacks(layout): + """prepare_stacks tasks call _prepare_single_tilt_series, not evaluate.""" + _write_manager_hb(layout) + spec = TaskSpec( + task_id="0000001-ts01", + model_checkpoint_path="", + tilt_series_path="/data/ts01.xml", + output_directory="/data/out", + setting="", + patch_size=0, + patch_overlap=0.0, + batch_size=0, + apply_ctf=False, + downsample=1, + init_fingerprint="", + task_type="prepare_stacks", + desired_pixel_size=10.0, + ) + write_pending(layout, spec) + + prepare_calls = [] + + def fake_prepare(xml_path, desired_pixel_size, device): + prepare_calls.append((str(xml_path), desired_pixel_size, device)) + + with patch( + "miss_alignment.distributed.worker._prepare_single_tilt_series", + side_effect=fake_prepare, + ): + run_worker_loop(layout, "worker-0", "cpu", manager_hb_timeout_s=30.0) + + assert len(prepare_calls) == 1 + assert prepare_calls[0][1] == 10.0 + assert (layout.done / "0000001-ts01.json").exists() + + +def test_worker_dispatches_cross_correlation(layout): + """cross_correlation tasks call _run_cross_correlation_single.""" + _write_manager_hb(layout) + spec = TaskSpec( + task_id="0000001-ts01", + model_checkpoint_path="", + tilt_series_path="/data/ts01.xml", + output_directory="/data/out", + setting="", + patch_size=0, + patch_overlap=0.0, + batch_size=0, + apply_ctf=False, + downsample=1, + init_fingerprint="", + task_type="cross_correlation", + lowpass_cutoff=0.25, + ) + write_pending(layout, spec) + + xcorr_calls = [] + + def fake_xcorr(xml_file, device, lowpass_cutoff): + xcorr_calls.append((str(xml_file), lowpass_cutoff)) + return None + + with patch( + "miss_alignment.distributed.worker._run_cross_correlation_single", + side_effect=fake_xcorr, + ): + run_worker_loop(layout, "worker-0", "cpu", manager_hb_timeout_s=30.0) + + assert len(xcorr_calls) == 1 + assert xcorr_calls[0][1] == 0.25 + assert (layout.done / "0000001-ts01.json").exists() diff --git a/tests/test_parallel.py b/tests/test_parallel.py index 92d6594a..b75c5fd3 100644 --- a/tests/test_parallel.py +++ b/tests/test_parallel.py @@ -1,47 +1,45 @@ -"""Tests for the shared one-process-per-device work-queue helper.""" - -import queue +"""Tests for the distributed worker provisioner (replaces _parallel.py tests).""" +from pathlib import Path +from unittest.mock import MagicMock, patch import pytest -from miss_alignment._parallel import run_device_pool - - -def _square_runner(device, task_queue, result_queue): - """Process int jobs by squaring them (device is ignored; no CUDA needed).""" - while True: - try: - n = task_queue.get_nowait() - except queue.Empty: - break - result_queue.put_nowait(n * n) - - -def _failing_runner(device, task_queue, result_queue): - """Raise on a specific job to exercise worker-failure detection.""" - while True: - try: - n = task_queue.get_nowait() - except queue.Empty: - break - if n == 3: - raise ValueError("boom") - result_queue.put_nowait(n * n) - - -@pytest.mark.filterwarnings("ignore") -def test_run_device_pool_processes_all_jobs(): - """Every job is processed exactly once across multiple worker processes.""" - jobs = [1, 2, 3, 4, 5] - results = run_device_pool( - jobs, _square_runner, runner_args=(), devices=[0, 1], desc="test" - ) - assert sorted(results) == [1, 4, 9, 16, 25] - - -@pytest.mark.filterwarnings("ignore") -def test_run_device_pool_raises_on_worker_failure(): - """A worker that dies with a non-zero exit code surfaces as RuntimeError.""" - jobs = [1, 2, 3, 4, 5] - with pytest.raises(RuntimeError, match="stopped unexpectedly"): - run_device_pool(jobs, _failing_runner, runner_args=(), devices=[0], desc="test") +from miss_alignment.distributed.provisioner import LocalProvisioner + + +def test_local_provisioner_spawns_worker_per_device(tmp_path): + with patch("miss_alignment.distributed.provisioner.subprocess.Popen") as mock_popen: + mock_proc = MagicMock() + mock_proc.poll.return_value = None + mock_popen.return_value = mock_proc + + p = LocalProvisioner(queue_dir=tmp_path, devices=[0, 1, 2]) + p.ensure_workers(n_workers=10) + + assert mock_popen.call_count == 3 + + +def test_local_provisioner_does_not_double_spawn(tmp_path): + with patch("miss_alignment.distributed.provisioner.subprocess.Popen") as mock_popen: + mock_proc = MagicMock() + mock_proc.poll.return_value = None + mock_popen.return_value = mock_proc + + p = LocalProvisioner(queue_dir=tmp_path, devices=[0]) + p.ensure_workers(n_workers=5) + p.ensure_workers(n_workers=5) + + assert mock_popen.call_count == 1 + + +def test_local_provisioner_shutdown_terminates(tmp_path): + with patch("miss_alignment.distributed.provisioner.subprocess.Popen") as mock_popen: + mock_proc = MagicMock() + mock_proc.poll.return_value = None + mock_popen.return_value = mock_proc + + p = LocalProvisioner(queue_dir=tmp_path, devices=[0]) + p.ensure_workers(n_workers=5) + p.shutdown() + + mock_proc.terminate.assert_called()