From 685ee4d471f08c6ff5ce3d3efe3f30a2b26e3174 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 17:03:18 -0700 Subject: [PATCH 01/33] docs: add distributed inference design spec Co-Authored-By: Claude Fable 5 --- ...2026-07-04-distributed-inference-design.md | 229 ++++++++++++++++++ 1 file changed, 229 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-04-distributed-inference-design.md diff --git a/docs/superpowers/specs/2026-07-04-distributed-inference-design.md b/docs/superpowers/specs/2026-07-04-distributed-inference-design.md new file mode 100644 index 0000000..84a6f07 --- /dev/null +++ b/docs/superpowers/specs/2026-07-04-distributed-inference-design.md @@ -0,0 +1,229 @@ +# Distributed Inference Design + +**Date:** 2026-07-04 +**Status:** Approved +**Scope:** Cluster distribution of the per-tilt-series alignment/inference phase + +--- + +## Problem + +With 300 tilt-series, the alignment phase (inference) takes 3× longer than the preceding training phase. Each series is independently optimizable — no data exchange between series — but the current implementation is limited to GPUs on the local machine. The head node has to finish all series sequentially across its local GPU pool before the next macro-iteration can begin. + +## Goal + +Distribute the per-series alignment tasks across a compute cluster so the head node fans out work, blocks until all series are done, and then continues to the next macro-iteration. Cluster mode is opt-in via environment variables; without them the system runs exactly as today (local multi-GPU pool). + +--- + +## Architecture + +A new `miss_alignment/distributed/` module sits between the existing `alignment/parallel.py` public API and the underlying `evaluate_tilt_series` function. `train.py` and `infer.py` are unchanged. + +`run_alignment_parallel` (in `alignment/parallel.py`) calls `load_cluster_config()`. If cluster config is present it delegates to the distributed manager + `ClusterProvisioner`; otherwise it uses the distributed manager + `LocalProvisioner` (replacing the current `_parallel.py` internals). In both cases the same queue layer is used. + +### Components + +| File | Role | +|---|---| +| `distributed/queue.py` | Queue directory layout, task JSON read/write, atomic rename claim | +| `distributed/manager.py` | Head-node coordinator: writes tasks, runs scheduler thread, blocks until done | +| `distributed/provisioner.py` | `WorkerProvisioner` interface + `LocalProvisioner` + `ClusterProvisioner` | +| `distributed/worker.py` | `miss-alignment worker` subcommand: claims tasks, runs inference, writes results | +| `distributed/config.py` | Reads env vars, returns `ClusterConfig` dataclass or `None` | + +--- + +## Task Format + +One JSON file per tilt-series, written to `/tasks/pending/-.json`: + +```json +{ + "task_id": "0000003-tilt_series_01", + "model_checkpoint_path": "/data/project/iter2/model.ckpt", + "tilt_series_path": "/data/project/tilt_series_01.xml", + "output_directory": "/data/project/iter2/", + "setting": "anchoring", + "patch_size": 96, + "patch_overlap": 0.1, + "batch_size": 32, + "apply_ctf": false, + "downsample": 2, + "init_fingerprint": "" +} +``` + +`init_fingerprint` covers the model checkpoint path and all alignment parameters. A worker skips reloading the model when consecutive tasks share the same fingerprint, amortizing checkpoint loading across many series in one macro-iteration. + +On completion, result fields are appended before writing to `done/`: +```json +{ "final_loss": 0.0312, "device": "cuda:0" } +``` + +On failure, error fields are appended before writing to `failed/`: +```json +{ "error": "CUDA out of memory...", "worker_id": "local-12345-gpu0" } +``` + +--- + +## Queue Directory Layout + +``` +/tasks/ +├── pending/ # one JSON per queued task +├── running/ +│ └── / # per-worker subdir +│ ├── .json # claimed task lives here during execution +│ └── hb- # worker heartbeat tick files (latest only) +├── done/ # completed task JSONs +├── failed/ # failed task JSONs +└── manager/ + └── hb- # manager heartbeat tick files (latest only) +``` + +The queue directory is always `/tasks/`. Since the training directory must already be on a shared filesystem (workers need the XML/MRC files), no additional configuration is needed for cluster nodes to access it. + +--- + +## Claim Protocol + +1. Worker lists `pending/`, **shuffles randomly** (avoids thundering herd on the same lexicographically-first file) +2. Attempts `os.rename(pending/.json, running//.json)` +3. `FileNotFoundError` means another worker won the race — try the next candidate +4. On task completion: write `done/.json`, then delete `running//.json` (publish-before-delete: a crash mid-step leaves an orphan in `running/` that the scheduler sweeps, not a lost task) +5. On task failure: write `failed/.json` with error, delete running copy, continue + +No lock files. The OS rename is the coordination primitive. + +--- + +## Worker Lifecycle + +**Startup** (`miss-alignment worker --queue-dir --device [--worker-id ]`): +1. Derive `worker_id` from `--worker-id` or `local--gpu` +2. Create `tasks/running//` +3. Start background heartbeat thread: write `tasks/running//hb-` every 5s, keep only the latest tick file +4. Enter claim loop + +**Claim loop:** +- Check manager heartbeat (`tasks/manager/hb-*`): if latest tick is >120s old → exit cleanly (manager is dead) +- List `pending/`, shuffle, attempt rename +- On empty queue: exit cleanly +- On claim: check if `init_fingerprint` matches last loaded model; if not, load checkpoint from `model_checkpoint_path` +- Call `evaluate_tilt_series(..., device=f"cuda:{device}")` — all internal LBFGS retries and optimization passes run unchanged +- Write result to `done/` or `failed/`, delete running copy, loop + +**No task-level retries.** Failures are deterministic (bad data, corrupt file, config error); the manager hard-fails after all remaining tasks complete. + +--- + +## Manager Lifecycle + +**Startup:** +1. Clear stale queue state from any prior run (delete `pending/`, `done/`, `failed/` contents; recover any `running/` orphans back to `pending/`) +2. Write all task JSONs to `pending/` +3. Start `ClusterProvisioner` or `LocalProvisioner` +4. Start scheduler background thread +5. Block in poll loop (500ms interval) until all tasks are in `done/` or `failed/` + +**Scheduler thread** (runs every ~10s): +- Write manager heartbeat to `tasks/manager/hb-` +- Sweep stalled workers: for each `running//`, check age of latest `hb-` file; if >120s stale → move task back to `pending/`, delete `running//` +- Call `provisioner.ensure_workers()` to respawn any dead local workers + +**Shutdown** (on completion or `KeyboardInterrupt`): +- Call `provisioner.shutdown()` (SIGTERM local children / cancel SLURM job IDs) +- Delete `tasks/` directory +- Return `dict[series_name → final_loss]` on full success, or raise with list of failed series if any task is in `failed/` + +**Failure policy:** if any series ends in `failed/`, the manager raises after all other tasks complete. `train.py` and `infer.py` propagate this as a hard failure — training stops. + +--- + +## Provisioners + +### `LocalProvisioner` + +Activated when neither `MISS_CLUSTER_CONFIG` nor `MISS_CLUSTER_SCRIPT` is set. + +- `ensure_workers(target)`: spawns `miss-alignment worker --queue-dir --device ` via `subprocess.Popen`, one per GPU in `devices_alignment` (same list as today: `list(range(torch.cuda.device_count()))`, respecting `CUDA_VISIBLE_DEVICES`) +- Respawns any that have exited prematurely (checked each scheduler tick) +- `shutdown()`: SIGTERM all children, short timeout, SIGKILL if needed + +This replaces `_parallel.py`'s `run_device_pool` / `mp.spawn` internals with no behavioral change for the local case. + +### `ClusterProvisioner` + +Activated when both `MISS_CLUSTER_CONFIG` and `MISS_CLUSTER_SCRIPT` are set. + +**`MISS_CLUSTER_CONFIG`** — path to a JSON file: +```json +{ + "submit": "sbatch {{script_path}}", + "submit_job_id_regex": "Submitted batch job (\\d+)", + "cancel": "scancel {{job_id}}" +} +``` + +**`MISS_CLUSTER_SCRIPT`** — path to a `.sh` template: +```bash +#!/bin/bash +#SBATCH --nodes=1 +#SBATCH --gres=gpu:1 +#SBATCH --time=04:00:00 + +conda activate miss-alignment +{{command}} +``` + +- `{{command}}` is filled with the `miss-alignment worker` invocation by the provisioner +- All cluster-specific settings (partition, memory, time limit, environment setup) are the user's responsibility in the template +- Additional `{{custom_var}}` placeholders filled via `MISS_CLUSTER_VAR_=` env vars +- Rendered scripts are written to `tasks/cluster/worker-.sh` +- One job submitted per tilt-series; job IDs stored for cancellation +- `shutdown()`: runs configured `cancel` command for each stored job ID; registers SIGINT/SIGTERM handlers so Ctrl-C on the head node cancels the cluster pool + +--- + +## Configuration + +All configuration via environment variables — no new CLI flags on `train` or `infer`: + +| Env var | Purpose | +|---|---| +| `MISS_CLUSTER_CONFIG` | Path to cluster scheduler JSON config. Activates cluster mode when set together with `MISS_CLUSTER_SCRIPT`. | +| `MISS_CLUSTER_SCRIPT` | Path to job submission shell script template. | +| `MISS_CLUSTER_VAR_` | Additional template variables, e.g. `MISS_CLUSTER_VAR_partition=gpu`. | + +`config.py` reads these at call time (not import time) and returns a `ClusterConfig` dataclass or `None`. + +--- + +## What Does Not Change + +- `evaluate_tilt_series` — called identically, all internal optimization passes unchanged +- `train.py` / `infer.py` — no changes; `run_alignment_parallel` signature unchanged +- LBFGS retries, anchoring iterations, coarse-to-fine spline passes — all internal to `evaluate_tilt_series` +- `CUDA_VISIBLE_DEVICES` still controls which local GPUs are used +- Single-GPU behavior is identical (one `LocalProvisioner` worker process per GPU) + +--- + +## File Changes Summary + +**New files:** +- `src/miss_alignment/distributed/__init__.py` +- `src/miss_alignment/distributed/queue.py` +- `src/miss_alignment/distributed/manager.py` +- `src/miss_alignment/distributed/provisioner.py` +- `src/miss_alignment/distributed/worker.py` +- `src/miss_alignment/distributed/config.py` + +**Modified files:** +- `src/miss_alignment/alignment/parallel.py` — replace `run_device_pool` delegation with distributed manager call +- `src/miss_alignment/_parallel.py` — kept for reference, internals superseded by `LocalProvisioner` +- `src/miss_alignment/_cli.py` — register `worker` subcommand + +**Deleted files:** none (old `_parallel.py` left in place until the new path is validated) From 77397e6b93be24767c04843b14246cff932d9ee8 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 17:04:02 -0700 Subject: [PATCH 02/33] docs: fix spec ambiguities (worker-id uniqueness, _parallel.py deletion) Co-Authored-By: Claude Fable 5 --- .../specs/2026-07-04-distributed-inference-design.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/superpowers/specs/2026-07-04-distributed-inference-design.md b/docs/superpowers/specs/2026-07-04-distributed-inference-design.md index 84a6f07..c36608a 100644 --- a/docs/superpowers/specs/2026-07-04-distributed-inference-design.md +++ b/docs/superpowers/specs/2026-07-04-distributed-inference-design.md @@ -182,6 +182,7 @@ conda activate miss-alignment - All cluster-specific settings (partition, memory, time limit, environment setup) are the user's responsibility in the template - Additional `{{custom_var}}` placeholders filled via `MISS_CLUSTER_VAR_=` env vars - Rendered scripts are written to `tasks/cluster/worker-.sh` +- The `{{command}}` is `miss-alignment worker --queue-dir --device 0 --worker-id "$(hostname)-$$-"` — the `$(hostname)` and `$$` are expanded by the compute node's shell at runtime, guaranteeing unique, stable worker IDs even across resubmissions - One job submitted per tilt-series; job IDs stored for cancellation - `shutdown()`: runs configured `cancel` command for each stored job ID; registers SIGINT/SIGTERM handlers so Ctrl-C on the head node cancels the cluster pool @@ -223,7 +224,7 @@ All configuration via environment variables — no new CLI flags on `train` or ` **Modified files:** - `src/miss_alignment/alignment/parallel.py` — replace `run_device_pool` delegation with distributed manager call -- `src/miss_alignment/_parallel.py` — kept for reference, internals superseded by `LocalProvisioner` - `src/miss_alignment/_cli.py` — register `worker` subcommand -**Deleted files:** none (old `_parallel.py` left in place until the new path is validated) +**Deleted files:** +- `src/miss_alignment/_parallel.py` — superseded by `LocalProvisioner`; deleted once the new path is validated in tests From bc1d2b15d7dcc365ee01e446e62aa0263043bdbc Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 17:19:44 -0700 Subject: [PATCH 03/33] docs: add distributed inference implementation plan Co-Authored-By: Claude Fable 5 --- .../plans/2026-07-04-distributed-inference.md | 1908 +++++++++++++++++ 1 file changed, 1908 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-04-distributed-inference.md diff --git a/docs/superpowers/plans/2026-07-04-distributed-inference.md b/docs/superpowers/plans/2026-07-04-distributed-inference.md new file mode 100644 index 0000000..2de3d7a --- /dev/null +++ b/docs/superpowers/plans/2026-07-04-distributed-inference.md @@ -0,0 +1,1908 @@ +# Distributed Inference Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add a disk-based task queue that distributes per-tilt-series alignment inference across cluster nodes (SLURM/PBS), while keeping the existing local multi-GPU behaviour unchanged. + +**Architecture:** A new `miss_alignment/distributed/` package implements a filesystem queue (atomic rename as the claim mutex), a manager that writes tasks and blocks until done, two provisioners (local subprocess and cluster batch scheduler), and a `miss-alignment worker` subcommand. `run_alignment_parallel` in `alignment/parallel.py` is updated to drive the manager instead of `_parallel.run_device_pool`. No changes to `train.py`, `infer.py`, or `evaluate_tilt_series`. + +**Tech Stack:** Python 3.10+, stdlib only (`pathlib`, `os`, `subprocess`, `threading`, `hashlib`, `json`, `re`, `time`, `signal`) — no new dependencies. + +## Global Constraints + +- Python ≥ 3.10 (project minimum). +- No new runtime dependencies beyond stdlib. +- `ruff check --fix && ruff format` must pass (line length 88, ignore E712). +- `pytest --color=yes` must pass with no warnings-as-errors regressions. +- All new files under `src/miss_alignment/distributed/`. +- Cluster mode activated only when **both** `MISS_CLUSTER_CONFIG` and `MISS_CLUSTER_SCRIPT` env vars are set; otherwise `LocalProvisioner` is used. +- `evaluate_tilt_series` is never modified. +- `train.py` and `infer.py` are never modified. + +--- + +## File Map + +| Path | Action | Responsibility | +|---|---|---| +| `src/miss_alignment/distributed/__init__.py` | Create | Re-export public API | +| `src/miss_alignment/distributed/queue.py` | Create | Directory layout, task JSON, atomic rename claim | +| `src/miss_alignment/distributed/manager.py` | Create | Head-node coordinator, scheduler thread, poll loop | +| `src/miss_alignment/distributed/provisioner.py` | Create | `WorkerProvisioner` ABC, `LocalProvisioner`, `ClusterProvisioner` | +| `src/miss_alignment/distributed/worker.py` | Create | `miss-alignment worker` subcommand logic | +| `src/miss_alignment/distributed/config.py` | Create | Read env vars, return `ClusterConfig \| None` | +| `src/miss_alignment/alignment/parallel.py` | Modify | Replace `run_device_pool` call with manager call | +| `src/miss_alignment/_cli.py` | Modify | Register `worker` subcommand | +| `src/miss_alignment/__init__.py` | Modify | Export `worker_miss_align` | +| `src/miss_alignment/_parallel.py` | Delete (after Task 6) | Superseded by `LocalProvisioner` | +| `tests/distributed/test_queue.py` | Create | Queue layer unit tests | +| `tests/distributed/test_manager.py` | Create | Manager + provisioner integration tests | +| `tests/distributed/test_worker.py` | Create | Worker claim loop unit tests | +| `tests/distributed/test_config.py` | Create | Config env-var parsing tests | +| `tests/test_parallel.py` | Modify | Update import from new path | + +--- + +## Task 1: Queue layer (`distributed/queue.py`) + +**Files:** +- Create: `src/miss_alignment/distributed/__init__.py` +- Create: `src/miss_alignment/distributed/queue.py` +- Create: `tests/distributed/__init__.py` +- Create: `tests/distributed/test_queue.py` + +**Interfaces:** +- Produces: + - `QueueLayout(root: Path)` — manages subdirectory creation; attributes `pending`, `running`, `done`, `failed`, `manager_hb`, each a `Path`. + - `TaskSpec` — `dataclass` with fields: `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`. + - `write_pending(layout: QueueLayout, spec: TaskSpec) -> None` — writes `layout.pending/.json`. + - `claim_one(layout: QueueLayout, worker_id: str) -> TaskSpec | None` — shuffled rename claim; returns `None` when queue empty. + - `mark_done(layout: QueueLayout, worker_id: str, spec: TaskSpec, final_loss: float, device: str) -> None` + - `mark_failed(layout: QueueLayout, worker_id: str, spec: TaskSpec, error: str) -> None` + - `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 hex. + - `clear_queue(layout: QueueLayout) -> None` — deletes all JSON files in `pending/`, `done/`, `failed/`; moves any `running//.json` back to `pending/` (orphan recovery). + +- [ ] **Step 1: Create directory skeleton and write failing tests** + +```bash +mkdir -p tests/distributed +touch tests/distributed/__init__.py +``` + +```python +# tests/distributed/test_queue.py +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, tmp_path): + """Two concurrent claimers: exactly one wins.""" + write_pending(layout, _spec()) + results = [] + results.append(claim_one(layout, "worker-0")) + results.append(claim_one(layout, "worker-1")) + claimed = [r for r in results 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, we clear and recover + clear_queue(layout) + # orphan should be back in pending + assert (layout.pending / "0000001-ts01.json").exists() + + +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 +``` + +- [ ] **Step 2: Run tests to verify they fail** + +```bash +cd /Users/tegunovd/dev/miss-alignment +pytest tests/distributed/test_queue.py -v 2>&1 | head -30 +``` + +Expected: `ModuleNotFoundError: No module named 'miss_alignment.distributed'` + +- [ ] **Step 3: Create `__init__.py` skeleton** + +```python +# src/miss_alignment/distributed/__init__.py +"""Disk-based distributed task queue for miss-alignment inference.""" +``` + +- [ ] **Step 4: Implement `queue.py`** + +```python +# src/miss_alignment/distributed/queue.py +"""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" + + def ensure_directories(self) -> None: + for d in ( + self.pending, + self.running, + self.done, + self.failed, + self.manager_hb, + self.cluster, + ): + 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 + + +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) + running_path = layout.worker_dir(worker_id) / f"{spec.task_id}.json" + running_path.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) + running_path = layout.worker_dir(worker_id) / f"{spec.task_id}.json" + running_path.unlink(missing_ok=True) + + +def clear_queue(layout: QueueLayout) -> None: + """Delete stale queue state from a prior run; recover running orphans to pending.""" + # recover orphaned running tasks back to pending + 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 + try: + worker_dir.rmdir() + except OSError: + pass # not empty; will be swept later + + for directory in (layout.pending, layout.done, layout.failed): + for f in directory.glob("*.json"): + f.unlink(missing_ok=True) +``` + +- [ ] **Step 5: Run tests to verify they pass** + +```bash +pytest tests/distributed/test_queue.py -v +``` + +Expected: all 9 tests PASS. + +- [ ] **Step 6: Commit** + +```bash +git add src/miss_alignment/distributed/__init__.py src/miss_alignment/distributed/queue.py tests/distributed/__init__.py tests/distributed/test_queue.py +git commit -m "feat: add distributed queue layer with atomic rename claim protocol" +``` + +--- + +## Task 2: Cluster configuration (`distributed/config.py`) + +**Files:** +- Create: `src/miss_alignment/distributed/config.py` +- Create: `tests/distributed/test_config.py` + +**Interfaces:** +- Consumes: nothing. +- Produces: + - `ClusterConfig` — `dataclass` with fields: `submit: str`, `submit_job_id_regex: str`, `cancel: str`, `script_path: Path`. + - `load_cluster_config() -> ClusterConfig | None` — reads `MISS_CLUSTER_CONFIG` and `MISS_CLUSTER_SCRIPT`; returns `None` if either is unset. + +- [ ] **Step 1: Write failing tests** + +```python +# tests/distributed/test_config.py +import os +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}}", + } + 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_returns_none_when_unset(monkeypatch): + monkeypatch.delenv("MISS_CLUSTER_CONFIG", raising=False) + monkeypatch.delenv("MISS_CLUSTER_SCRIPT", raising=False) + assert load_cluster_config() is None + + +def test_load_cluster_config_returns_none_when_only_one_set(monkeypatch, cluster_json): + monkeypatch.setenv("MISS_CLUSTER_CONFIG", str(cluster_json)) + monkeypatch.delenv("MISS_CLUSTER_SCRIPT", raising=False) + assert load_cluster_config() is None + + +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 + + +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() +``` + +- [ ] **Step 2: Run tests to verify they fail** + +```bash +pytest tests/distributed/test_config.py -v 2>&1 | head -20 +``` + +Expected: `ImportError` or `ModuleNotFoundError`. + +- [ ] **Step 3: Implement `config.py`** + +```python +# src/miss_alignment/distributed/config.py +"""Read cluster configuration from environment variables. + +Cluster mode is activated when both MISS_CLUSTER_CONFIG and MISS_CLUSTER_SCRIPT +are set. If either is absent, load_cluster_config() returns None and +LocalProvisioner is used instead. +""" + +from __future__ import annotations + +import json +import os +from dataclasses import dataclass +from pathlib import Path + + +@dataclass +class ClusterConfig: + submit: str + submit_job_id_regex: str + cancel: str + script_path: Path + + +def load_cluster_config() -> ClusterConfig | None: + """Return ClusterConfig if both env vars are set, else None.""" + config_path_str = os.environ.get("MISS_CLUSTER_CONFIG") + script_path_str = os.environ.get("MISS_CLUSTER_SCRIPT") + + if not config_path_str or not script_path_str: + return None + + 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"], + script_path=script_path, + ) +``` + +- [ ] **Step 4: Run tests to verify they pass** + +```bash +pytest tests/distributed/test_config.py -v +``` + +Expected: all 5 tests PASS. + +- [ ] **Step 5: Commit** + +```bash +git add src/miss_alignment/distributed/config.py tests/distributed/test_config.py +git commit -m "feat: add cluster config reader (MISS_CLUSTER_CONFIG + MISS_CLUSTER_SCRIPT)" +``` + +--- + +## Task 3: Worker subcommand (`distributed/worker.py`) + +**Files:** +- Create: `src/miss_alignment/distributed/worker.py` +- Create: `tests/distributed/test_worker.py` + +**Interfaces:** +- Consumes: + - `QueueLayout`, `TaskSpec`, `claim_one`, `mark_done`, `mark_failed` from `distributed/queue.py` +- Produces: + - `worker_miss_align(queue_dir: Path, device: int, worker_id: str | None)` — Typer command entry point; loops until queue empty or manager heartbeat stale. + - `run_worker_loop(layout: QueueLayout, worker_id: str, device: str, manager_hb_timeout_s: float) -> None` — testable loop body. + +- [ ] **Step 1: Write failing tests** + +```python +# tests/distributed/test_worker.py +"""Unit tests for the worker claim loop. + +These tests do NOT call evaluate_tilt_series; they mock it to keep +tests fast and free of CUDA/warpylib dependencies. +""" +import json +import time +from pathlib import Path +from unittest.mock import 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): + """Write a fresh manager heartbeat tick.""" + 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"): + 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_worker_processes_task_and_writes_done(layout, tmp_path): + _write_manager_hb(layout) + write_pending(layout, _spec()) + + fake_loss = [0.5, 0.3, 0.1] + with patch( + "miss_alignment.distributed.worker.evaluate_tilt_series", + return_value=(Path("/data/ts01.xml"), fake_loss), + ): + 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"), + ): + 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 manager heartbeat that is already old + hb_file = layout.manager_hb / "hb-0" + hb_file.write_text("") + # Make it appear 200 seconds old by back-dating mtime + old_time = time.time() - 200 + import os + 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), + ): + run_worker_loop(layout, "worker-0", "cpu", manager_hb_timeout_s=120.0) + + # Worker should exit without processing the task + 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) + spec1 = _spec("0000001-ts01") + spec2 = TaskSpec( + task_id="0000002-ts02", + model_checkpoint_path="/data/model.ckpt", + tilt_series_path="/data/ts02.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", # same fingerprint + ) + write_pending(layout, spec1) + write_pending(layout, spec2) + + load_calls = [] + + def fake_evaluate(**kwargs): + return (Path(kwargs["tilt_series_path"]), [0.1]) + + with patch( + "miss_alignment.distributed.worker.evaluate_tilt_series", + side_effect=fake_evaluate, + ): + with patch( + "miss_alignment.distributed.worker.MissAlignment.load_from_checkpoint", + ) as mock_load: + mock_load.return_value = mock_load # return self as stub + run_worker_loop(layout, "worker-0", "cpu", manager_hb_timeout_s=30.0) + # Model should only be loaded once despite two tasks + assert mock_load.call_count == 1 +``` + +- [ ] **Step 2: Run tests to verify they fail** + +```bash +pytest tests/distributed/test_worker.py -v 2>&1 | head -20 +``` + +Expected: `ImportError` for `miss_alignment.distributed.worker`. + +- [ ] **Step 3: Implement `worker.py`** + +```python +# src/miss_alignment/distributed/worker.py +"""Worker subcommand: claims tasks from the queue and runs evaluate_tilt_series. + +Usage (launched by provisioner): + miss-alignment worker --queue-dir --device [--worker-id ] +""" + +from __future__ import annotations + +import os +import sys +import time +import traceback +from pathlib import Path + +import torch +import typer + +from ..alignment.tilt_series import evaluate_tilt_series +from ..models.models import MissAlignment +from .queue import ( + QueueLayout, + TaskSpec, + claim_one, + mark_done, + mark_failed, +) + +# Seconds without a manager heartbeat tick before the worker exits. +_MANAGER_HB_TIMEOUT_S = 120.0 +# Seconds between heartbeat writes. +_HB_INTERVAL_S = 5.0 + + +def _write_worker_hb(worker_dir: Path, seq: int) -> None: + """Write a new heartbeat tick, removing the previous one.""" + new_hb = worker_dir / f"hb-{seq}" + new_hb.write_text("") + if seq > 0: + old_hb = worker_dir / f"hb-{seq - 1}" + old_hb.unlink(missing_ok=True) + + +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 run_worker_loop( + layout: QueueLayout, + worker_id: str, + device: str, + manager_hb_timeout_s: float = _MANAGER_HB_TIMEOUT_S, +) -> None: + """Main worker loop: claim → check heartbeat → evaluate → write result. + + Separated from the Typer command for testability. + """ + worker_dir = layout.worker_dir(worker_id) + worker_dir.mkdir(parents=True, exist_ok=True) + + last_fingerprint: str | None = None + loaded_model = None + hb_seq = 0 + last_hb_time = 0.0 + + while True: + # Check manager heartbeat before every claim attempt. + age = _manager_hb_age_s(layout) + if age > manager_hb_timeout_s: + print( + f"[{worker_id}] Manager heartbeat stale ({age:.0f}s > " + f"{manager_hb_timeout_s:.0f}s). Exiting.", + file=sys.stderr, + ) + return + + # Write our own heartbeat if due. + now = time.time() + if now - last_hb_time >= _HB_INTERVAL_S: + _write_worker_hb(worker_dir, hb_seq) + hb_seq += 1 + last_hb_time = now + + spec = claim_one(layout, worker_id) + if spec is None: + return # queue empty, exit cleanly + + print(f"[{worker_id}] Claimed {spec.task_id}", file=sys.stderr) + + # Load model only when fingerprint changes. + if spec.init_fingerprint != last_fingerprint: + loaded_model = MissAlignment.load_from_checkpoint( + spec.model_checkpoint_path, map_location="cpu" + ) + last_fingerprint = spec.init_fingerprint + + try: + _, 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=spec.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, + ) + final_loss = float(loss_values[-1]) if loss_values else float("nan") + mark_done(layout, worker_id, spec, final_loss=final_loss, device=device) + print( + f"[{worker_id}] Done {spec.task_id} loss={final_loss:.4f}", + file=sys.stderr, + ) + except Exception: + error = traceback.format_exc() + mark_failed(layout, worker_id, spec, error=error) + print( + f"[{worker_id}] Failed {spec.task_id}:\n{error}", + file=sys.stderr, + ) + + +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: str | None = 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" + + run_worker_loop(layout, worker_id, cuda_device) +``` + +- [ ] **Step 4: Run tests to verify they pass** + +```bash +pytest tests/distributed/test_worker.py -v +``` + +Expected: all 4 tests PASS. + +- [ ] **Step 5: Commit** + +```bash +git add src/miss_alignment/distributed/worker.py tests/distributed/test_worker.py +git commit -m "feat: add worker subcommand with claim loop and model fingerprint reuse" +``` + +--- + +## Task 4: Provisioners (`distributed/provisioner.py`) + +**Files:** +- Create: `src/miss_alignment/distributed/provisioner.py` +- Create: `tests/distributed/test_provisioner.py` + +**Interfaces:** +- Consumes: `ClusterConfig` from `distributed/config.py`. +- Produces: + - `WorkerProvisioner` — ABC with `ensure_workers(n_tasks: int) -> None` and `shutdown() -> None`. + - `LocalProvisioner(queue_dir: Path, devices: list[int])` — spawns `miss-alignment worker` child processes. + - `ClusterProvisioner(queue_dir: Path, config: ClusterConfig, n_tasks: int)` — submits cluster jobs. + +- [ ] **Step 1: Write failing tests** + +```python +# tests/distributed/test_provisioner.py +"""Tests for LocalProvisioner and ClusterProvisioner.""" +import subprocess +import sys +from pathlib import Path +from unittest.mock import MagicMock, call, patch + +import pytest + +from miss_alignment.distributed.config import ClusterConfig +from miss_alignment.distributed.provisioner import ClusterProvisioner, LocalProvisioner + + +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 # still running + mock_popen.return_value = mock_proc + + p = LocalProvisioner(queue_dir=tmp_path, devices=[0, 1]) + p.ensure_workers(n_tasks=10) + + assert mock_popen.call_count == 2 + # Each call should pass --device 0 and --device 1 + calls_str = [str(c) for c in mock_popen.call_args_list] + assert any("--device" in s and "0" in s for s in calls_str) + assert any("--device" in s and "1" in s for s in calls_str) + + +def test_local_provisioner_shutdown_terminates_processes(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_tasks=5) + p.shutdown() + + mock_proc.terminate.assert_called() + + +def test_local_provisioner_does_not_respawn_running_processes(tmp_path): + with patch("miss_alignment.distributed.provisioner.subprocess.Popen") as mock_popen: + mock_proc = MagicMock() + mock_proc.poll.return_value = None # still running + mock_popen.return_value = mock_proc + + p = LocalProvisioner(queue_dir=tmp_path, devices=[0]) + p.ensure_workers(n_tasks=5) + p.ensure_workers(n_tasks=5) # second call should not spawn again + + assert mock_popen.call_count == 1 + + +def test_cluster_provisioner_submits_one_job_per_task(tmp_path): + script = tmp_path / "worker.sh" + script.write_text("#!/bin/bash\n{{command}}\n") + cfg = ClusterConfig( + submit="sbatch {{script_path}}", + submit_job_id_regex=r"Submitted batch job (\d+)", + cancel="scancel {{job_id}}", + script_path=script, + ) + + submitted = [] + + def fake_run(cmd, **kwargs): + submitted.append(cmd) + result = MagicMock() + result.stdout = "Submitted batch job 12345\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_tasks=3) + + assert len(submitted) == 3 + + +def test_cluster_provisioner_cancels_jobs_on_shutdown(tmp_path): + script = tmp_path / "worker.sh" + script.write_text("#!/bin/bash\n{{command}}\n") + cfg = ClusterConfig( + submit="sbatch {{script_path}}", + submit_job_id_regex=r"Submitted batch job (\d+)", + cancel="scancel {{job_id}}", + script_path=script, + ) + + cancel_calls = [] + + def fake_run(cmd, **kwargs): + if "sbatch" in cmd: + result = MagicMock() + result.stdout = "Submitted batch job 99999\n" + 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_tasks=2) + p.shutdown() + + assert len(cancel_calls) == 2 + assert all("scancel" in c for c in cancel_calls) +``` + +- [ ] **Step 2: Run tests to verify they fail** + +```bash +pytest tests/distributed/test_provisioner.py -v 2>&1 | head -20 +``` + +Expected: `ImportError` for `miss_alignment.distributed.provisioner`. + +- [ ] **Step 3: Implement `provisioner.py`** + +```python +# src/miss_alignment/distributed/provisioner.py +"""Worker provisioners: spawn local child processes or submit cluster jobs.""" + +from __future__ import annotations + +import re +import subprocess +import sys +from abc import ABC, abstractmethod +from pathlib import Path +from string import Template + +from .config import ClusterConfig + + +class WorkerProvisioner(ABC): + @abstractmethod + def ensure_workers(self, n_tasks: int) -> None: + """Ensure sufficient workers are running for n_tasks tasks.""" + + @abstractmethod + def shutdown(self) -> None: + """Terminate all managed workers.""" + + +class LocalProvisioner(WorkerProvisioner): + """Spawns miss-alignment worker child processes, one 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] = {} # device -> process + + def ensure_workers(self, n_tasks: int) -> None: + 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() + + +class ClusterProvisioner(WorkerProvisioner): + """Submits one cluster job per task via a configurable submit command.""" + + def __init__(self, queue_dir: Path, config: ClusterConfig) -> None: + self._queue_dir = queue_dir + self._config = config + self._job_ids: list[str] = [] + self._scripts_dir = queue_dir / "cluster" + self._scripts_dir.mkdir(parents=True, exist_ok=True) + + def _render_script(self, index: int) -> Path: + """Render the .sh template for one worker and write it to tasks/cluster/.""" + template_text = self._config.script_path.read_text() + # The {{command}} in the template uses shell-evaluated $(hostname) and $$ + # so worker IDs are unique per compute node at runtime. + command = ( + f"miss-alignment worker" + f" --queue-dir {self._queue_dir}" + f" --device 0" + f' --worker-id "$(hostname)-$$-{index}"' + ) + # Replace {{command}} and any {{MISS_CLUSTER_VAR_*}} env var placeholders. + import os + + rendered = template_text.replace("{{command}}", command) + 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 ensure_workers(self, n_tasks: int) -> None: + already = len(self._job_ids) + for i in range(already, n_tasks): + script_path = self._render_script(i) + submit_cmd = self._config.submit.replace( + "{{script_path}}", str(script_path) + ) + result = subprocess.run( + submit_cmd, + shell=True, + capture_output=False, + stdout=subprocess.PIPE, + stderr=sys.stderr, + text=True, + ) + match = re.search(self._config.submit_job_id_regex, result.stdout) + if match: + self._job_ids.append(match.group(1)) + + def shutdown(self) -> None: + for job_id in self._job_ids: + cancel_cmd = self._config.cancel.replace("{{job_id}}", job_id) + subprocess.run(cancel_cmd, shell=True, stderr=subprocess.DEVNULL) + self._job_ids.clear() +``` + +- [ ] **Step 4: Run tests to verify they pass** + +```bash +pytest tests/distributed/test_provisioner.py -v +``` + +Expected: all 5 tests PASS. + +- [ ] **Step 5: Commit** + +```bash +git add src/miss_alignment/distributed/provisioner.py tests/distributed/test_provisioner.py +git commit -m "feat: add LocalProvisioner and ClusterProvisioner" +``` + +--- + +## Task 5: Manager (`distributed/manager.py`) + +**Files:** +- Create: `src/miss_alignment/distributed/manager.py` +- Create: `tests/distributed/test_manager.py` + +**Interfaces:** +- Consumes: + - `QueueLayout`, `TaskSpec`, `write_pending`, `clear_queue`, `compute_fingerprint` from `distributed/queue.py` + - `WorkerProvisioner` from `distributed/provisioner.py` +- Produces: + - `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], queue_root: Path, cluster_config: ClusterConfig | None) -> dict[str, float]` + +- [ ] **Step 1: Write failing tests** + +```python +# tests/distributed/test_manager.py +"""Integration tests for the manager coordinator.""" +import json +import threading +import time +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from miss_alignment.distributed.manager import run_distributed +from miss_alignment.distributed.queue import QueueLayout, mark_done, mark_failed + + +def _fake_provisioner_class(layout): + """Returns a provisioner that writes done files for each pending task.""" + + class FakeProvisioner: + def __init__(self, *a, **kw): + pass + + def ensure_workers(self, n_tasks): + pass + + def shutdown(self): + pass + + return FakeProvisioner + + +def _make_xml(tmp_path, name): + p = tmp_path / f"{name}.xml" + p.write_text(f"{name}") + return p + + +def test_run_distributed_returns_losses(tmp_path): + """Manager resolves all tasks completed by a simulated worker thread.""" + 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" + + # Simulate a worker: poll pending/ and write done/ files + def fake_worker(layout_root): + layout = QueueLayout(layout_root) + deadline = time.time() + 10 + done_count = 0 + while done_count < 2 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) + import os + try: + os.rename(f, running_dir / f.name) + except FileNotFoundError: + continue + 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) + + worker_thread = threading.Thread(target=fake_worker, args=(queue_root,), daemon=True) + worker_thread.start() + + with patch( + "miss_alignment.distributed.manager.LocalProvisioner", + _fake_provisioner_class(None), + ): + 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], + queue_root=queue_root, + cluster_config=None, + ) + + assert set(losses.keys()) == {"ts01", "ts02"} + assert all(v == pytest.approx(0.01) for v in losses.values()) + + +def test_run_distributed_raises_if_any_task_fails(tmp_path): + """Manager raises RuntimeError if any series ends in failed/.""" + xml1 = _make_xml(tmp_path, "ts01") + ckpt = tmp_path / "model.ckpt" + ckpt.write_text("") + + queue_root = tmp_path / "tasks" + + def fake_failing_worker(layout_root): + layout = QueueLayout(layout_root) + deadline = time.time() + 10 + while 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) + import os + try: + os.rename(f, running_dir / f.name) + except FileNotFoundError: + continue + fail_data = {**data, "error": "boom", "worker_id": "fake-worker"} + (layout.failed / f"{task_id}.json").write_text( + json.dumps(fail_data) + ) + (running_dir / f"{task_id}.json").unlink(missing_ok=True) + return + time.sleep(0.05) + + worker_thread = threading.Thread( + target=fake_failing_worker, args=(queue_root,), daemon=True + ) + worker_thread.start() + + with patch( + "miss_alignment.distributed.manager.LocalProvisioner", + _fake_provisioner_class(None), + ): + 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], + queue_root=queue_root, + cluster_config=None, + ) +``` + +- [ ] **Step 2: Run tests to verify they fail** + +```bash +pytest tests/distributed/test_manager.py -v 2>&1 | head -20 +``` + +Expected: `ImportError` for `miss_alignment.distributed.manager`. + +- [ ] **Step 3: Implement `manager.py`** + +```python +# src/miss_alignment/distributed/manager.py +"""Head-node coordinator: writes tasks, starts provisioner, blocks until done.""" + +from __future__ import annotations + +import json +import sys +import threading +import time +from pathlib import Path + +import tqdm + +from .config import ClusterConfig +from .provisioner import ClusterProvisioner, 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: + old_hb = layout.manager_hb / f"hb-{seq - 1}" + old_hb.unlink(missing_ok=True) + + +def _sweep_stalled_workers(layout: QueueLayout) -> None: + """Move tasks from stalled worker dirs back to pending/.""" + for worker_dir in layout.running.iterdir(): + if not worker_dir.is_dir(): + continue + ticks = list(worker_dir.glob("hb-*")) + if ticks: + latest = max(ticks, key=lambda p: p.stat().st_mtime) + age = time.time() - latest.stat().st_mtime + else: + # No heartbeat yet — use dir mtime as proxy + age = time.time() - worker_dir.stat().st_mtime + + if age <= _WORKER_STALL_TIMEOUT_S: + continue + + # Worker is stalled — recover its tasks + for task_file in worker_dir.glob("*.json"): + dest = layout.pending / task_file.name + try: + import os + os.rename(task_file, dest) + print( + f"[manager] Recovered stalled task {task_file.name} to pending", + file=sys.stderr, + ) + except FileNotFoundError: + pass + # Clean up heartbeat files + for hb in worker_dir.glob("hb-*"): + hb.unlink(missing_ok=True) + try: + worker_dir.rmdir() + except OSError: + pass + + +def _scheduler_thread( + layout: QueueLayout, + provisioner: WorkerProvisioner, + n_tasks: int, + stop_event: threading.Event, +) -> None: + hb_seq = 0 + last_hb = 0.0 + last_sweep = 0.0 + + 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) + provisioner.ensure_workers(n_tasks) + last_sweep = now + + stop_event.wait(timeout=1.0) + + +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], + queue_root: Path, + cluster_config: ClusterConfig | None, +) -> dict[str, float]: + """Write tasks, provision workers, block until all tasks are terminal. + + Returns a dict mapping tilt-series name to final loss. + Raises RuntimeError listing all failed series if any task ends in failed/. + """ + layout = QueueLayout(queue_root) + layout.ensure_directories() + clear_queue(layout) + + # Fingerprint is the same for all tasks in one alignment phase. + 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, + ) + + 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), + 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, + ) + write_pending(layout, spec) + + n_tasks = len(tilt_series_list) + if cluster_config is not None: + provisioner: WorkerProvisioner = ClusterProvisioner( + queue_dir=queue_root, config=cluster_config + ) + else: + provisioner = LocalProvisioner(queue_dir=queue_root, devices=devices) + + stop_event = threading.Event() + scheduler = threading.Thread( + target=_scheduler_thread, + args=(layout, provisioner, n_tasks, stop_event), + daemon=True, + ) + scheduler.start() + provisioner.ensure_workers(n_tasks) + + pending_ids = set(task_ids) + losses: dict[str, float] = {} + failed_series: list[str] = [] + + pbar = tqdm.tqdm(total=n_tasks, desc="Tilt series alignment", file=sys.stdout) + try: + while pending_ids: + time.sleep(_POLL_INTERVAL_S) + + 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() + + if failed_series: + raise RuntimeError( + f"Alignment failed for {len(failed_series)} tilt series: " + + ", ".join(failed_series) + ) + + return losses +``` + +- [ ] **Step 4: Run tests to verify they pass** + +```bash +pytest tests/distributed/test_manager.py -v +``` + +Expected: both tests PASS. + +- [ ] **Step 5: Run all distributed tests together** + +```bash +pytest tests/distributed/ -v +``` + +Expected: all tests PASS. + +- [ ] **Step 6: Commit** + +```bash +git add src/miss_alignment/distributed/manager.py tests/distributed/test_manager.py +git commit -m "feat: add distributed manager with scheduler thread and poll loop" +``` + +--- + +## Task 6: Wire up `alignment/parallel.py` and CLI; delete `_parallel.py` + +**Files:** +- Modify: `src/miss_alignment/alignment/parallel.py` +- Modify: `src/miss_alignment/_cli.py` +- Modify: `src/miss_alignment/__init__.py` +- Modify: `tests/test_parallel.py` +- Delete: `src/miss_alignment/_parallel.py` + +**Interfaces:** +- Consumes: + - `run_distributed` from `distributed/manager.py` + - `load_cluster_config` from `distributed/config.py` + - `worker_miss_align` from `distributed/worker.py` +- Produces: `run_alignment_parallel` — same signature as before, same return type `dict[str, float]`. + +- [ ] **Step 1: Update `tests/test_parallel.py` to use the new import** + +The existing `test_parallel.py` tests `_parallel.run_device_pool` directly. Since `_parallel.py` is being deleted, update the tests to verify the equivalent behaviour through `LocalProvisioner` instead. Replace the file entirely: + +```python +# tests/test_parallel.py +"""Tests for the distributed worker provisioner (replaces _parallel.py tests).""" +import subprocess +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from miss_alignment.distributed.provisioner import LocalProvisioner + + +def test_local_provisioner_spawns_worker_per_device(tmp_path): + """LocalProvisioner starts one worker process per GPU device.""" + 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_tasks=10) + + assert mock_popen.call_count == 3 + + +def test_local_provisioner_does_not_double_spawn(tmp_path): + """Calling ensure_workers twice does not spawn extra processes.""" + 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_tasks=5) + p.ensure_workers(n_tasks=5) + + assert mock_popen.call_count == 1 + + +@pytest.mark.filterwarnings("ignore") +def test_local_provisioner_shutdown_terminates(tmp_path): + """shutdown() terminates all spawned processes.""" + 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_tasks=3) + p.shutdown() + + mock_proc.terminate.assert_called() +``` + +- [ ] **Step 2: Run updated tests to verify they pass before touching sources** + +```bash +pytest tests/test_parallel.py -v +``` + +Expected: all 3 tests PASS (they import from `distributed.provisioner` which already exists). + +- [ ] **Step 3: Update `alignment/parallel.py`** + +Replace the file entirely: + +```python +# src/miss_alignment/alignment/parallel.py +from pathlib import Path + +from ..distributed.config import load_cluster_config +from ..distributed.manager import run_distributed + + +def run_alignment_parallel( + model_checkpoint: Path, + tilt_series_list: list[Path], + output_directory: Path, + setting: str | tuple[int, int] | tuple[int, int, int, int], + patch_size: int, + patch_overlap: float, + batch_size: int, + apply_ctf: bool, + downsample: int, + devices_list: list[int], +) -> dict[str, float]: + """Distribute per-tilt-series alignment across local GPUs or a cluster. + + With no cluster env vars set, workers are spawned as local child processes + (one per GPU in devices_list). Set MISS_CLUSTER_CONFIG and MISS_CLUSTER_SCRIPT + to fan work out to a batch scheduler instead. + + Returns a dict mapping tilt-series stem names to their final loss values. + """ + cluster_config = load_cluster_config() + # output_directory is the training directory in train.py and data_directory in + # infer.py — both are the top-level data dir, so tasks/ lives alongside the XMLs. + 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, + queue_root=queue_root, + cluster_config=cluster_config, + ) +``` + +- [ ] **Step 4: Register `worker` subcommand in `_cli.py`** + +```python +# src/miss_alignment/_cli.py +from click import Context +import typer +from typer.core import TyperGroup + + +class OrderCommands(TyperGroup): + def list_commands(self, ctx: Context): + """Return list of commands in the order appear.""" + return list(self.commands) # get commands using self.commands + + +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) +``` + +- [ ] **Step 5: Export `worker_miss_align` from `__init__.py`** + +```python +# src/miss_alignment/__init__.py +"""She has a chaotic good alignment for tilt-series.""" + +from importlib.metadata import PackageNotFoundError, version + +try: + __version__ = version("miss_alignment") +except PackageNotFoundError: + __version__ = "uninstalled" + +__author__ = "Marten Chaillet" +__email__ = "martenchaillet@gmail.com" +__all__ = [ + "__version__", + "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 +``` + +- [ ] **Step 6: Delete `_parallel.py`** + +```bash +git rm src/miss_alignment/_parallel.py +``` + +- [ ] **Step 7: Run the full test suite** + +```bash +pytest --color=yes -v +``` + +Expected: all tests PASS, no warnings-as-errors regressions. If `test_infer.py` or `test_train.py` import `_parallel` directly, fix those imports to remove them (the module no longer exists). + +- [ ] **Step 8: Run linter** + +```bash +ruff check --fix src/miss_alignment/ +ruff format src/miss_alignment/ +``` + +Expected: no errors. + +- [ ] **Step 9: Commit** + +```bash +git add src/miss_alignment/alignment/parallel.py src/miss_alignment/_cli.py src/miss_alignment/__init__.py tests/test_parallel.py +git commit -m "feat: wire distributed queue into run_alignment_parallel; add worker subcommand; delete _parallel.py" +``` + +--- + +## Task 7: Update `distributed/__init__.py` and run full suite + +**Files:** +- Modify: `src/miss_alignment/distributed/__init__.py` + +**Interfaces:** +- Produces: public re-exports for any consumer that imports directly from `miss_alignment.distributed`. + +- [ ] **Step 1: Update `__init__.py`** + +```python +# src/miss_alignment/distributed/__init__.py +"""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, 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", + "LocalProvisioner", + "WorkerProvisioner", + "QueueLayout", + "TaskSpec", + "claim_one", + "clear_queue", + "compute_fingerprint", + "mark_done", + "mark_failed", + "write_pending", + "run_worker_loop", + "worker_miss_align", +] +``` + +- [ ] **Step 2: Run full test suite and linter** + +```bash +pytest --color=yes --cov --cov-report=term-missing +ruff check src/miss_alignment/ +ruff format --check src/miss_alignment/ +``` + +Expected: all tests PASS, coverage report shows `distributed/` coverage, no ruff errors. + +- [ ] **Step 3: Verify `miss-alignment worker --help` works** + +```bash +miss-alignment worker --help +``` + +Expected output includes `--queue-dir`, `--device`, `--worker-id` options. + +- [ ] **Step 4: Final commit** + +```bash +git add src/miss_alignment/distributed/__init__.py +git commit -m "feat: export distributed public API from __init__.py" +``` From 105714c29e4d0ae08b1d85b58a6124747e2bb90e Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 17:58:06 -0700 Subject: [PATCH 04/33] =?UTF-8?q?docs:=20update=20spec=20and=20plan=20?= =?UTF-8?q?=E2=80=94=20n-cluster-workers=20trigger,=20multi-task=20workers?= =?UTF-8?q?,=20model=20param?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Cluster mode triggered by --n-cluster-workers, not env vars alone - Env vars required (not optional) when cluster mode is active - ClusterProvisioner submits N jobs, each draining the queue - evaluate_tilt_series gains optional model param for checkpoint reuse - clear_queue bug fixed (delete before recover) - manager heartbeat written before workers start (race fix) - tasks/ deleted on manager shutdown - __main__.py added for python -m miss_alignment launch Co-Authored-By: Claude Fable 5 --- .../plans/2026-07-04-distributed-inference.md | 935 +++++++++++------- ...2026-07-04-distributed-inference-design.md | 115 ++- 2 files changed, 645 insertions(+), 405 deletions(-) diff --git a/docs/superpowers/plans/2026-07-04-distributed-inference.md b/docs/superpowers/plans/2026-07-04-distributed-inference.md index 2de3d7a..c20119f 100644 --- a/docs/superpowers/plans/2026-07-04-distributed-inference.md +++ b/docs/superpowers/plans/2026-07-04-distributed-inference.md @@ -4,20 +4,21 @@ **Goal:** Add a disk-based task queue that distributes per-tilt-series alignment inference across cluster nodes (SLURM/PBS), while keeping the existing local multi-GPU behaviour unchanged. -**Architecture:** A new `miss_alignment/distributed/` package implements a filesystem queue (atomic rename as the claim mutex), a manager that writes tasks and blocks until done, two provisioners (local subprocess and cluster batch scheduler), and a `miss-alignment worker` subcommand. `run_alignment_parallel` in `alignment/parallel.py` is updated to drive the manager instead of `_parallel.run_device_pool`. No changes to `train.py`, `infer.py`, or `evaluate_tilt_series`. +**Architecture:** A new `miss_alignment/distributed/` package implements a filesystem queue (atomic rename as the claim mutex), a manager that writes tasks and blocks until done, two provisioners (local subprocess and cluster batch scheduler), and a `miss-alignment worker` subcommand. Each worker processes many series per run; checkpoint loading is amortized by passing the resident model into `evaluate_tilt_series`. Cluster mode is triggered by a new `--n-cluster-workers` CLI arg; without it local multi-GPU mode is unchanged. -**Tech Stack:** Python 3.10+, stdlib only (`pathlib`, `os`, `subprocess`, `threading`, `hashlib`, `json`, `re`, `time`, `signal`) — no new dependencies. +**Tech Stack:** Python 3.10+, stdlib only for `distributed/` (`pathlib`, `os`, `subprocess`, `threading`, `hashlib`, `json`, `re`, `time`, `signal`). `tqdm` (already a dependency) used in manager. `typer` (already a dependency) used in worker CLI. ## Global Constraints - Python ≥ 3.10 (project minimum). -- No new runtime dependencies beyond stdlib. +- No new runtime dependencies beyond stdlib + existing deps. - `ruff check --fix && ruff format` must pass (line length 88, ignore E712). - `pytest --color=yes` must pass with no warnings-as-errors regressions. -- All new files under `src/miss_alignment/distributed/`. -- Cluster mode activated only when **both** `MISS_CLUSTER_CONFIG` and `MISS_CLUSTER_SCRIPT` env vars are set; otherwise `LocalProvisioner` is used. -- `evaluate_tilt_series` is never modified. -- `train.py` and `infer.py` are never modified. +- All new queue infrastructure under `src/miss_alignment/distributed/`. +- Cluster mode activated only by `--n-cluster-workers N`; without it local mode runs unchanged. +- `MISS_CLUSTER_CONFIG` and `MISS_CLUSTER_SCRIPT` are **required** when `--n-cluster-workers` is set; `config.py` raises `RuntimeError` if either is absent (never silently falls back to local mode). +- `evaluate_tilt_series` gains one optional `model` parameter and is otherwise unchanged. +- `train.py` and `infer.py` are modified only to add the `--n-cluster-workers` option and pass it through. --- @@ -25,27 +26,34 @@ | Path | Action | Responsibility | |---|---|---| -| `src/miss_alignment/distributed/__init__.py` | Create | Re-export public API | +| `src/miss_alignment/__main__.py` | Create | `python -m miss_alignment` entry point for `LocalProvisioner` subprocess launch | +| `src/miss_alignment/distributed/__init__.py` | Create | Public re-exports | | `src/miss_alignment/distributed/queue.py` | Create | Directory layout, task JSON, atomic rename claim | | `src/miss_alignment/distributed/manager.py` | Create | Head-node coordinator, scheduler thread, poll loop | | `src/miss_alignment/distributed/provisioner.py` | Create | `WorkerProvisioner` ABC, `LocalProvisioner`, `ClusterProvisioner` | | `src/miss_alignment/distributed/worker.py` | Create | `miss-alignment worker` subcommand logic | -| `src/miss_alignment/distributed/config.py` | Create | Read env vars, return `ClusterConfig \| None` | -| `src/miss_alignment/alignment/parallel.py` | Modify | Replace `run_device_pool` call with manager call | +| `src/miss_alignment/distributed/config.py` | Create | Read env vars, return `ClusterConfig` or raise | +| `src/miss_alignment/alignment/tilt_series.py` | Modify | Add optional `model` parameter to `evaluate_tilt_series` | +| `src/miss_alignment/alignment/parallel.py` | Modify | Replace `run_device_pool` call with manager; add `n_cluster_workers` param | +| `src/miss_alignment/train.py` | Modify | Add `--n-cluster-workers` option; pass to `run_alignment_parallel` | +| `src/miss_alignment/infer.py` | Modify | Add `--n-cluster-workers` option; pass to `run_alignment_parallel` | | `src/miss_alignment/_cli.py` | Modify | Register `worker` subcommand | | `src/miss_alignment/__init__.py` | Modify | Export `worker_miss_align` | -| `src/miss_alignment/_parallel.py` | Delete (after Task 6) | Superseded by `LocalProvisioner` | +| `src/miss_alignment/_parallel.py` | Delete (Task 7) | Superseded by `LocalProvisioner` | +| `tests/distributed/__init__.py` | Create | Test package | | `tests/distributed/test_queue.py` | Create | Queue layer unit tests | -| `tests/distributed/test_manager.py` | Create | Manager + provisioner integration tests | +| `tests/distributed/test_manager.py` | Create | Manager integration tests | | `tests/distributed/test_worker.py` | Create | Worker claim loop unit tests | | `tests/distributed/test_config.py` | Create | Config env-var parsing tests | -| `tests/test_parallel.py` | Modify | Update import from new path | +| `tests/distributed/test_provisioner.py` | Create | Provisioner unit tests | +| `tests/test_parallel.py` | Modify | Update to test new `LocalProvisioner` path | --- ## Task 1: Queue layer (`distributed/queue.py`) **Files:** +- Create: `src/miss_alignment/__main__.py` - Create: `src/miss_alignment/distributed/__init__.py` - Create: `src/miss_alignment/distributed/queue.py` - Create: `tests/distributed/__init__.py` @@ -53,22 +61,24 @@ **Interfaces:** - Produces: - - `QueueLayout(root: Path)` — manages subdirectory creation; attributes `pending`, `running`, `done`, `failed`, `manager_hb`, each a `Path`. - - `TaskSpec` — `dataclass` with fields: `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`. - - `write_pending(layout: QueueLayout, spec: TaskSpec) -> None` — writes `layout.pending/.json`. - - `claim_one(layout: QueueLayout, worker_id: str) -> TaskSpec | None` — shuffled rename claim; returns `None` when queue empty. + - `QueueLayout(root: Path)` — dataclass; `ensure_directories() -> None`; properties: `pending`, `running`, `done`, `failed`, `manager_hb`, `cluster` each returning `Path`; `worker_dir(worker_id: str) -> Path`. + - `TaskSpec` — dataclass with fields: `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`. + - `compute_fingerprint(model_checkpoint_path, setting, patch_size, patch_overlap, batch_size, apply_ctf, downsample) -> str` — SHA-256 hex. + - `write_pending(layout: QueueLayout, spec: TaskSpec) -> None` + - `claim_one(layout: QueueLayout, worker_id: str) -> TaskSpec | None` - `mark_done(layout: QueueLayout, worker_id: str, spec: TaskSpec, final_loss: float, device: str) -> None` - `mark_failed(layout: QueueLayout, worker_id: str, spec: TaskSpec, error: str) -> None` - - `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 hex. - - `clear_queue(layout: QueueLayout) -> None` — deletes all JSON files in `pending/`, `done/`, `failed/`; moves any `running//.json` back to `pending/` (orphan recovery). + - `clear_queue(layout: QueueLayout) -> None` — deletes `pending/done/failed` contents first, then recovers `running/` orphans into the now-empty `pending/`. -- [ ] **Step 1: Create directory skeleton and write failing tests** +- [ ] **Step 1: Create directory skeleton** ```bash mkdir -p tests/distributed touch tests/distributed/__init__.py ``` +- [ ] **Step 2: Write failing tests** + ```python # tests/distributed/test_queue.py import json @@ -128,13 +138,12 @@ def test_claim_one_returns_none_when_empty(layout): assert claim_one(layout, "worker-0") is None -def test_claim_one_exclusive(layout, tmp_path): - """Two concurrent claimers: exactly one wins.""" +def test_claim_one_exclusive(layout): + """Two sequential claimers: exactly one wins.""" write_pending(layout, _spec()) - results = [] - results.append(claim_one(layout, "worker-0")) - results.append(claim_one(layout, "worker-1")) - claimed = [r for r in results if r is not None] + 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 @@ -167,12 +176,27 @@ def test_clear_queue_recovers_orphans(layout): spec = _spec() write_pending(layout, spec) claim_one(layout, "worker-0") - # simulate crash: running file remains, we clear and recover + # simulate crash: running file remains; clear should put it back in pending clear_queue(layout) - # orphan should be back in pending 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) @@ -186,23 +210,30 @@ def test_compute_fingerprint_differs_on_change(): assert fp1 != fp2 ``` -- [ ] **Step 2: Run tests to verify they fail** +- [ ] **Step 3: Run tests to verify they fail** ```bash cd /Users/tegunovd/dev/miss-alignment -pytest tests/distributed/test_queue.py -v 2>&1 | head -30 +pytest tests/distributed/test_queue.py -v 2>&1 | head -20 ``` Expected: `ModuleNotFoundError: No module named 'miss_alignment.distributed'` -- [ ] **Step 3: Create `__init__.py` skeleton** +- [ ] **Step 4: Create `__main__.py` and `distributed/__init__.py` skeleton** + +```python +# src/miss_alignment/__main__.py +from miss_alignment import cli + +cli() +``` ```python # src/miss_alignment/distributed/__init__.py """Disk-based distributed task queue for miss-alignment inference.""" ``` -- [ ] **Step 4: Implement `queue.py`** +- [ ] **Step 5: Implement `queue.py`** ```python # src/miss_alignment/distributed/queue.py @@ -361,8 +392,7 @@ def mark_done( data["final_loss"] = final_loss data["device"] = device _atomic_write(layout.done / f"{spec.task_id}.json", data) - running_path = layout.worker_dir(worker_id) / f"{spec.task_id}.json" - running_path.unlink(missing_ok=True) + (layout.worker_dir(worker_id) / f"{spec.task_id}.json").unlink(missing_ok=True) def mark_failed( @@ -376,13 +406,19 @@ def mark_failed( data["error"] = error data["worker_id"] = worker_id _atomic_write(layout.failed / f"{spec.task_id}.json", data) - running_path = layout.worker_dir(worker_id) / f"{spec.task_id}.json" - running_path.unlink(missing_ok=True) + (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.""" - # recover orphaned running tasks back to pending + """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 @@ -392,28 +428,30 @@ def clear_queue(layout: QueueLayout) -> None: 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; will be swept later - - for directory in (layout.pending, layout.done, layout.failed): - for f in directory.glob("*.json"): - f.unlink(missing_ok=True) + pass # not empty yet; scheduler will sweep it ``` -- [ ] **Step 5: Run tests to verify they pass** +- [ ] **Step 6: Run tests to verify they pass** ```bash pytest tests/distributed/test_queue.py -v ``` -Expected: all 9 tests PASS. +Expected: all 10 tests PASS. -- [ ] **Step 6: Commit** +- [ ] **Step 7: Commit** ```bash -git add src/miss_alignment/distributed/__init__.py src/miss_alignment/distributed/queue.py tests/distributed/__init__.py tests/distributed/test_queue.py +git add src/miss_alignment/__main__.py \ + src/miss_alignment/distributed/__init__.py \ + src/miss_alignment/distributed/queue.py \ + tests/distributed/__init__.py \ + tests/distributed/test_queue.py git commit -m "feat: add distributed queue layer with atomic rename claim protocol" ``` @@ -426,16 +464,14 @@ git commit -m "feat: add distributed queue layer with atomic rename claim protoc - Create: `tests/distributed/test_config.py` **Interfaces:** -- Consumes: nothing. - Produces: - - `ClusterConfig` — `dataclass` with fields: `submit: str`, `submit_job_id_regex: str`, `cancel: str`, `script_path: Path`. - - `load_cluster_config() -> ClusterConfig | None` — reads `MISS_CLUSTER_CONFIG` and `MISS_CLUSTER_SCRIPT`; returns `None` if either is unset. + - `ClusterConfig` — dataclass with fields: `submit: str`, `submit_job_id_regex: str`, `cancel: str`, `script_path: Path`. + - `load_cluster_config() -> ClusterConfig` — reads `MISS_CLUSTER_CONFIG` and `MISS_CLUSTER_SCRIPT`; raises `RuntimeError` if either is unset, `FileNotFoundError` if a path doesn't exist, `KeyError` if a required JSON key is missing. - [ ] **Step 1: Write failing tests** ```python # tests/distributed/test_config.py -import os import json import pytest from pathlib import Path @@ -461,16 +497,18 @@ def cluster_script(tmp_path): return p -def test_load_cluster_config_returns_none_when_unset(monkeypatch): +def test_load_cluster_config_raises_when_config_unset(monkeypatch): monkeypatch.delenv("MISS_CLUSTER_CONFIG", raising=False) monkeypatch.delenv("MISS_CLUSTER_SCRIPT", raising=False) - assert load_cluster_config() is None + with pytest.raises(RuntimeError, match="MISS_CLUSTER_CONFIG"): + load_cluster_config() -def test_load_cluster_config_returns_none_when_only_one_set(monkeypatch, cluster_json): +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) - assert load_cluster_config() is None + with pytest.raises(RuntimeError, match="MISS_CLUSTER_SCRIPT"): + load_cluster_config() def test_load_cluster_config_returns_config(monkeypatch, cluster_json, cluster_script): @@ -505,7 +543,7 @@ def test_load_cluster_config_raises_on_missing_key(monkeypatch, tmp_path, cluste pytest tests/distributed/test_config.py -v 2>&1 | head -20 ``` -Expected: `ImportError` or `ModuleNotFoundError`. +Expected: `ImportError` for `miss_alignment.distributed.config`. - [ ] **Step 3: Implement `config.py`** @@ -513,9 +551,9 @@ Expected: `ImportError` or `ModuleNotFoundError`. # src/miss_alignment/distributed/config.py """Read cluster configuration from environment variables. -Cluster mode is activated when both MISS_CLUSTER_CONFIG and MISS_CLUSTER_SCRIPT -are set. If either is absent, load_cluster_config() returns None and -LocalProvisioner is used instead. +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 @@ -534,13 +572,23 @@ class ClusterConfig: script_path: Path -def load_cluster_config() -> ClusterConfig | None: - """Return ClusterConfig if both env vars are set, else None.""" +def load_cluster_config() -> ClusterConfig: + """Return ClusterConfig. Raises RuntimeError if env vars are missing.""" config_path_str = os.environ.get("MISS_CLUSTER_CONFIG") - script_path_str = os.environ.get("MISS_CLUSTER_SCRIPT") + 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', and 'cancel' keys." + ) - if not config_path_str or not script_path_str: - return None + 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) @@ -571,7 +619,7 @@ Expected: all 5 tests PASS. ```bash git add src/miss_alignment/distributed/config.py tests/distributed/test_config.py -git commit -m "feat: add cluster config reader (MISS_CLUSTER_CONFIG + MISS_CLUSTER_SCRIPT)" +git commit -m "feat: add cluster config reader (raises if env vars missing)" ``` --- @@ -585,9 +633,11 @@ git commit -m "feat: add cluster config reader (MISS_CLUSTER_CONFIG + MISS_CLUST **Interfaces:** - Consumes: - `QueueLayout`, `TaskSpec`, `claim_one`, `mark_done`, `mark_failed` from `distributed/queue.py` + - `MissAlignment` from `..models.models` + - `evaluate_tilt_series` from `..alignment.tilt_series` (after Task 4 adds the `model` param) - Produces: - - `worker_miss_align(queue_dir: Path, device: int, worker_id: str | None)` — Typer command entry point; loops until queue empty or manager heartbeat stale. - - `run_worker_loop(layout: QueueLayout, worker_id: str, device: str, manager_hb_timeout_s: float) -> None` — testable loop body. + - `run_worker_loop(layout, worker_id, device, manager_hb_timeout_s) -> None` — testable loop body. + - `worker_miss_align(queue_dir, device, worker_id)` — Typer command. - [ ] **Step 1: Write failing tests** @@ -595,21 +645,18 @@ git commit -m "feat: add cluster config reader (MISS_CLUSTER_CONFIG + MISS_CLUST # tests/distributed/test_worker.py """Unit tests for the worker claim loop. -These tests do NOT call evaluate_tilt_series; they mock it to keep -tests fast and free of CUDA/warpylib dependencies. +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 patch +from unittest.mock import MagicMock, patch import pytest -from miss_alignment.distributed.queue import ( - QueueLayout, - TaskSpec, - write_pending, -) +from miss_alignment.distributed.queue import QueueLayout, TaskSpec, write_pending from miss_alignment.distributed.worker import run_worker_loop @@ -621,17 +668,16 @@ def layout(tmp_path): def _write_manager_hb(layout, seq=0): - """Write a fresh manager heartbeat tick.""" 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"): +def _spec(task_id="0000001-ts01", fingerprint="abc123"): return TaskSpec( task_id=task_id, model_checkpoint_path="/data/model.ckpt", - tilt_series_path="/data/ts01.xml", + tilt_series_path=f"/data/{task_id}.xml", output_directory="/data/out", setting="anchoring", patch_size=96, @@ -639,18 +685,20 @@ def _spec(task_id="0000001-ts01"): batch_size=32, apply_ctf=False, downsample=2, - init_fingerprint="abc123", + init_fingerprint=fingerprint, ) -def test_worker_processes_task_and_writes_done(layout, tmp_path): +def test_worker_processes_task_and_writes_done(layout): _write_manager_hb(layout) write_pending(layout, _spec()) - fake_loss = [0.5, 0.3, 0.1] with patch( "miss_alignment.distributed.worker.evaluate_tilt_series", - return_value=(Path("/data/ts01.xml"), fake_loss), + 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) @@ -667,6 +715,9 @@ def test_worker_writes_failed_on_exception(layout): 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) @@ -677,12 +728,10 @@ def test_worker_writes_failed_on_exception(layout): def test_worker_exits_when_manager_hb_stale(layout): - # Write a manager heartbeat that is already old + # Write a heartbeat file that is 200 seconds old hb_file = layout.manager_hb / "hb-0" hb_file.write_text("") - # Make it appear 200 seconds old by back-dating mtime old_time = time.time() - 200 - import os os.utime(hb_file, (old_time, old_time)) write_pending(layout, _spec()) @@ -691,49 +740,67 @@ def test_worker_exits_when_manager_hb_stale(layout): 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) - # Worker should exit without processing the task 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) - spec1 = _spec("0000001-ts01") - spec2 = TaskSpec( - task_id="0000002-ts02", - model_checkpoint_path="/data/model.ckpt", - tilt_series_path="/data/ts02.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", # same fingerprint - ) - write_pending(layout, spec1) - write_pending(layout, spec2) + 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, ): - with patch( - "miss_alignment.distributed.worker.MissAlignment.load_from_checkpoint", - ) as mock_load: - mock_load.return_value = mock_load # return self as stub - run_worker_loop(layout, "worker-0", "cpu", manager_hb_timeout_s=30.0) - # Model should only be loaded once despite two tasks - assert mock_load.call_count == 1 + run_worker_loop(layout, "worker-0", "cpu", manager_hb_timeout_s=30.0) + + assert len(load_calls) == 2 ``` - [ ] **Step 2: Run tests to verify they fail** @@ -750,6 +817,10 @@ Expected: `ImportError` for `miss_alignment.distributed.worker`. # src/miss_alignment/distributed/worker.py """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 ] """ @@ -775,19 +846,15 @@ from .queue import ( mark_failed, ) -# Seconds without a manager heartbeat tick before the worker exits. _MANAGER_HB_TIMEOUT_S = 120.0 -# Seconds between heartbeat writes. _HB_INTERVAL_S = 5.0 def _write_worker_hb(worker_dir: Path, seq: int) -> None: - """Write a new heartbeat tick, removing the previous one.""" new_hb = worker_dir / f"hb-{seq}" new_hb.write_text("") if seq > 0: - old_hb = worker_dir / f"hb-{seq - 1}" - old_hb.unlink(missing_ok=True) + (worker_dir / f"hb-{seq - 1}").unlink(missing_ok=True) def _manager_hb_age_s(layout: QueueLayout) -> float: @@ -799,36 +866,38 @@ def _manager_hb_age_s(layout: QueueLayout) -> float: 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 run_worker_loop( layout: QueueLayout, worker_id: str, device: str, manager_hb_timeout_s: float = _MANAGER_HB_TIMEOUT_S, ) -> None: - """Main worker loop: claim → check heartbeat → evaluate → write result. - - Separated from the Typer command for testability. - """ + """Main worker loop: claim → evaluate → write result. Repeat until queue empty.""" worker_dir = layout.worker_dir(worker_id) worker_dir.mkdir(parents=True, exist_ok=True) last_fingerprint: str | None = None - loaded_model = None + cached_model: MissAlignment | None = None hb_seq = 0 last_hb_time = 0.0 while True: - # Check manager heartbeat before every claim attempt. age = _manager_hb_age_s(layout) if age > manager_hb_timeout_s: print( - f"[{worker_id}] Manager heartbeat stale ({age:.0f}s > " - f"{manager_hb_timeout_s:.0f}s). Exiting.", + f"[{worker_id}] Manager heartbeat stale ({age:.0f}s). Exiting.", file=sys.stderr, ) return - # Write our own heartbeat if due. now = time.time() if now - last_hb_time >= _HB_INTERVAL_S: _write_worker_hb(worker_dir, hb_seq) @@ -841,25 +910,28 @@ def run_worker_loop( print(f"[{worker_id}] Claimed {spec.task_id}", file=sys.stderr) - # Load model only when fingerprint changes. if spec.init_fingerprint != last_fingerprint: - loaded_model = MissAlignment.load_from_checkpoint( - spec.model_checkpoint_path, map_location="cpu" - ) + cached_model = _load_model(spec.model_checkpoint_path) last_fingerprint = spec.init_fingerprint + # Convert setting back to tuple if it was serialized as a list. + setting = ( + tuple(spec.setting) if isinstance(spec.setting, list) else spec.setting + ) + try: _, 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=spec.setting, + 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, ) final_loss = float(loss_values[-1]) if loss_values else float("nan") mark_done(layout, worker_id, spec, final_loss=final_loss, device=device) @@ -891,7 +963,6 @@ def worker_miss_align( layout.ensure_directories() cuda_device = f"cuda:{device}" if torch.cuda.is_available() else "cpu" - run_worker_loop(layout, worker_id, cuda_device) ``` @@ -901,18 +972,99 @@ def worker_miss_align( pytest tests/distributed/test_worker.py -v ``` -Expected: all 4 tests PASS. +Expected: all 5 tests PASS. - [ ] **Step 5: Commit** ```bash git add src/miss_alignment/distributed/worker.py tests/distributed/test_worker.py -git commit -m "feat: add worker subcommand with claim loop and model fingerprint reuse" +git commit -m "feat: add worker subcommand with model fingerprint reuse across tasks" +``` + +--- + +## Task 4: Add `model` parameter to `evaluate_tilt_series` + +**Files:** +- Modify: `src/miss_alignment/alignment/tilt_series.py` + +**Interfaces:** +- Produces: `evaluate_tilt_series(..., model: MissAlignment | None = None)` — when `model` is provided, uses it directly instead of loading from disk. All existing callers unaffected. + +- [ ] **Step 1: Read the current model-loading block** + +Open `src/miss_alignment/alignment/tilt_series.py` at lines 118–186. The relevant section is: + +```python +# line 131 ends the signature +) -> tuple[Path, list[float]]: + ... + # line 172 + model = MissAlignment.load_from_checkpoint( + model_checkpoint_path, + map_location="cpu", + ) + # line 184 + if hasattr(model.net, "_orig_mod"): + model.net = model.net._orig_mod +``` + +- [ ] **Step 2: Add the `model` parameter to the signature** + +In `src/miss_alignment/alignment/tilt_series.py`, add `model` as the last parameter before the closing `)`: + +```python +# Before: + n_control_points: int = 7, +) -> tuple[Path, list[float]]: + +# After: + n_control_points: int = 7, + model: "MissAlignment | None" = None, +) -> tuple[Path, list[float]]: +``` + +Use a string annotation to avoid a circular import (the type is already imported at the top of the file — verify with `grep "MissAlignment" src/miss_alignment/alignment/tilt_series.py` before committing; if already imported, use the bare type). + +- [ ] **Step 3: Replace the model-loading block** + +```python +# Before (lines ~172-185): + model = MissAlignment.load_from_checkpoint( + model_checkpoint_path, + map_location="cpu", + ) + if hasattr(model.net, "_orig_mod"): + model.net = model.net._orig_mod + +# After: + if model is None: + model = MissAlignment.load_from_checkpoint( + model_checkpoint_path, + map_location="cpu", + ) + if hasattr(model.net, "_orig_mod"): + model.net = model.net._orig_mod +``` + +- [ ] **Step 4: Run the existing alignment tests to verify nothing broke** + +```bash +pytest tests/alignment/ -v +``` + +Expected: all tests PASS. + +- [ ] **Step 5: Commit** + +```bash +git add src/miss_alignment/alignment/tilt_series.py +git commit -m "feat: add optional model param to evaluate_tilt_series for checkpoint reuse" ``` --- -## Task 4: Provisioners (`distributed/provisioner.py`) +## Task 5: Provisioners (`distributed/provisioner.py`) **Files:** - Create: `src/miss_alignment/distributed/provisioner.py` @@ -921,19 +1073,17 @@ git commit -m "feat: add worker subcommand with claim loop and model fingerprint **Interfaces:** - Consumes: `ClusterConfig` from `distributed/config.py`. - Produces: - - `WorkerProvisioner` — ABC with `ensure_workers(n_tasks: int) -> None` and `shutdown() -> None`. - - `LocalProvisioner(queue_dir: Path, devices: list[int])` — spawns `miss-alignment worker` child processes. - - `ClusterProvisioner(queue_dir: Path, config: ClusterConfig, n_tasks: int)` — submits cluster jobs. + - `WorkerProvisioner` — ABC with `ensure_workers(n_workers: int) -> None` and `shutdown() -> None`. + - `LocalProvisioner(queue_dir: Path, devices: list[int])` — spawns `python -m miss_alignment worker` child processes, one per device. Ignores `n_workers`. + - `ClusterProvisioner(queue_dir: Path, config: ClusterConfig)` — submits exactly `n_workers` cluster jobs on the first `ensure_workers` call. - [ ] **Step 1: Write failing tests** ```python # tests/distributed/test_provisioner.py -"""Tests for LocalProvisioner and ClusterProvisioner.""" -import subprocess import sys from pathlib import Path -from unittest.mock import MagicMock, call, patch +from unittest.mock import MagicMock, patch import pytest @@ -944,46 +1094,58 @@ from miss_alignment.distributed.provisioner import ClusterProvisioner, LocalProv 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 # still running + mock_proc.poll.return_value = None mock_popen.return_value = mock_proc - p = LocalProvisioner(queue_dir=tmp_path, devices=[0, 1]) - p.ensure_workers(n_tasks=10) + p = LocalProvisioner(queue_dir=tmp_path, devices=[0, 1, 2]) + p.ensure_workers(n_workers=10) - assert mock_popen.call_count == 2 - # Each call should pass --device 0 and --device 1 - calls_str = [str(c) for c in mock_popen.call_args_list] - assert any("--device" in s and "0" in s for s in calls_str) - assert any("--device" in s and "1" in s for s in calls_str) + assert mock_popen.call_count == 3 + # verify device args + all_args = [str(c) for c in mock_popen.call_args_list] + assert any("'0'" in s or '"0"' in s or "0" in s for s in all_args) -def test_local_provisioner_shutdown_terminates_processes(tmp_path): +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_tasks=5) - p.shutdown() + p.ensure_workers(n_workers=5) + p.ensure_workers(n_workers=5) - mock_proc.terminate.assert_called() + assert mock_popen.call_count == 1 -def test_local_provisioner_does_not_respawn_running_processes(tmp_path): +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 = None # still running + 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_tasks=5) - p.ensure_workers(n_tasks=5) # second call should not spawn again + p.ensure_workers(n_workers=5) + p.ensure_workers(n_workers=5) # should respawn because poll() != None - assert mock_popen.call_count == 1 + 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() -def test_cluster_provisioner_submits_one_job_per_task(tmp_path): + mock_proc.terminate.assert_called() + + +def test_cluster_provisioner_submits_n_workers_jobs(tmp_path): script = tmp_path / "worker.sh" script.write_text("#!/bin/bash\n{{command}}\n") cfg = ClusterConfig( @@ -1001,14 +1163,16 @@ def test_cluster_provisioner_submits_one_job_per_task(tmp_path): result.stdout = "Submitted batch job 12345\n" return result - with patch("miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run): + with patch( + "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run + ): p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) - p.ensure_workers(n_tasks=3) + p.ensure_workers(n_workers=4) - assert len(submitted) == 3 + assert len(submitted) == 4 -def test_cluster_provisioner_cancels_jobs_on_shutdown(tmp_path): +def test_cluster_provisioner_cancels_on_shutdown(tmp_path): script = tmp_path / "worker.sh" script.write_text("#!/bin/bash\n{{command}}\n") cfg = ClusterConfig( @@ -1028,12 +1192,14 @@ def test_cluster_provisioner_cancels_jobs_on_shutdown(tmp_path): cancel_calls.append(cmd) return MagicMock() - with patch("miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run): + with patch( + "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run + ): p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) - p.ensure_workers(n_tasks=2) + p.ensure_workers(n_workers=3) p.shutdown() - assert len(cancel_calls) == 2 + assert len(cancel_calls) == 3 assert all("scancel" in c for c in cancel_calls) ``` @@ -1053,20 +1219,20 @@ Expected: `ImportError` for `miss_alignment.distributed.provisioner`. from __future__ import annotations +import os import re import subprocess import sys from abc import ABC, abstractmethod from pathlib import Path -from string import Template from .config import ClusterConfig class WorkerProvisioner(ABC): @abstractmethod - def ensure_workers(self, n_tasks: int) -> None: - """Ensure sufficient workers are running for n_tasks tasks.""" + def ensure_workers(self, n_workers: int) -> None: + """Ensure workers are running. Called once at startup and each scheduler tick.""" @abstractmethod def shutdown(self) -> None: @@ -1074,14 +1240,14 @@ class WorkerProvisioner(ABC): class LocalProvisioner(WorkerProvisioner): - """Spawns miss-alignment worker child processes, one per GPU device.""" + """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] = {} # device -> process + self._procs: dict[int, subprocess.Popen] = {} - def ensure_workers(self, n_tasks: int) -> None: + def ensure_workers(self, n_workers: int) -> None: for device in self._devices: proc = self._procs.get(device) if proc is not None and proc.poll() is None: @@ -1115,7 +1281,7 @@ class LocalProvisioner(WorkerProvisioner): class ClusterProvisioner(WorkerProvisioner): - """Submits one cluster job per task via a configurable submit command.""" + """Submits exactly n_workers cluster jobs, each running until the queue drains.""" def __init__(self, queue_dir: Path, config: ClusterConfig) -> None: self._queue_dir = queue_dir @@ -1125,19 +1291,14 @@ class ClusterProvisioner(WorkerProvisioner): self._scripts_dir.mkdir(parents=True, exist_ok=True) def _render_script(self, index: int) -> Path: - """Render the .sh template for one worker and write it to tasks/cluster/.""" template_text = self._config.script_path.read_text() - # The {{command}} in the template uses shell-evaluated $(hostname) and $$ - # so worker IDs are unique per compute node at runtime. + # $(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}"' ) - # Replace {{command}} and any {{MISS_CLUSTER_VAR_*}} env var placeholders. - import os - rendered = template_text.replace("{{command}}", command) for key, value in os.environ.items(): if key.startswith("MISS_CLUSTER_VAR_"): @@ -1148,9 +1309,9 @@ class ClusterProvisioner(WorkerProvisioner): script_path.write_text(rendered) return script_path - def ensure_workers(self, n_tasks: int) -> None: + def ensure_workers(self, n_workers: int) -> None: already = len(self._job_ids) - for i in range(already, n_tasks): + for i in range(already, n_workers): script_path = self._render_script(i) submit_cmd = self._config.submit.replace( "{{script_path}}", str(script_path) @@ -1158,7 +1319,6 @@ class ClusterProvisioner(WorkerProvisioner): result = subprocess.run( submit_cmd, shell=True, - capture_output=False, stdout=subprocess.PIPE, stderr=sys.stderr, text=True, @@ -1180,7 +1340,7 @@ class ClusterProvisioner(WorkerProvisioner): pytest tests/distributed/test_provisioner.py -v ``` -Expected: all 5 tests PASS. +Expected: all 6 tests PASS. - [ ] **Step 5: Commit** @@ -1191,7 +1351,7 @@ git commit -m "feat: add LocalProvisioner and ClusterProvisioner" --- -## Task 5: Manager (`distributed/manager.py`) +## Task 6: Manager (`distributed/manager.py`) **Files:** - Create: `src/miss_alignment/distributed/manager.py` @@ -1202,39 +1362,28 @@ git commit -m "feat: add LocalProvisioner and ClusterProvisioner" - `QueueLayout`, `TaskSpec`, `write_pending`, `clear_queue`, `compute_fingerprint` from `distributed/queue.py` - `WorkerProvisioner` from `distributed/provisioner.py` - Produces: - - `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], queue_root: Path, cluster_config: ClusterConfig | None) -> dict[str, float]` + - `run_distributed(tilt_series_list, model_checkpoint, output_directory, setting, patch_size, patch_overlap, batch_size, apply_ctf, downsample, devices, n_cluster_workers, queue_root) -> dict[str, float]` - [ ] **Step 1: Write failing tests** ```python # tests/distributed/test_manager.py -"""Integration tests for the manager coordinator.""" +"""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 MagicMock, patch +from unittest.mock import patch import pytest from miss_alignment.distributed.manager import run_distributed -from miss_alignment.distributed.queue import QueueLayout, mark_done, mark_failed - - -def _fake_provisioner_class(layout): - """Returns a provisioner that writes done files for each pending task.""" - - class FakeProvisioner: - def __init__(self, *a, **kw): - pass - - def ensure_workers(self, n_tasks): - pass - - def shutdown(self): - pass - - return FakeProvisioner +from miss_alignment.distributed.queue import QueueLayout def _make_xml(tmp_path, name): @@ -1243,45 +1392,61 @@ def _make_xml(tmp_path, name): return p -def test_run_distributed_returns_losses(tmp_path): - """Manager resolves all tasks completed by a simulated worker thread.""" - 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" +def _fake_worker_thread(queue_root, n_tasks, fail=False): + """Simulates a worker: claims pending tasks, writes done or failed.""" - # Simulate a worker: poll pending/ and write done/ files - def fake_worker(layout_root): - layout = QueueLayout(layout_root) - deadline = time.time() + 10 + def _run(): + layout = QueueLayout(queue_root) done_count = 0 - while done_count < 2 and time.time() < deadline: + 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) - import os try: os.rename(f, running_dir / f.name) except FileNotFoundError: continue - done_data = {**data, "final_loss": 0.01, "device": "cpu"} - (layout.done / f"{task_id}.json").write_text( - json.dumps(done_data) - ) + 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) - worker_thread = threading.Thread(target=fake_worker, args=(queue_root,), daemon=True) - worker_thread.start() + return threading.Thread(target=_run, daemon=True) + + +class _NoOpProvisioner: + def ensure_workers(self, n_workers): + pass + + def shutdown(self): + pass + + +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", - _fake_provisioner_class(None), + return_value=_NoOpProvisioner(), ): losses = run_distributed( tilt_series_list=[xml1, xml2], @@ -1294,52 +1459,26 @@ def test_run_distributed_returns_losses(tmp_path): apply_ctf=False, downsample=2, devices=[0], + n_cluster_workers=None, queue_root=queue_root, - cluster_config=None, ) assert set(losses.keys()) == {"ts01", "ts02"} assert all(v == pytest.approx(0.01) for v in losses.values()) -def test_run_distributed_raises_if_any_task_fails(tmp_path): - """Manager raises RuntimeError if any series ends in failed/.""" +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" - def fake_failing_worker(layout_root): - layout = QueueLayout(layout_root) - deadline = time.time() + 10 - while 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) - import os - try: - os.rename(f, running_dir / f.name) - except FileNotFoundError: - continue - fail_data = {**data, "error": "boom", "worker_id": "fake-worker"} - (layout.failed / f"{task_id}.json").write_text( - json.dumps(fail_data) - ) - (running_dir / f"{task_id}.json").unlink(missing_ok=True) - return - time.sleep(0.05) - - worker_thread = threading.Thread( - target=fake_failing_worker, args=(queue_root,), daemon=True - ) - worker_thread.start() + worker = _fake_worker_thread(queue_root, n_tasks=1, fail=True) + worker.start() with patch( "miss_alignment.distributed.manager.LocalProvisioner", - _fake_provisioner_class(None), + return_value=_NoOpProvisioner(), ): with pytest.raises(RuntimeError, match="ts01"): run_distributed( @@ -1353,9 +1492,40 @@ def test_run_distributed_raises_if_any_task_fails(tmp_path): apply_ctf=False, downsample=2, devices=[0], + n_cluster_workers=None, queue_root=queue_root, - cluster_config=None, ) + + +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() ``` - [ ] **Step 2: Run tests to verify they fail** @@ -1375,6 +1545,7 @@ Expected: `ImportError` for `miss_alignment.distributed.manager`. from __future__ import annotations import json +import shutil import sys import threading import time @@ -1382,7 +1553,7 @@ from pathlib import Path import tqdm -from .config import ClusterConfig +from .config import load_cluster_config from .provisioner import ClusterProvisioner, LocalProvisioner, WorkerProvisioner from .queue import ( QueueLayout, @@ -1406,27 +1577,22 @@ def _write_manager_hb(layout: QueueLayout, seq: int) -> None: new_hb = layout.manager_hb / f"hb-{seq}" new_hb.write_text("") if seq > 0: - old_hb = layout.manager_hb / f"hb-{seq - 1}" - old_hb.unlink(missing_ok=True) + (layout.manager_hb / f"hb-{seq - 1}").unlink(missing_ok=True) def _sweep_stalled_workers(layout: QueueLayout) -> None: - """Move tasks from stalled worker dirs back to pending/.""" for worker_dir in layout.running.iterdir(): if not worker_dir.is_dir(): continue ticks = list(worker_dir.glob("hb-*")) if ticks: - latest = max(ticks, key=lambda p: p.stat().st_mtime) - age = time.time() - latest.stat().st_mtime + age = time.time() - max(ticks, key=lambda p: p.stat().st_mtime).stat().st_mtime else: - # No heartbeat yet — use dir mtime as proxy age = time.time() - worker_dir.stat().st_mtime if age <= _WORKER_STALL_TIMEOUT_S: continue - # Worker is stalled — recover its tasks for task_file in worker_dir.glob("*.json"): dest = layout.pending / task_file.name try: @@ -1438,7 +1604,6 @@ def _sweep_stalled_workers(layout: QueueLayout) -> None: ) except FileNotFoundError: pass - # Clean up heartbeat files for hb in worker_dir.glob("hb-*"): hb.unlink(missing_ok=True) try: @@ -1450,11 +1615,11 @@ def _sweep_stalled_workers(layout: QueueLayout) -> None: def _scheduler_thread( layout: QueueLayout, provisioner: WorkerProvisioner, - n_tasks: int, + n_workers: int, stop_event: threading.Event, ) -> None: - hb_seq = 0 - last_hb = 0.0 + hb_seq = 1 # seq 0 written before thread starts + last_hb = time.time() last_sweep = 0.0 while not stop_event.is_set(): @@ -1467,7 +1632,7 @@ def _scheduler_thread( if now - last_sweep >= _SCHEDULER_INTERVAL_S: _sweep_stalled_workers(layout) - provisioner.ensure_workers(n_tasks) + provisioner.ensure_workers(n_workers) last_sweep = now stop_event.wait(timeout=1.0) @@ -1484,19 +1649,18 @@ def run_distributed( apply_ctf: bool, downsample: int, devices: list[int], + n_cluster_workers: int | None, queue_root: Path, - cluster_config: ClusterConfig | None, ) -> dict[str, float]: """Write tasks, provision workers, block until all tasks are terminal. - Returns a dict mapping tilt-series name to final loss. - Raises RuntimeError listing all failed series if any task ends in failed/. + Returns dict[series_name → final_loss]. Raises RuntimeError listing all + failed series if any task ends in failed/. Deletes queue_root on exit. """ layout = QueueLayout(queue_root) layout.ensure_directories() clear_queue(layout) - # Fingerprint is the same for all tasks in one alignment phase. fingerprint = compute_fingerprint( model_checkpoint_path=str(model_checkpoint), setting=setting if isinstance(setting, str) else list(setting), @@ -1526,28 +1690,34 @@ def run_distributed( ) write_pending(layout, spec) - n_tasks = len(tilt_series_list) - if cluster_config is not None: + # Write the first manager heartbeat before starting workers so workers + # never see a missing heartbeat on startup. + _write_manager_hb(layout, seq=0) + + if n_cluster_workers is not None: + cluster_config = load_cluster_config() provisioner: WorkerProvisioner = ClusterProvisioner( queue_dir=queue_root, config=cluster_config ) + n_workers = n_cluster_workers else: provisioner = LocalProvisioner(queue_dir=queue_root, devices=devices) + n_workers = len(devices) stop_event = threading.Event() scheduler = threading.Thread( target=_scheduler_thread, - args=(layout, provisioner, n_tasks, stop_event), + args=(layout, provisioner, n_workers, stop_event), daemon=True, ) scheduler.start() - provisioner.ensure_workers(n_tasks) + provisioner.ensure_workers(n_workers) pending_ids = set(task_ids) losses: dict[str, float] = {} failed_series: list[str] = [] - pbar = tqdm.tqdm(total=n_tasks, desc="Tilt series alignment", file=sys.stdout) + pbar = tqdm.tqdm(total=len(task_ids), desc="Tilt series alignment", file=sys.stdout) try: while pending_ids: time.sleep(_POLL_INTERVAL_S) @@ -1578,6 +1748,7 @@ def run_distributed( stop_event.set() scheduler.join(timeout=5.0) provisioner.shutdown() + shutil.rmtree(queue_root, ignore_errors=True) if failed_series: raise RuntimeError( @@ -1588,15 +1759,7 @@ def run_distributed( return losses ``` -- [ ] **Step 4: Run tests to verify they pass** - -```bash -pytest tests/distributed/test_manager.py -v -``` - -Expected: both tests PASS. - -- [ ] **Step 5: Run all distributed tests together** +- [ ] **Step 4: Run all distributed tests** ```bash pytest tests/distributed/ -v @@ -1604,40 +1767,38 @@ pytest tests/distributed/ -v Expected: all tests PASS. -- [ ] **Step 6: Commit** +- [ ] **Step 5: Commit** ```bash git add src/miss_alignment/distributed/manager.py tests/distributed/test_manager.py -git commit -m "feat: add distributed manager with scheduler thread and poll loop" +git commit -m "feat: add distributed manager with scheduler thread, poll loop, and cleanup" ``` --- -## Task 6: Wire up `alignment/parallel.py` and CLI; delete `_parallel.py` +## Task 7: Wire up CLI, parallel, train, infer; delete `_parallel.py` **Files:** - Modify: `src/miss_alignment/alignment/parallel.py` +- Modify: `src/miss_alignment/train.py` +- Modify: `src/miss_alignment/infer.py` - Modify: `src/miss_alignment/_cli.py` - Modify: `src/miss_alignment/__init__.py` +- Modify: `src/miss_alignment/distributed/__init__.py` - Modify: `tests/test_parallel.py` - Delete: `src/miss_alignment/_parallel.py` **Interfaces:** -- Consumes: - - `run_distributed` from `distributed/manager.py` - - `load_cluster_config` from `distributed/config.py` - - `worker_miss_align` from `distributed/worker.py` -- Produces: `run_alignment_parallel` — same signature as before, same return type `dict[str, float]`. +- `run_alignment_parallel` gains `n_cluster_workers: int | None = None` parameter. +- `train_miss_align` and `infer_miss_align` each gain `n_cluster_workers: Optional[int] = typer.Option(None, ...)`. -- [ ] **Step 1: Update `tests/test_parallel.py` to use the new import** +- [ ] **Step 1: Update `tests/test_parallel.py`** -The existing `test_parallel.py` tests `_parallel.run_device_pool` directly. Since `_parallel.py` is being deleted, update the tests to verify the equivalent behaviour through `LocalProvisioner` instead. Replace the file entirely: +Replace the file entirely to test through the new `LocalProvisioner` path: ```python # tests/test_parallel.py """Tests for the distributed worker provisioner (replaces _parallel.py tests).""" -import subprocess -import sys from pathlib import Path from unittest.mock import MagicMock, patch @@ -1647,64 +1808,57 @@ from miss_alignment.distributed.provisioner import LocalProvisioner def test_local_provisioner_spawns_worker_per_device(tmp_path): - """LocalProvisioner starts one worker process per GPU device.""" 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_tasks=10) + p.ensure_workers(n_workers=10) assert mock_popen.call_count == 3 def test_local_provisioner_does_not_double_spawn(tmp_path): - """Calling ensure_workers twice does not spawn extra processes.""" 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_tasks=5) - p.ensure_workers(n_tasks=5) + p.ensure_workers(n_workers=5) + p.ensure_workers(n_workers=5) assert mock_popen.call_count == 1 -@pytest.mark.filterwarnings("ignore") def test_local_provisioner_shutdown_terminates(tmp_path): - """shutdown() terminates all spawned processes.""" 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_tasks=3) + p.ensure_workers(n_workers=5) p.shutdown() mock_proc.terminate.assert_called() ``` -- [ ] **Step 2: Run updated tests to verify they pass before touching sources** +- [ ] **Step 2: Run updated `test_parallel.py` to verify it passes now** ```bash pytest tests/test_parallel.py -v ``` -Expected: all 3 tests PASS (they import from `distributed.provisioner` which already exists). +Expected: all 3 tests PASS. - [ ] **Step 3: Update `alignment/parallel.py`** -Replace the file entirely: - ```python # src/miss_alignment/alignment/parallel.py from pathlib import Path -from ..distributed.config import load_cluster_config from ..distributed.manager import run_distributed @@ -1719,18 +1873,17 @@ def run_alignment_parallel( apply_ctf: bool, downsample: int, devices_list: list[int], + n_cluster_workers: int | None = None, ) -> dict[str, float]: """Distribute per-tilt-series alignment across local GPUs or a cluster. - With no cluster env vars set, workers are spawned as local child processes - (one per GPU in devices_list). Set MISS_CLUSTER_CONFIG and MISS_CLUSTER_SCRIPT - to fan work out to a batch scheduler instead. + 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 a dict mapping tilt-series stem names to their final loss values. + Returns dict mapping tilt-series stem names to their final loss values. """ - cluster_config = load_cluster_config() - # output_directory is the training directory in train.py and data_directory in - # infer.py — both are the top-level data dir, so tasks/ lives alongside the XMLs. queue_root = output_directory / "tasks" return run_distributed( @@ -1744,12 +1897,76 @@ def run_alignment_parallel( apply_ctf=apply_ctf, downsample=downsample, devices=devices_list, + n_cluster_workers=n_cluster_workers, queue_root=queue_root, - cluster_config=cluster_config, ) ``` -- [ ] **Step 4: Register `worker` subcommand in `_cli.py`** +- [ ] **Step 4: Add `--n-cluster-workers` to `train.py`** + +Add the new option to `train_miss_align`'s signature. Find the `preprocess: bool` option (last existing option, around line 308) and add after it: + +```python + 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.", + ), +``` + +Then pass it through to `run_alignment_parallel` in the call at line ~487: + +```python + run_alignment_parallel( + model_checkpoint=str(training_model_path), + tilt_series_list=tilt_series_list, + output_directory=training_directory, + setting=iteration_settings["alignment"], + patch_size=alignment_config["patch_size"], + patch_overlap=alignment_config["patch_overlap"], + batch_size=alignment_config["batch_size"], + apply_ctf=general_config["apply_ctf"], + downsample=iteration_settings["downsample"], + devices_list=devices_alignment, + n_cluster_workers=n_cluster_workers, + ) +``` + +- [ ] **Step 5: Add `--n-cluster-workers` to `infer.py`** + +Add the same option to `infer_miss_align`'s signature after `preprocess: bool` (around line 41): + +```python + 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.", + ), +``` + +Then pass it through to the `run_alignment_parallel` call at line ~161: + +```python + run_alignment_parallel( + model_checkpoint=str(model_checkpoint), + tilt_series_list=tilt_series_list, + output_directory=data_directory, + setting=iteration_settings["alignment"], + patch_size=alignment_config["patch_size"], + patch_overlap=alignment_config["patch_overlap"], + batch_size=alignment_config["batch_size"], + apply_ctf=general_config["apply_ctf"], + downsample=iteration_settings["downsample"], + devices_list=devices_alignment, + n_cluster_workers=n_cluster_workers, + ) +``` + +- [ ] **Step 6: Register the `worker` subcommand in `_cli.py`** ```python # src/miss_alignment/_cli.py @@ -1761,7 +1978,7 @@ from typer.core import TyperGroup class OrderCommands(TyperGroup): def list_commands(self, ctx: Context): """Return list of commands in the order appear.""" - return list(self.commands) # get commands using self.commands + return list(self.commands) cli = typer.Typer(cls=OrderCommands, add_completion=False, no_args_is_help=True) @@ -1772,7 +1989,7 @@ from .distributed.worker import worker_miss_align # noqa: E402 cli.command(name="worker")(worker_miss_align) ``` -- [ ] **Step 5: Export `worker_miss_align` from `__init__.py`** +- [ ] **Step 7: Export `worker_miss_align` from `__init__.py`** ```python # src/miss_alignment/__init__.py @@ -1801,47 +2018,7 @@ from .infer import infer_miss_align from .distributed.worker import worker_miss_align ``` -- [ ] **Step 6: Delete `_parallel.py`** - -```bash -git rm src/miss_alignment/_parallel.py -``` - -- [ ] **Step 7: Run the full test suite** - -```bash -pytest --color=yes -v -``` - -Expected: all tests PASS, no warnings-as-errors regressions. If `test_infer.py` or `test_train.py` import `_parallel` directly, fix those imports to remove them (the module no longer exists). - -- [ ] **Step 8: Run linter** - -```bash -ruff check --fix src/miss_alignment/ -ruff format src/miss_alignment/ -``` - -Expected: no errors. - -- [ ] **Step 9: Commit** - -```bash -git add src/miss_alignment/alignment/parallel.py src/miss_alignment/_cli.py src/miss_alignment/__init__.py tests/test_parallel.py -git commit -m "feat: wire distributed queue into run_alignment_parallel; add worker subcommand; delete _parallel.py" -``` - ---- - -## Task 7: Update `distributed/__init__.py` and run full suite - -**Files:** -- Modify: `src/miss_alignment/distributed/__init__.py` - -**Interfaces:** -- Produces: public re-exports for any consumer that imports directly from `miss_alignment.distributed`. - -- [ ] **Step 1: Update `__init__.py`** +- [ ] **Step 8: Update `distributed/__init__.py`** ```python # src/miss_alignment/distributed/__init__.py @@ -1882,27 +2059,41 @@ __all__ = [ ] ``` -- [ ] **Step 2: Run full test suite and linter** +- [ ] **Step 9: Delete `_parallel.py`** ```bash -pytest --color=yes --cov --cov-report=term-missing -ruff check src/miss_alignment/ -ruff format --check src/miss_alignment/ +git rm src/miss_alignment/_parallel.py +``` + +- [ ] **Step 10: Run full test suite and linter** + +```bash +pytest --color=yes -v +ruff check --fix src/miss_alignment/ +ruff format src/miss_alignment/ ``` -Expected: all tests PASS, coverage report shows `distributed/` coverage, no ruff errors. +Expected: all tests PASS. If any test imports `miss_alignment._parallel` directly, fix that import (the module no longer exists). No ruff errors. -- [ ] **Step 3: Verify `miss-alignment worker --help` works** +- [ ] **Step 11: Verify the worker subcommand is registered** ```bash +miss-alignment --help miss-alignment worker --help ``` -Expected output includes `--queue-dir`, `--device`, `--worker-id` options. +Expected: `worker` appears in the command list; `--queue-dir`, `--device`, `--worker-id` appear in its help. -- [ ] **Step 4: Final commit** +- [ ] **Step 12: Commit** ```bash -git add src/miss_alignment/distributed/__init__.py -git commit -m "feat: export distributed public API from __init__.py" +git add \ + src/miss_alignment/alignment/parallel.py \ + src/miss_alignment/train.py \ + src/miss_alignment/infer.py \ + src/miss_alignment/_cli.py \ + src/miss_alignment/__init__.py \ + src/miss_alignment/distributed/__init__.py \ + tests/test_parallel.py +git commit -m "feat: wire distributed queue into run_alignment_parallel; add --n-cluster-workers; register worker subcommand; delete _parallel.py" ``` diff --git a/docs/superpowers/specs/2026-07-04-distributed-inference-design.md b/docs/superpowers/specs/2026-07-04-distributed-inference-design.md index c36608a..d77214a 100644 --- a/docs/superpowers/specs/2026-07-04-distributed-inference-design.md +++ b/docs/superpowers/specs/2026-07-04-distributed-inference-design.md @@ -1,7 +1,7 @@ # Distributed Inference Design **Date:** 2026-07-04 -**Status:** Approved +**Status:** Approved (revised) **Scope:** Cluster distribution of the per-tilt-series alignment/inference phase --- @@ -12,15 +12,17 @@ With 300 tilt-series, the alignment phase (inference) takes 3× longer than the ## Goal -Distribute the per-series alignment tasks across a compute cluster so the head node fans out work, blocks until all series are done, and then continues to the next macro-iteration. Cluster mode is opt-in via environment variables; without them the system runs exactly as today (local multi-GPU pool). +Distribute the per-series alignment tasks across a compute cluster so the head node fans out work, blocks until all series are done, and then continues to the next macro-iteration. Cluster mode is opt-in via a new `--n-cluster-workers` CLI argument; without it the system runs exactly as today (local multi-GPU pool). --- ## Architecture -A new `miss_alignment/distributed/` module sits between the existing `alignment/parallel.py` public API and the underlying `evaluate_tilt_series` function. `train.py` and `infer.py` are unchanged. +A new `miss_alignment/distributed/` module sits between the existing `alignment/parallel.py` public API and the underlying `evaluate_tilt_series` function. -`run_alignment_parallel` (in `alignment/parallel.py`) calls `load_cluster_config()`. If cluster config is present it delegates to the distributed manager + `ClusterProvisioner`; otherwise it uses the distributed manager + `LocalProvisioner` (replacing the current `_parallel.py` internals). In both cases the same queue layer is used. +`run_alignment_parallel` accepts a new `n_cluster_workers: int | None` parameter. When it is `None`, `LocalProvisioner` is used (one subprocess per GPU, unchanged behavior). When it is set, `ClusterProvisioner` is used and submits exactly `n_cluster_workers` jobs; each worker runs the claim loop and processes as many series as it can grab until the queue drains. + +`evaluate_tilt_series` gains an optional `model: MissAlignment | None = None` parameter. When provided, the worker's resident model is used directly instead of loading from disk — this amortizes checkpoint loading across all series a worker processes in one macro-iteration. Callers that omit the parameter are unaffected. ### Components @@ -30,7 +32,7 @@ A new `miss_alignment/distributed/` module sits between the existing `alignment/ | `distributed/manager.py` | Head-node coordinator: writes tasks, runs scheduler thread, blocks until done | | `distributed/provisioner.py` | `WorkerProvisioner` interface + `LocalProvisioner` + `ClusterProvisioner` | | `distributed/worker.py` | `miss-alignment worker` subcommand: claims tasks, runs inference, writes results | -| `distributed/config.py` | Reads env vars, returns `ClusterConfig` dataclass or `None` | +| `distributed/config.py` | Reads env vars, returns `ClusterConfig` dataclass or raises if missing | --- @@ -43,7 +45,7 @@ One JSON file per tilt-series, written to `/tasks/pending/- "task_id": "0000003-tilt_series_01", "model_checkpoint_path": "/data/project/iter2/model.ckpt", "tilt_series_path": "/data/project/tilt_series_01.xml", - "output_directory": "/data/project/iter2/", + "output_directory": "/data/project/", "setting": "anchoring", "patch_size": 96, "patch_overlap": 0.1, @@ -54,7 +56,9 @@ One JSON file per tilt-series, written to `/tasks/pending/- } ``` -`init_fingerprint` covers the model checkpoint path and all alignment parameters. A worker skips reloading the model when consecutive tasks share the same fingerprint, amortizing checkpoint loading across many series in one macro-iteration. +`init_fingerprint` covers the model checkpoint path and all alignment parameters. A worker skips reloading the model when consecutive tasks share the same fingerprint, amortizing checkpoint loading across the many series each worker processes in one macro-iteration. + +`setting` is serialized to JSON as a string or array. Workers convert array back to `tuple` before passing to `evaluate_tilt_series`. On completion, result fields are appended before writing to `done/`: ```json @@ -63,7 +67,7 @@ On completion, result fields are appended before writing to `done/`: On failure, error fields are appended before writing to `failed/`: ```json -{ "error": "CUDA out of memory...", "worker_id": "local-12345-gpu0" } +{ "error": "CUDA out of memory...", "worker_id": "cluster-node01-12345-0" } ``` --- @@ -79,8 +83,10 @@ On failure, error fields are appended before writing to `failed/`: │ └── hb- # worker heartbeat tick files (latest only) ├── done/ # completed task JSONs ├── failed/ # failed task JSONs -└── manager/ - └── hb- # manager heartbeat tick files (latest only) +├── manager/ +│ └── hb- # manager heartbeat tick files (latest only) +└── cluster/ + └── worker-.sh # rendered cluster submission scripts ``` The queue directory is always `/tasks/`. Since the training directory must already be on a shared filesystem (workers need the XML/MRC files), no additional configuration is needed for cluster nodes to access it. @@ -111,9 +117,11 @@ No lock files. The OS rename is the coordination primitive. - Check manager heartbeat (`tasks/manager/hb-*`): if latest tick is >120s old → exit cleanly (manager is dead) - List `pending/`, shuffle, attempt rename - On empty queue: exit cleanly -- On claim: check if `init_fingerprint` matches last loaded model; if not, load checkpoint from `model_checkpoint_path` -- Call `evaluate_tilt_series(..., device=f"cuda:{device}")` — all internal LBFGS retries and optimization passes run unchanged -- Write result to `done/` or `failed/`, delete running copy, loop +- On claim: check if `init_fingerprint` matches last loaded model; if not, load checkpoint from `model_checkpoint_path` and cache it +- Call `evaluate_tilt_series(..., model=cached_model, device=f"cuda:{device}")` — all internal LBFGS retries and optimization passes run unchanged +- Write result to `done/` or `failed/`, delete running copy, loop back + +Each worker processes **many series** per run. The checkpoint is loaded only once per macro-iteration (all tasks in a phase share the same `init_fingerprint`). **No task-level retries.** Failures are deterministic (bad data, corrupt file, config error); the manager hard-fails after all remaining tasks complete. @@ -122,21 +130,22 @@ No lock files. The OS rename is the coordination primitive. ## Manager Lifecycle **Startup:** -1. Clear stale queue state from any prior run (delete `pending/`, `done/`, `failed/` contents; recover any `running/` orphans back to `pending/`) +1. Clear stale queue state from any prior run: first delete `pending/`, `done/`, `failed/` contents; then recover any `running/` orphans back to `pending/` 2. Write all task JSONs to `pending/` -3. Start `ClusterProvisioner` or `LocalProvisioner` +3. Write the first manager heartbeat tick immediately (before starting workers, to avoid a race where workers start and find no heartbeat) 4. Start scheduler background thread -5. Block in poll loop (500ms interval) until all tasks are in `done/` or `failed/` +5. Start `ClusterProvisioner` or `LocalProvisioner`, call `ensure_workers(n_workers)` +6. Block in poll loop (500ms interval) until all tasks are in `done/` or `failed/` **Scheduler thread** (runs every ~10s): - Write manager heartbeat to `tasks/manager/hb-` -- Sweep stalled workers: for each `running//`, check age of latest `hb-` file; if >120s stale → move task back to `pending/`, delete `running//` +- Sweep stalled workers: for each `running//`, check age of latest `hb-` file; if >120s stale → move task back to `pending/`, clean up dir - Call `provisioner.ensure_workers()` to respawn any dead local workers **Shutdown** (on completion or `KeyboardInterrupt`): - Call `provisioner.shutdown()` (SIGTERM local children / cancel SLURM job IDs) -- Delete `tasks/` directory -- Return `dict[series_name → final_loss]` on full success, or raise with list of failed series if any task is in `failed/` +- Delete `tasks/` directory (clean state for next macro-iteration) +- Return `dict[series_name → final_loss]` on full success, or raise `RuntimeError` listing all failed series if any task is in `failed/` **Failure policy:** if any series ends in `failed/`, the manager raises after all other tasks complete. `train.py` and `infer.py` propagate this as a hard failure — training stops. @@ -146,9 +155,9 @@ No lock files. The OS rename is the coordination primitive. ### `LocalProvisioner` -Activated when neither `MISS_CLUSTER_CONFIG` nor `MISS_CLUSTER_SCRIPT` is set. +Activated when `n_cluster_workers` is `None` (default). -- `ensure_workers(target)`: spawns `miss-alignment worker --queue-dir --device ` via `subprocess.Popen`, one per GPU in `devices_alignment` (same list as today: `list(range(torch.cuda.device_count()))`, respecting `CUDA_VISIBLE_DEVICES`) +- `ensure_workers(n_workers)`: spawns `miss-alignment worker --queue-dir --device ` via `subprocess.Popen`, one per GPU in `devices_alignment` (same list as today: `list(range(torch.cuda.device_count()))`, respecting `CUDA_VISIBLE_DEVICES`). `n_workers` is ignored; the number of local workers is always equal to the number of devices. - Respawns any that have exited prematurely (checked each scheduler tick) - `shutdown()`: SIGTERM all children, short timeout, SIGKILL if needed @@ -156,7 +165,9 @@ This replaces `_parallel.py`'s `run_device_pool` / `mp.spawn` internals with no ### `ClusterProvisioner` -Activated when both `MISS_CLUSTER_CONFIG` and `MISS_CLUSTER_SCRIPT` are set. +Activated when `n_cluster_workers` is set. + +Requires both `MISS_CLUSTER_CONFIG` and `MISS_CLUSTER_SCRIPT` env vars to be set; raises `RuntimeError` immediately if either is absent, so the user gets a clear error rather than a silent fallback to local mode. **`MISS_CLUSTER_CONFIG`** — path to a JSON file: ```json @@ -178,37 +189,70 @@ conda activate miss-alignment {{command}} ``` -- `{{command}}` is filled with the `miss-alignment worker` invocation by the provisioner +- `{{command}}` is filled by the provisioner with the `miss-alignment worker` invocation - All cluster-specific settings (partition, memory, time limit, environment setup) are the user's responsibility in the template - Additional `{{custom_var}}` placeholders filled via `MISS_CLUSTER_VAR_=` env vars - Rendered scripts are written to `tasks/cluster/worker-.sh` -- The `{{command}}` is `miss-alignment worker --queue-dir --device 0 --worker-id "$(hostname)-$$-"` — the `$(hostname)` and `$$` are expanded by the compute node's shell at runtime, guaranteeing unique, stable worker IDs even across resubmissions -- One job submitted per tilt-series; job IDs stored for cancellation +- The injected `{{command}}` is: `miss-alignment worker --queue-dir --device 0 --worker-id "$(hostname)-$$-"` — `$(hostname)` and `$$` are expanded by the compute node's shell at runtime, guaranteeing unique, stable worker IDs +- `ensure_workers(n_workers)` submits exactly `n_workers` jobs; each worker pulls tasks from the shared queue until it drains - `shutdown()`: runs configured `cancel` command for each stored job ID; registers SIGINT/SIGTERM handlers so Ctrl-C on the head node cancels the cluster pool --- ## Configuration -All configuration via environment variables — no new CLI flags on `train` or `infer`: +### CLI changes + +`miss-alignment train` and `miss-alignment infer` each gain one new optional argument: + +| Argument | Type | Default | Meaning | +|---|---|---|---| +| `--n-cluster-workers` | `int \| None` | `None` | Number of cluster jobs to submit. When set, activates cluster mode. When absent, local multi-GPU mode is used (unchanged). | + +### Environment variables (cluster mode only) | Env var | Purpose | |---|---| -| `MISS_CLUSTER_CONFIG` | Path to cluster scheduler JSON config. Activates cluster mode when set together with `MISS_CLUSTER_SCRIPT`. | -| `MISS_CLUSTER_SCRIPT` | Path to job submission shell script template. | +| `MISS_CLUSTER_CONFIG` | Path to cluster scheduler JSON config. **Required** when `--n-cluster-workers` is set. | +| `MISS_CLUSTER_SCRIPT` | Path to job submission shell script template. **Required** when `--n-cluster-workers` is set. | | `MISS_CLUSTER_VAR_` | Additional template variables, e.g. `MISS_CLUSTER_VAR_partition=gpu`. | -`config.py` reads these at call time (not import time) and returns a `ClusterConfig` dataclass or `None`. +`config.py` is called only when `n_cluster_workers` is not `None`. It raises `RuntimeError` if either env var is absent, with a message naming which one is missing. + +--- + +## Changes to `evaluate_tilt_series` + +A single optional parameter is added: + +```python +def evaluate_tilt_series( + model_checkpoint_path: Path, + tilt_series_path: Path, + output_directory: Path, + setting: str | tuple[int, int] | tuple[int, int, int, int] = "anchoring", + patch_size: int = 96, + patch_overlap: float = 0.1, + batch_size: int = 16, + apply_ctf: bool = True, + downsample: int = 1, + device: str = "cpu", + initial_reliable_fraction: float = 1 / 2, + n_control_points: int = 7, + model: MissAlignment | None = None, # NEW +) -> tuple[Path, list[float]]: +``` + +When `model` is provided, the function uses it directly instead of loading from `model_checkpoint_path`. This is the only change to the function. All existing callers (which omit `model`) are unaffected. --- ## What Does Not Change -- `evaluate_tilt_series` — called identically, all internal optimization passes unchanged -- `train.py` / `infer.py` — no changes; `run_alignment_parallel` signature unchanged +- `evaluate_tilt_series` internal logic — only the signature gains one optional parameter - LBFGS retries, anchoring iterations, coarse-to-fine spline passes — all internal to `evaluate_tilt_series` - `CUDA_VISIBLE_DEVICES` still controls which local GPUs are used -- Single-GPU behavior is identical (one `LocalProvisioner` worker process per GPU) +- Single-GPU local behavior is identical --- @@ -221,10 +265,15 @@ All configuration via environment variables — no new CLI flags on `train` or ` - `src/miss_alignment/distributed/provisioner.py` - `src/miss_alignment/distributed/worker.py` - `src/miss_alignment/distributed/config.py` +- `src/miss_alignment/__main__.py` — enables `python -m miss_alignment` for `LocalProvisioner` subprocess launch **Modified files:** -- `src/miss_alignment/alignment/parallel.py` — replace `run_device_pool` delegation with distributed manager call +- `src/miss_alignment/alignment/tilt_series.py` — add optional `model` parameter to `evaluate_tilt_series` +- `src/miss_alignment/alignment/parallel.py` — replace `run_device_pool` call with distributed manager; add `n_cluster_workers` parameter +- `src/miss_alignment/train.py` — accept and pass through `--n-cluster-workers` +- `src/miss_alignment/infer.py` — accept and pass through `--n-cluster-workers` - `src/miss_alignment/_cli.py` — register `worker` subcommand +- `src/miss_alignment/__init__.py` — export `worker_miss_align` **Deleted files:** - `src/miss_alignment/_parallel.py` — superseded by `LocalProvisioner`; deleted once the new path is validated in tests From 93faad271953e28f38aeb32b35b22cc7d4d764fc Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 18:15:53 -0700 Subject: [PATCH 05/33] feat: add distributed queue layer with atomic rename claim protocol Co-Authored-By: Claude Fable 5 --- src/miss_alignment/__main__.py | 3 + src/miss_alignment/distributed/__init__.py | 1 + src/miss_alignment/distributed/queue.py | 197 +++++++++++++++++++++ tests/distributed/__init__.py | 0 tests/distributed/test_queue.py | 127 +++++++++++++ 5 files changed, 328 insertions(+) create mode 100644 src/miss_alignment/__main__.py create mode 100644 src/miss_alignment/distributed/__init__.py create mode 100644 src/miss_alignment/distributed/queue.py create mode 100644 tests/distributed/__init__.py create mode 100644 tests/distributed/test_queue.py diff --git a/src/miss_alignment/__main__.py b/src/miss_alignment/__main__.py new file mode 100644 index 0000000..7f9e2d3 --- /dev/null +++ b/src/miss_alignment/__main__.py @@ -0,0 +1,3 @@ +from miss_alignment import cli + +cli() diff --git a/src/miss_alignment/distributed/__init__.py b/src/miss_alignment/distributed/__init__.py new file mode 100644 index 0000000..d431841 --- /dev/null +++ b/src/miss_alignment/distributed/__init__.py @@ -0,0 +1 @@ +"""Disk-based distributed task queue for miss-alignment inference.""" diff --git a/src/miss_alignment/distributed/queue.py b/src/miss_alignment/distributed/queue.py new file mode 100644 index 0000000..4444a72 --- /dev/null +++ b/src/miss_alignment/distributed/queue.py @@ -0,0 +1,197 @@ +"""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" + + def ensure_directories(self) -> None: + for d in ( + self.pending, + self.running, + self.done, + self.failed, + self.manager_hb, + self.cluster, + ): + 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 + + +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/tests/distributed/__init__.py b/tests/distributed/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/distributed/test_queue.py b/tests/distributed/test_queue.py new file mode 100644 index 0000000..e86451e --- /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 From 7fdf6ff1e4249f40ccf98bcde9c786b8de6480c9 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 18:20:33 -0700 Subject: [PATCH 06/33] feat: add cluster config reader (raises if env vars missing) Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/config.py | 56 +++++++++++++++++++++ tests/distributed/test_config.py | 63 ++++++++++++++++++++++++ 2 files changed, 119 insertions(+) create mode 100644 src/miss_alignment/distributed/config.py create mode 100644 tests/distributed/test_config.py diff --git a/src/miss_alignment/distributed/config.py b/src/miss_alignment/distributed/config.py new file mode 100644 index 0000000..219ebd0 --- /dev/null +++ b/src/miss_alignment/distributed/config.py @@ -0,0 +1,56 @@ +"""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 +from pathlib import Path + + +@dataclass +class ClusterConfig: + submit: str + submit_job_id_regex: str + cancel: str + script_path: Path + + +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', and 'cancel' 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"], + script_path=script_path, + ) diff --git a/tests/distributed/test_config.py b/tests/distributed/test_config.py new file mode 100644 index 0000000..db033ab --- /dev/null +++ b/tests/distributed/test_config.py @@ -0,0 +1,63 @@ +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}}", + } + 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 + + +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() From b3a470988b73057d2eea8d9204289022693b71ab Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 18:26:21 -0700 Subject: [PATCH 07/33] feat: add worker subcommand with model fingerprint reuse across tasks Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/worker.py | 150 +++++++++++++++++++++ tests/distributed/test_worker.py | 158 +++++++++++++++++++++++ 2 files changed, 308 insertions(+) create mode 100644 src/miss_alignment/distributed/worker.py create mode 100644 tests/distributed/test_worker.py diff --git a/src/miss_alignment/distributed/worker.py b/src/miss_alignment/distributed/worker.py new file mode 100644 index 0000000..7629426 --- /dev/null +++ b/src/miss_alignment/distributed/worker.py @@ -0,0 +1,150 @@ +"""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 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 .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}" + new_hb.write_text("") + if seq > 0: + (worker_dir / f"hb-{seq - 1}").unlink(missing_ok=True) + + +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 run_worker_loop( + layout: QueueLayout, + worker_id: str, + device: str, + manager_hb_timeout_s: float = _MANAGER_HB_TIMEOUT_S, +) -> None: + """Main worker loop: claim → evaluate → write result. Repeat until queue empty.""" + worker_dir = layout.worker_dir(worker_id) + worker_dir.mkdir(parents=True, exist_ok=True) + + last_fingerprint: str | None = None + cached_model: MissAlignment | None = None + hb_seq = 0 + last_hb_time = 0.0 + + while True: + age = _manager_hb_age_s(layout) + if age > manager_hb_timeout_s: + print( + f"[{worker_id}] Manager heartbeat stale ({age:.0f}s). Exiting.", + file=sys.stderr, + ) + return + + now = time.time() + if now - last_hb_time >= _HB_INTERVAL_S: + _write_worker_hb(worker_dir, hb_seq) + hb_seq += 1 + last_hb_time = now + + spec = claim_one(layout, worker_id) + if spec is None: + return # queue empty, exit cleanly + + print(f"[{worker_id}] Claimed {spec.task_id}", file=sys.stderr) + + if spec.init_fingerprint != last_fingerprint: + cached_model = _load_model(spec.model_checkpoint_path) + last_fingerprint = spec.init_fingerprint + + # Convert setting back to tuple if it was serialized as a list. + setting = ( + tuple(spec.setting) if isinstance(spec.setting, list) else spec.setting + ) + + try: + _, 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, + ) + final_loss = float(loss_values[-1]) if loss_values else float("nan") + mark_done(layout, worker_id, spec, final_loss=final_loss, device=device) + print( + f"[{worker_id}] Done {spec.task_id} loss={final_loss:.4f}", + file=sys.stderr, + ) + except Exception: + error = traceback.format_exc() + mark_failed(layout, worker_id, spec, error=error) + print( + f"[{worker_id}] Failed {spec.task_id}:\n{error}", + file=sys.stderr, + ) + + +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" + run_worker_loop(layout, worker_id, cuda_device) diff --git a/tests/distributed/test_worker.py b/tests/distributed/test_worker.py new file mode 100644 index 0000000..144cf0f --- /dev/null +++ b/tests/distributed/test_worker.py @@ -0,0 +1,158 @@ +"""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 From 590af8fe49e9e06d6d90d9ac5b94458c8c289134 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 18:44:45 -0700 Subject: [PATCH 08/33] feat: add optional model param to evaluate_tilt_series for checkpoint reuse Co-Authored-By: Claude Fable 5 --- src/miss_alignment/alignment/tilt_series.py | 33 +++++++++++---------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/src/miss_alignment/alignment/tilt_series.py b/src/miss_alignment/alignment/tilt_series.py index d55267c..dcdec76 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) From e6b5a20b5f722dc5f9e62a6c6141261fff70b84a Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 18:48:54 -0700 Subject: [PATCH 09/33] feat: add LocalProvisioner and ClusterProvisioner Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/provisioner.py | 117 +++++++++++++++++ tests/distributed/test_provisioner.py | 120 ++++++++++++++++++ 2 files changed, 237 insertions(+) create mode 100644 src/miss_alignment/distributed/provisioner.py create mode 100644 tests/distributed/test_provisioner.py diff --git a/src/miss_alignment/distributed/provisioner.py b/src/miss_alignment/distributed/provisioner.py new file mode 100644 index 0000000..ce023cf --- /dev/null +++ b/src/miss_alignment/distributed/provisioner.py @@ -0,0 +1,117 @@ +"""Worker provisioners: spawn local child processes or submit cluster jobs.""" + +from __future__ import annotations + +import os +import re +import subprocess +import sys +from abc import ABC, abstractmethod +from pathlib import Path + +from .config import ClusterConfig + + +class WorkerProvisioner(ABC): + @abstractmethod + def ensure_workers(self, n_workers: int) -> None: + """Ensure workers are running. Called once at startup and each scheduler tick.""" + + @abstractmethod + def shutdown(self) -> None: + """Terminate all managed workers.""" + + +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] = {} + + def ensure_workers(self, n_workers: int) -> None: + 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() + + +class ClusterProvisioner(WorkerProvisioner): + """Submits exactly n_workers cluster jobs, each running until the queue drains.""" + + def __init__(self, queue_dir: Path, config: ClusterConfig) -> None: + self._queue_dir = queue_dir + self._config = config + self._job_ids: list[str] = [] + self._scripts_dir = queue_dir / "cluster" + self._scripts_dir.mkdir(parents=True, exist_ok=True) + + 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}"' + ) + rendered = template_text.replace("{{command}}", command) + 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 ensure_workers(self, n_workers: int) -> None: + already = len(self._job_ids) + for i in range(already, n_workers): + script_path = self._render_script(i) + 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_ids.append(match.group(1)) + + def shutdown(self) -> None: + for job_id in self._job_ids: + cancel_cmd = self._config.cancel.replace("{{job_id}}", job_id) + subprocess.run(cancel_cmd, shell=True, stderr=subprocess.DEVNULL) + self._job_ids.clear() diff --git a/tests/distributed/test_provisioner.py b/tests/distributed/test_provisioner.py new file mode 100644 index 0000000..f06f574 --- /dev/null +++ b/tests/distributed/test_provisioner.py @@ -0,0 +1,120 @@ +import sys +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, LocalProvisioner + + +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 + # verify device args appear in calls + 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) # should respawn because poll() != None + + 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() + + +def test_cluster_provisioner_submits_n_workers_jobs(tmp_path): + script = tmp_path / "worker.sh" + script.write_text("#!/bin/bash\n{{command}}\n") + cfg = ClusterConfig( + submit="sbatch {{script_path}}", + submit_job_id_regex=r"Submitted batch job (\d+)", + cancel="scancel {{job_id}}", + script_path=script, + ) + + submitted = [] + + def fake_run(cmd, **kwargs): + submitted.append(cmd) + result = MagicMock() + result.stdout = "Submitted batch job 12345\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(submitted) == 4 + + +def test_cluster_provisioner_cancels_on_shutdown(tmp_path): + script = tmp_path / "worker.sh" + script.write_text("#!/bin/bash\n{{command}}\n") + cfg = ClusterConfig( + submit="sbatch {{script_path}}", + submit_job_id_regex=r"Submitted batch job (\d+)", + cancel="scancel {{job_id}}", + script_path=script, + ) + + cancel_calls = [] + + def fake_run(cmd, **kwargs): + if "sbatch" in cmd: + result = MagicMock() + result.stdout = "Submitted batch job 99999\n" + 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) From 1f2c90f10eab82babb76372386a60ef3f28d9959 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 18:56:34 -0700 Subject: [PATCH 10/33] feat: add distributed manager with scheduler thread, poll loop, and cleanup Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/manager.py | 217 ++++++++++++++++++++++ tests/distributed/test_manager.py | 158 ++++++++++++++++ 2 files changed, 375 insertions(+) create mode 100644 src/miss_alignment/distributed/manager.py create mode 100644 tests/distributed/test_manager.py diff --git a/src/miss_alignment/distributed/manager.py b/src/miss_alignment/distributed/manager.py new file mode 100644 index 0000000..d64c249 --- /dev/null +++ b/src/miss_alignment/distributed/manager.py @@ -0,0 +1,217 @@ +"""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, 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-*")) + if ticks: + age = time.time() - max(ticks, key=lambda p: p.stat().st_mtime).stat().st_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"): + dest = layout.pending / task_file.name + try: + os.rename(task_file, dest) + print( + f"[manager] Recovered stalled task {task_file.name} to pending", + file=sys.stderr, + ) + except FileNotFoundError: + pass + for hb in worker_dir.glob("hb-*"): + hb.unlink(missing_ok=True) + try: + worker_dir.rmdir() + except OSError: + pass + + +def _scheduler_thread( + layout: QueueLayout, + provisioner: WorkerProvisioner, + n_workers: int, + stop_event: threading.Event, +) -> None: + hb_seq = 1 # seq 0 written before thread starts + last_hb = time.time() + last_sweep = 0.0 + + 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) + provisioner.ensure_workers(n_workers) + last_sweep = now + + stop_event.wait(timeout=1.0) + + +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, +) -> 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. + """ + layout = QueueLayout(queue_root) + layout.ensure_directories() + clear_queue(layout) + + 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, + ) + + 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), + 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, + ) + 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() + provisioner: WorkerProvisioner = ClusterProvisioner( + queue_dir=queue_root, config=cluster_config + ) + n_workers = n_cluster_workers + else: + provisioner = LocalProvisioner(queue_dir=queue_root, devices=devices) + n_workers = len(devices) + + stop_event = threading.Event() + scheduler = threading.Thread( + target=_scheduler_thread, + args=(layout, provisioner, n_workers, stop_event), + daemon=True, + ) + scheduler.start() + provisioner.ensure_workers(n_workers) + + pending_ids = set(task_ids) + losses: dict[str, float] = {} + failed_series: list[str] = [] + + pbar = tqdm.tqdm(total=len(task_ids), desc="Tilt series alignment", file=sys.stdout) + try: + while pending_ids: + time.sleep(_POLL_INTERVAL_S) + + 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/tests/distributed/test_manager.py b/tests/distributed/test_manager.py new file mode 100644 index 0000000..55b6ff2 --- /dev/null +++ b/tests/distributed/test_manager.py @@ -0,0 +1,158 @@ +"""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 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() From 889ee597c08232d95908c217e8990a0d7ad2737e Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 19:02:45 -0700 Subject: [PATCH 11/33] feat: wire distributed queue into run_alignment_parallel; add --n-cluster-workers; register worker subcommand; delete _parallel.py Also migrate prepare_stacks.py and preprocessing.py off _parallel import. Co-Authored-By: Claude Fable 5 --- src/miss_alignment/__init__.py | 2 + src/miss_alignment/_cli.py | 4 + src/miss_alignment/_parallel.py | 83 ------------- src/miss_alignment/alignment/parallel.py | 110 ++++-------------- src/miss_alignment/distributed/__init__.py | 34 ++++++ src/miss_alignment/distributed/provisioner.py | 7 +- src/miss_alignment/infer.py | 8 ++ src/miss_alignment/prepare_stacks.py | 48 +++++++- src/miss_alignment/preprocessing.py | 50 +++++++- src/miss_alignment/train.py | 8 ++ tests/test_parallel.py | 86 +++++++------- 11 files changed, 219 insertions(+), 221 deletions(-) delete mode 100644 src/miss_alignment/_parallel.py diff --git a/src/miss_alignment/__init__.py b/src/miss_alignment/__init__.py index 8745d7f..dface4e 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/_cli.py b/src/miss_alignment/_cli.py index d89206d..a304641 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 e756d43..0000000 --- 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 a216132..ba691ca 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/distributed/__init__.py b/src/miss_alignment/distributed/__init__.py index d431841..9d2d21d 100644 --- a/src/miss_alignment/distributed/__init__.py +++ b/src/miss_alignment/distributed/__init__.py @@ -1 +1,35 @@ """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, 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", + "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/provisioner.py b/src/miss_alignment/distributed/provisioner.py index ce023cf..bd16cb1 100644 --- a/src/miss_alignment/distributed/provisioner.py +++ b/src/miss_alignment/distributed/provisioner.py @@ -15,7 +15,10 @@ class WorkerProvisioner(ABC): @abstractmethod def ensure_workers(self, n_workers: int) -> None: - """Ensure workers are running. Called once at startup and each scheduler tick.""" + """Ensure workers are running. + + Called once at startup and on each scheduler tick to respawn dead workers. + """ @abstractmethod def shutdown(self) -> None: @@ -85,7 +88,7 @@ def _render_script(self, index: int) -> Path: rendered = template_text.replace("{{command}}", command) for key, value in os.environ.items(): if key.startswith("MISS_CLUSTER_VAR_"): - var_name = key[len("MISS_CLUSTER_VAR_"):].lower() + var_name = key[len("MISS_CLUSTER_VAR_") :].lower() rendered = rendered.replace(f"{{{{{var_name}}}}}", value) script_path = self._scripts_dir / f"worker-{index}.sh" diff --git a/src/miss_alignment/infer.py b/src/miss_alignment/infer.py index ecdae73..1474975 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 with pretilt estimation.", ), + 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. @@ -159,6 +166,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 8bc0b6a..23bcc8a 100644 --- a/src/miss_alignment/prepare_stacks.py +++ b/src/miss_alignment/prepare_stacks.py @@ -4,14 +4,58 @@ preprocessed tilt stacks ready for training. """ +import multiprocessing as mp import queue +import sys +import time from pathlib import Path import mrcfile +import tqdm from warpylib import TiltSeries from warpylib.movie import Movie -from ._parallel import run_device_pool + +def _run_device_pool(jobs, runner, runner_args, devices, desc): + """Minimal one-process-per-GPU work queue. Internal to prepare_stacks.""" + 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 = [] + 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: + 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 def _get_original_pixel_size(tilt_series: TiltSeries) -> float: @@ -162,7 +206,7 @@ def prepare_stacks_parallel( f"Preparing stacks for {len(xml_files)} tilt series at {desired_pixel_size} Å" ) - run_device_pool( + _run_device_pool( jobs=xml_files, runner=_prepare_stacks_runner, runner_args=(desired_pixel_size,), diff --git a/src/miss_alignment/preprocessing.py b/src/miss_alignment/preprocessing.py index 120ccc1..d809774 100644 --- a/src/miss_alignment/preprocessing.py +++ b/src/miss_alignment/preprocessing.py @@ -1,12 +1,58 @@ """Preprocessing utilities for tilt-series alignment.""" +import multiprocessing as mp import queue +import sys +import time from pathlib import Path -from ._parallel import run_device_pool +import tqdm + from .data.io import TiltSeriesData +def _run_device_pool(jobs, runner, runner_args, devices, desc): + """Minimal one-process-per-GPU work queue. Internal to preprocessing.""" + 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 = [] + 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: + 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 + + def _run_cross_correlation_single( xml_file: Path, device: int | None, @@ -138,7 +184,7 @@ def run_cross_correlation_alignment_parallel( else: print(" Using default device assignment\n") - run_device_pool( + _run_device_pool( jobs=xml_files, runner=_cross_correlation_runner, runner_args=(lowpass_cutoff, pretilt_search_range), diff --git a/src/miss_alignment/train.py b/src/miss_alignment/train.py index 9c4ad3d..ef73313 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 with pretilt estimation.", ), + 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. @@ -485,6 +492,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/test_parallel.py b/tests/test_parallel.py index 92d6594..b75c5fd 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() From 35c11a3206c6137437105a67bd9708267b18dd4a Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 19:14:09 -0700 Subject: [PATCH 12/33] feat: extend distributed queue to prepare_stacks and cross_correlation tasks - TaskSpec gains task_type, desired_pixel_size, lowpass_cutoff, pretilt_search_range - Worker dispatches on task_type; alignment/prepare_stacks/cross_correlation all supported - prepare_stacks_parallel and run_cross_correlation_alignment_parallel migrated off local _run_device_pool copies onto run_distributed - n_cluster_workers threaded through to all three task types in train.py and infer.py - Duplicate _run_device_pool code removed from prepare_stacks.py and preprocessing.py Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/manager.py | 46 ++++++++--- src/miss_alignment/distributed/queue.py | 9 ++- src/miss_alignment/distributed/worker.py | 84 ++++++++++++++----- src/miss_alignment/infer.py | 2 + src/miss_alignment/prepare_stacks.py | 98 ++++++----------------- src/miss_alignment/preprocessing.py | 98 +++++------------------ src/miss_alignment/train.py | 2 + tests/distributed/test_manager.py | 50 ++++++++++++ tests/distributed/test_worker.py | 74 +++++++++++++++++ 9 files changed, 281 insertions(+), 182 deletions(-) diff --git a/src/miss_alignment/distributed/manager.py b/src/miss_alignment/distributed/manager.py index d64c249..7670564 100644 --- a/src/miss_alignment/distributed/manager.py +++ b/src/miss_alignment/distributed/manager.py @@ -45,7 +45,10 @@ def _sweep_stalled_workers(layout: QueueLayout) -> None: continue ticks = list(worker_dir.glob("hb-*")) if ticks: - age = time.time() - max(ticks, key=lambda p: p.stat().st_mtime).stat().st_mtime + age = ( + time.time() + - max(ticks, key=lambda p: p.stat().st_mtime).stat().st_mtime + ) else: age = time.time() - worker_dir.stat().st_mtime @@ -109,25 +112,38 @@ def run_distributed( 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, + pretilt_search_range: tuple | 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 = 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, - ) + # 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): @@ -135,7 +151,9 @@ def run_distributed( task_ids.append(task_id) spec = TaskSpec( task_id=task_id, - model_checkpoint_path=str(model_checkpoint), + 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), @@ -145,6 +163,12 @@ def run_distributed( apply_ctf=apply_ctf, downsample=downsample, init_fingerprint=fingerprint, + task_type=task_type, + desired_pixel_size=desired_pixel_size, + lowpass_cutoff=lowpass_cutoff, + pretilt_search_range=( + list(pretilt_search_range) if pretilt_search_range is not None else None + ), ) write_pending(layout, spec) diff --git a/src/miss_alignment/distributed/queue.py b/src/miss_alignment/distributed/queue.py index 4444a72..5322d2d 100644 --- a/src/miss_alignment/distributed/queue.py +++ b/src/miss_alignment/distributed/queue.py @@ -75,6 +75,11 @@ class TaskSpec: 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 + pretilt_search_range: list | None = None def _atomic_write(path: Path, data: dict) -> None: @@ -86,7 +91,9 @@ def _atomic_write(path: Path, data: dict) -> None: 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__}) + return TaskSpec( + **{k: v for k, v in data.items() if k in TaskSpec.__dataclass_fields__} + ) def compute_fingerprint( diff --git a/src/miss_alignment/distributed/worker.py b/src/miss_alignment/distributed/worker.py index 7629426..1bc0ccc 100644 --- a/src/miss_alignment/distributed/worker.py +++ b/src/miss_alignment/distributed/worker.py @@ -22,6 +22,8 @@ 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, @@ -58,6 +60,65 @@ def _load_model(checkpoint_path: str) -> MissAlignment: 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": + pretilt_range = ( + tuple(spec.pretilt_search_range) + if spec.pretilt_search_range is not None + else (-30.0, 30.0) + ) + _run_cross_correlation_single( + xml_file=Path(spec.tilt_series_path), + device=_device_int(device), + lowpass_cutoff=spec.lowpass_cutoff or 0.25, + pretilt_search_range=pretilt_range, + ) + return 0.0 + + else: + raise ValueError(f"Unknown task_type: {spec.task_type!r}") + + def run_worker_loop( layout: QueueLayout, worker_id: str, @@ -94,30 +155,13 @@ def run_worker_loop( print(f"[{worker_id}] Claimed {spec.task_id}", file=sys.stderr) - if spec.init_fingerprint != last_fingerprint: + # For alignment tasks, (re)load the model when the 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 - # Convert setting back to tuple if it was serialized as a list. - setting = ( - tuple(spec.setting) if isinstance(spec.setting, list) else spec.setting - ) - try: - _, 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, - ) - final_loss = float(loss_values[-1]) if loss_values else float("nan") + final_loss = _execute_task(spec, device, cached_model) mark_done(layout, worker_id, spec, final_loss=final_loss, device=device) print( f"[{worker_id}] Done {spec.task_id} loss={final_loss:.4f}", diff --git a/src/miss_alignment/infer.py b/src/miss_alignment/infer.py index 1474975..28d49d7 100644 --- a/src/miss_alignment/infer.py +++ b/src/miss_alignment/infer.py @@ -100,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 @@ -119,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 diff --git a/src/miss_alignment/prepare_stacks.py b/src/miss_alignment/prepare_stacks.py index 23bcc8a..f4ce53d 100644 --- a/src/miss_alignment/prepare_stacks.py +++ b/src/miss_alignment/prepare_stacks.py @@ -4,58 +4,13 @@ preprocessed tilt stacks ready for training. """ -import multiprocessing as mp -import queue -import sys -import time from pathlib import Path import mrcfile -import tqdm from warpylib import TiltSeries from warpylib.movie import Movie - -def _run_device_pool(jobs, runner, runner_args, devices, desc): - """Minimal one-process-per-GPU work queue. Internal to prepare_stacks.""" - 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 = [] - 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: - 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 +from .distributed.manager import run_distributed def _get_original_pixel_size(tilt_series: TiltSeries) -> float: @@ -133,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, @@ -150,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. @@ -188,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: @@ -203,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 d809774..174dc17 100644 --- a/src/miss_alignment/preprocessing.py +++ b/src/miss_alignment/preprocessing.py @@ -1,56 +1,9 @@ """Preprocessing utilities for tilt-series alignment.""" -import multiprocessing as mp -import queue -import sys -import time from pathlib import Path -import tqdm - from .data.io import TiltSeriesData - - -def _run_device_pool(jobs, runner, runner_args, devices, desc): - """Minimal one-process-per-GPU work queue. Internal to preprocessing.""" - 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 = [] - 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: - 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 +from .distributed.manager import run_distributed def _run_cross_correlation_single( @@ -120,33 +73,12 @@ def _run_cross_correlation_single( return float(pretilt) -def _cross_correlation_runner( - device: int | None, - task_queue, - result_queue, - lowpass_cutoff: float, - pretilt_search_range: tuple[float, float], -) -> None: - """Pull tilt-series off the queue and align them on a single device.""" - import torch - - 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, pretilt_search_range - ) - 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, pretilt_search_range: tuple[float, float] = (-30.0, 30.0), + n_cluster_workers: int | None = None, ) -> None: """ Run cross-correlation based alignment with pretilt estimation in parallel. @@ -166,6 +98,9 @@ def run_cross_correlation_alignment_parallel( Low-pass filter cutoff frequency (default: 0.25). pretilt_search_range : tuple[float, float], optional Search range for pretilt estimation in degrees (default: (-30.0, 30.0)). + 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")) @@ -184,12 +119,23 @@ 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, pretilt_search_range), - 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, + pretilt_search_range=pretilt_search_range, ) print("\nCross-correlation alignment complete!\n") diff --git a/src/miss_alignment/train.py b/src/miss_alignment/train.py index ef73313..96615bc 100644 --- a/src/miss_alignment/train.py +++ b/src/miss_alignment/train.py @@ -394,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 @@ -417,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 diff --git a/tests/distributed/test_manager.py b/tests/distributed/test_manager.py index 55b6ff2..cdbdb42 100644 --- a/tests/distributed/test_manager.py +++ b/tests/distributed/test_manager.py @@ -156,3 +156,53 @@ def test_run_distributed_cleans_up_tasks_dir(tmp_path): ) 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_worker.py b/tests/distributed/test_worker.py index 144cf0f..a74a558 100644 --- a/tests/distributed/test_worker.py +++ b/tests/distributed/test_worker.py @@ -156,3 +156,77 @@ def fake_load(path, map_location=None): 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_tilt_series.""" + _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, + pretilt_search_range=[-30.0, 30.0], + ) + write_pending(layout, spec) + + xcorr_calls = [] + + def fake_xcorr(xml_file, device, lowpass_cutoff, pretilt_search_range): + xcorr_calls.append((str(xml_file), lowpass_cutoff)) + return 2.5 # pretilt degrees + + 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() From 8d7f87a5136643fd47b05bf48bfc944a41e799d7 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 19:17:36 -0700 Subject: [PATCH 13/33] docs: add cluster_example with SLURM config, submission script, and README Co-Authored-By: Claude Fable 5 --- cluster_example/README.md | 66 +++++++++++++++++++++++++++++ cluster_example/cluster_config.json | 5 +++ cluster_example/worker.sh | 18 ++++++++ 3 files changed, 89 insertions(+) create mode 100644 cluster_example/README.md create mode 100644 cluster_example/cluster_config.json create mode 100644 cluster_example/worker.sh diff --git a/cluster_example/README.md b/cluster_example/README.md new file mode 100644 index 0000000..091c9e0 --- /dev/null +++ b/cluster_example/README.md @@ -0,0 +1,66 @@ +# 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 + +Replace the three fields in `cluster_config.json`: + +| Field | Purpose | PBS/Torque example | +|---|---|---| +| `submit` | Command to submit `{{script_path}}` | `"qsub {{script_path}}"` | +| `submit_job_id_regex` | Regex capturing the job ID from submit stdout | `"(\\d+)\\..*"` | +| `cancel` | Command to cancel `{{job_id}}` | `"qdel {{job_id}}"` | + +Update `worker.sh` with the corresponding scheduler directives (`#PBS` instead of +`#SBATCH`, etc.). diff --git a/cluster_example/cluster_config.json b/cluster_example/cluster_config.json new file mode 100644 index 0000000..0383e97 --- /dev/null +++ b/cluster_example/cluster_config.json @@ -0,0 +1,5 @@ +{ + "submit": "sbatch {{script_path}}", + "submit_job_id_regex": "Submitted batch job (\\d+)", + "cancel": "scancel {{job_id}}" +} diff --git a/cluster_example/worker.sh b/cluster_example/worker.sh new file mode 100644 index 0000000..2aaedae --- /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=slurm-%j.out +#SBATCH --error=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}} From 176919b662365c7a0bb82ce668766cdc61c2bac4 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 19:39:25 -0700 Subject: [PATCH 14/33] chore: remove worker progress noise from stderr Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/manager.py | 4 ---- src/miss_alignment/distributed/worker.py | 6 ------ 2 files changed, 10 deletions(-) diff --git a/src/miss_alignment/distributed/manager.py b/src/miss_alignment/distributed/manager.py index 7670564..c840e1d 100644 --- a/src/miss_alignment/distributed/manager.py +++ b/src/miss_alignment/distributed/manager.py @@ -59,10 +59,6 @@ def _sweep_stalled_workers(layout: QueueLayout) -> None: dest = layout.pending / task_file.name try: os.rename(task_file, dest) - print( - f"[manager] Recovered stalled task {task_file.name} to pending", - file=sys.stderr, - ) except FileNotFoundError: pass for hb in worker_dir.glob("hb-*"): diff --git a/src/miss_alignment/distributed/worker.py b/src/miss_alignment/distributed/worker.py index 1bc0ccc..17f164a 100644 --- a/src/miss_alignment/distributed/worker.py +++ b/src/miss_alignment/distributed/worker.py @@ -153,8 +153,6 @@ def run_worker_loop( if spec is None: return # queue empty, exit cleanly - print(f"[{worker_id}] Claimed {spec.task_id}", file=sys.stderr) - # For alignment tasks, (re)load the model when the fingerprint changes. if spec.task_type == "alignment" and spec.init_fingerprint != last_fingerprint: cached_model = _load_model(spec.model_checkpoint_path) @@ -163,10 +161,6 @@ def run_worker_loop( try: final_loss = _execute_task(spec, device, cached_model) mark_done(layout, worker_id, spec, final_loss=final_loss, device=device) - print( - f"[{worker_id}] Done {spec.task_id} loss={final_loss:.4f}", - file=sys.stderr, - ) except Exception: error = traceback.format_exc() mark_failed(layout, worker_id, spec, error=error) From 2951cee8b76fbc2844546fc9f8b144ea67a51c0d Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 19:49:32 -0700 Subject: [PATCH 15/33] fix: scheduler thread double-submits workers on startup last_sweep=0.0 caused the scheduler thread to call ensure_workers() immediately, racing with the explicit startup call in run_distributed and doubling the number of submitted jobs (e.g. 40 instead of 20). Initialize last_sweep=time.time() so the scheduler's first sweep fires after _SCHEDULER_INTERVAL_S, leaving the explicit startup call as the sole first submission. Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/manager.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/miss_alignment/distributed/manager.py b/src/miss_alignment/distributed/manager.py index c840e1d..e42e829 100644 --- a/src/miss_alignment/distributed/manager.py +++ b/src/miss_alignment/distributed/manager.py @@ -77,7 +77,9 @@ def _scheduler_thread( ) -> None: hb_seq = 1 # seq 0 written before thread starts last_hb = time.time() - last_sweep = 0.0 + # 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() while not stop_event.is_set(): now = time.time() From 6372bb602404aaa4d45f9719a57ea78bc33deae4 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 20:56:29 -0700 Subject: [PATCH 16/33] fix: progress bar description reflects actual task type Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/manager.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/miss_alignment/distributed/manager.py b/src/miss_alignment/distributed/manager.py index e42e829..ffe4ff3 100644 --- a/src/miss_alignment/distributed/manager.py +++ b/src/miss_alignment/distributed/manager.py @@ -197,7 +197,12 @@ def run_distributed( losses: dict[str, float] = {} failed_series: list[str] = [] - pbar = tqdm.tqdm(total=len(task_ids), desc="Tilt series alignment", file=sys.stdout) + _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) From 5a6c7314986eeb775a8269b1c972250762c83c14 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 22:08:26 -0700 Subject: [PATCH 17/33] feat: write exit reason to tasks/logs/.exit on worker shutdown Records why each worker stopped: 'queue empty', 'manager heartbeat stale', or the traceback on an unhandled exception. Includes done/failed counts. Equivalent to WarpTools' logs/.exit files. Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/queue.py | 5 ++++ src/miss_alignment/distributed/worker.py | 33 ++++++++++++++++++------ 2 files changed, 30 insertions(+), 8 deletions(-) diff --git a/src/miss_alignment/distributed/queue.py b/src/miss_alignment/distributed/queue.py index 5322d2d..b314519 100644 --- a/src/miss_alignment/distributed/queue.py +++ b/src/miss_alignment/distributed/queue.py @@ -47,6 +47,10 @@ def manager_hb(self) -> Path: 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, @@ -55,6 +59,7 @@ def ensure_directories(self) -> None: self.failed, self.manager_hb, self.cluster, + self.logs, ): d.mkdir(parents=True, exist_ok=True) diff --git a/src/miss_alignment/distributed/worker.py b/src/miss_alignment/distributed/worker.py index 17f164a..714cb7b 100644 --- a/src/miss_alignment/distributed/worker.py +++ b/src/miss_alignment/distributed/worker.py @@ -124,8 +124,11 @@ def run_worker_loop( worker_id: str, device: str, manager_hb_timeout_s: float = _MANAGER_HB_TIMEOUT_S, -) -> None: - """Main worker loop: claim → evaluate → write result. Repeat until queue empty.""" +) -> 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) @@ -133,15 +136,16 @@ def run_worker_loop( cached_model: MissAlignment | None = None hb_seq = 0 last_hb_time = 0.0 + tasks_done = 0 + tasks_failed = 0 while True: age = _manager_hb_age_s(layout) if age > manager_hb_timeout_s: - print( - f"[{worker_id}] Manager heartbeat stale ({age:.0f}s). Exiting.", - file=sys.stderr, + return ( + f"manager heartbeat stale ({age:.0f}s > {manager_hb_timeout_s:.0f}s); " + f"done={tasks_done} failed={tasks_failed}" ) - return now = time.time() if now - last_hb_time >= _HB_INTERVAL_S: @@ -151,7 +155,7 @@ def run_worker_loop( spec = claim_one(layout, worker_id) if spec is None: - return # queue empty, exit cleanly + return f"queue empty; done={tasks_done} failed={tasks_failed}" # For alignment tasks, (re)load the model when the fingerprint changes. if spec.task_type == "alignment" and spec.init_fingerprint != last_fingerprint: @@ -161,9 +165,11 @@ def run_worker_loop( 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, @@ -185,4 +191,15 @@ def worker_miss_align( layout.ensure_directories() cuda_device = f"cuda:{device}" if torch.cuda.is_available() else "cpu" - run_worker_loop(layout, worker_id, cuda_device) + 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 From c3a69a01e8ead6952cc861a9d26780d529d4b14a Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 22:14:18 -0700 Subject: [PATCH 18/33] fix: move worker heartbeat to background thread Heartbeat was written between task claims, so a single long-running series (>120s) would trigger the manager's stall sweep, deleting the worker's running/ directory. The worker would then crash writing the next heartbeat tick (FileNotFoundError), even though it had been doing useful work the whole time. Fix: _start_heartbeat_thread() ticks every 5s in a daemon thread, independent of task execution. The running/ dir is also recreated if the manager sweeps it mid-task. Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/worker.py | 107 +++++++++++++++-------- 1 file changed, 70 insertions(+), 37 deletions(-) diff --git a/src/miss_alignment/distributed/worker.py b/src/miss_alignment/distributed/worker.py index 714cb7b..4513918 100644 --- a/src/miss_alignment/distributed/worker.py +++ b/src/miss_alignment/distributed/worker.py @@ -12,6 +12,7 @@ import os import sys +import threading import time import traceback from pathlib import Path @@ -38,11 +39,36 @@ def _write_worker_hb(worker_dir: Path, seq: int) -> None: new_hb = worker_dir / f"hb-{seq}" - new_hb.write_text("") + 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-*")) @@ -132,48 +158,55 @@ def run_worker_loop( 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 - hb_seq = 0 - last_hb_time = 0.0 tasks_done = 0 tasks_failed = 0 - while True: - age = _manager_hb_age_s(layout) - if age > manager_hb_timeout_s: - return ( - f"manager heartbeat stale ({age:.0f}s > {manager_hb_timeout_s:.0f}s); " - f"done={tasks_done} failed={tasks_failed}" - ) - - now = time.time() - if now - last_hb_time >= _HB_INTERVAL_S: - _write_worker_hb(worker_dir, hb_seq) - hb_seq += 1 - last_hb_time = now - - 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 the model when the 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: + 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() - 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, - ) + # Unreachable — loop only exits via return statements above. + return f"done={tasks_done} failed={tasks_failed}" def worker_miss_align( From 8ee9766390ed5510e981944e56e3a414589255a4 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 22:17:55 -0700 Subject: [PATCH 19/33] fix: route SLURM stdout/stderr into tasks/logs/ instead of working dir Provisioner now substitutes {{logs_dir}} (tasks/logs/) and {{tasks_dir}} in the submission script template alongside {{command}}. The example worker.sh uses {{logs_dir}} for --output and --error so all slurm-*.out and slurm-*.err files land in the same directory as the worker .exit files. Co-Authored-By: Claude Fable 5 --- cluster_example/worker.sh | 4 ++-- src/miss_alignment/distributed/provisioner.py | 4 ++++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/cluster_example/worker.sh b/cluster_example/worker.sh index 2aaedae..bf8a8b7 100644 --- a/cluster_example/worker.sh +++ b/cluster_example/worker.sh @@ -7,8 +7,8 @@ #SBATCH --mem=32G #SBATCH --time=08:00:00 #SBATCH --partition={{partition}} -#SBATCH --output=slurm-%j.out -#SBATCH --error=slurm-%j.err +#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. diff --git a/src/miss_alignment/distributed/provisioner.py b/src/miss_alignment/distributed/provisioner.py index bd16cb1..04c42cc 100644 --- a/src/miss_alignment/distributed/provisioner.py +++ b/src/miss_alignment/distributed/provisioner.py @@ -85,7 +85,11 @@ def _render_script(self, index: int) -> Path: 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() From a279e4527f49f142c69fc873b42448bef6f67edf Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sat, 4 Jul 2026 22:31:02 -0700 Subject: [PATCH 20/33] fix: stall sweep skips tasks already in done/ or failed/ When a long-running series caused the old (pre-background-thread) sweep to re-pend a task that the worker subsequently completed, the task ended up in both pending/ and done/ simultaneously. Now the sweep checks for a matching file in done/ or failed/ before re-pending, and discards the running/ copy if found. Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/manager.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/miss_alignment/distributed/manager.py b/src/miss_alignment/distributed/manager.py index ffe4ff3..20676c1 100644 --- a/src/miss_alignment/distributed/manager.py +++ b/src/miss_alignment/distributed/manager.py @@ -56,6 +56,13 @@ def _sweep_stalled_workers(layout: QueueLayout) -> None: 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) From 05552f74c5240a3026a03f1ac522e6aa2c3805d2 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sun, 5 Jul 2026 08:04:31 -0700 Subject: [PATCH 21/33] fix: clarify pool configuration log message '12 workers -> 8 partitions' implied the workers were divided into partitions. The arrow is misleading: reconstruction workers and partitions are sized independently. Use a comma instead. Co-Authored-By: Claude Fable 5 --- src/miss_alignment/data/training_datamodule.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/miss_alignment/data/training_datamodule.py b/src/miss_alignment/data/training_datamodule.py index 7a36f76..f20c3db 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)" ) From ef5b212852fc54bb4133f65d1463d5bc20391fcd Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sun, 5 Jul 2026 08:05:50 -0700 Subject: [PATCH 22/33] fix: TOCTOU race in stall sweep reading heartbeat mtime MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit max(...).stat().st_mtime called stat() twice on the same path: once inside the key function to find the newest tick, and again on the result to read its mtime. The heartbeat thread could delete the file between those two calls (hb-43 → hb-44), causing FileNotFoundError in the scheduler thread and crashing the whole run. Fix: iterate ticks once, catching FileNotFoundError per-tick, and track the newest mtime seen without re-stating the winning file. Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/manager.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/src/miss_alignment/distributed/manager.py b/src/miss_alignment/distributed/manager.py index 20676c1..5534ef3 100644 --- a/src/miss_alignment/distributed/manager.py +++ b/src/miss_alignment/distributed/manager.py @@ -44,11 +44,16 @@ def _sweep_stalled_workers(layout: QueueLayout) -> None: if not worker_dir.is_dir(): continue ticks = list(worker_dir.glob("hb-*")) - if ticks: - age = ( - time.time() - - max(ticks, key=lambda p: p.stat().st_mtime).stat().st_mtime - ) + 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 From 9759a5094c7548b1b9f7bf1e05a8957fbcdd1680 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sun, 5 Jul 2026 08:09:28 -0700 Subject: [PATCH 23/33] fix: surface scheduler thread crashes to the main thread An unhandled exception in _scheduler_thread previously killed it silently while the poll loop hung forever waiting for tasks to complete. Now: - _scheduler_thread catches all exceptions, stores them in error_box, and sets stop_event to wake the poll loop immediately - The poll loop checks error_box each iteration and re-raises as RuntimeError with the original exception chained Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/manager.py | 35 +++++++++++++++-------- 1 file changed, 23 insertions(+), 12 deletions(-) diff --git a/src/miss_alignment/distributed/manager.py b/src/miss_alignment/distributed/manager.py index 5534ef3..4b6ed44 100644 --- a/src/miss_alignment/distributed/manager.py +++ b/src/miss_alignment/distributed/manager.py @@ -86,6 +86,7 @@ def _scheduler_thread( provisioner: WorkerProvisioner, n_workers: int, stop_event: threading.Event, + error_box: list, ) -> None: hb_seq = 1 # seq 0 written before thread starts last_hb = time.time() @@ -93,20 +94,24 @@ def _scheduler_thread( # run_distributed calls ensure_workers() explicitly at startup instead. last_sweep = time.time() - while not stop_event.is_set(): - now = 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_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) - provisioner.ensure_workers(n_workers) - last_sweep = now + if now - last_sweep >= _SCHEDULER_INTERVAL_S: + _sweep_stalled_workers(layout) + provisioner.ensure_workers(n_workers) + last_sweep = now - stop_event.wait(timeout=1.0) + 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( @@ -197,9 +202,10 @@ def run_distributed( n_workers = len(devices) stop_event = threading.Event() + scheduler_errors: list = [] scheduler = threading.Thread( target=_scheduler_thread, - args=(layout, provisioner, n_workers, stop_event), + args=(layout, provisioner, n_workers, stop_event, scheduler_errors), daemon=True, ) scheduler.start() @@ -219,6 +225,11 @@ def run_distributed( while pending_ids: time.sleep(_POLL_INTERVAL_S) + if scheduler_errors: + raise RuntimeError( + "Scheduler thread crashed" + ) from scheduler_errors[0] + for done_file in layout.done.glob("*.json"): data = json.loads(done_file.read_text()) tid = data["task_id"] From c946b99ffd8e7b5b26fecbba2d44ac36848492be Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Sun, 5 Jul 2026 22:14:26 -0700 Subject: [PATCH 24/33] Remove spec and plan for worker pool feature --- .../plans/2026-07-04-distributed-inference.md | 2099 ----------------- ...2026-07-04-distributed-inference-design.md | 279 --- 2 files changed, 2378 deletions(-) delete mode 100644 docs/superpowers/plans/2026-07-04-distributed-inference.md delete mode 100644 docs/superpowers/specs/2026-07-04-distributed-inference-design.md diff --git a/docs/superpowers/plans/2026-07-04-distributed-inference.md b/docs/superpowers/plans/2026-07-04-distributed-inference.md deleted file mode 100644 index c20119f..0000000 --- a/docs/superpowers/plans/2026-07-04-distributed-inference.md +++ /dev/null @@ -1,2099 +0,0 @@ -# Distributed Inference Implementation Plan - -> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. - -**Goal:** Add a disk-based task queue that distributes per-tilt-series alignment inference across cluster nodes (SLURM/PBS), while keeping the existing local multi-GPU behaviour unchanged. - -**Architecture:** A new `miss_alignment/distributed/` package implements a filesystem queue (atomic rename as the claim mutex), a manager that writes tasks and blocks until done, two provisioners (local subprocess and cluster batch scheduler), and a `miss-alignment worker` subcommand. Each worker processes many series per run; checkpoint loading is amortized by passing the resident model into `evaluate_tilt_series`. Cluster mode is triggered by a new `--n-cluster-workers` CLI arg; without it local multi-GPU mode is unchanged. - -**Tech Stack:** Python 3.10+, stdlib only for `distributed/` (`pathlib`, `os`, `subprocess`, `threading`, `hashlib`, `json`, `re`, `time`, `signal`). `tqdm` (already a dependency) used in manager. `typer` (already a dependency) used in worker CLI. - -## Global Constraints - -- Python ≥ 3.10 (project minimum). -- No new runtime dependencies beyond stdlib + existing deps. -- `ruff check --fix && ruff format` must pass (line length 88, ignore E712). -- `pytest --color=yes` must pass with no warnings-as-errors regressions. -- All new queue infrastructure under `src/miss_alignment/distributed/`. -- Cluster mode activated only by `--n-cluster-workers N`; without it local mode runs unchanged. -- `MISS_CLUSTER_CONFIG` and `MISS_CLUSTER_SCRIPT` are **required** when `--n-cluster-workers` is set; `config.py` raises `RuntimeError` if either is absent (never silently falls back to local mode). -- `evaluate_tilt_series` gains one optional `model` parameter and is otherwise unchanged. -- `train.py` and `infer.py` are modified only to add the `--n-cluster-workers` option and pass it through. - ---- - -## File Map - -| Path | Action | Responsibility | -|---|---|---| -| `src/miss_alignment/__main__.py` | Create | `python -m miss_alignment` entry point for `LocalProvisioner` subprocess launch | -| `src/miss_alignment/distributed/__init__.py` | Create | Public re-exports | -| `src/miss_alignment/distributed/queue.py` | Create | Directory layout, task JSON, atomic rename claim | -| `src/miss_alignment/distributed/manager.py` | Create | Head-node coordinator, scheduler thread, poll loop | -| `src/miss_alignment/distributed/provisioner.py` | Create | `WorkerProvisioner` ABC, `LocalProvisioner`, `ClusterProvisioner` | -| `src/miss_alignment/distributed/worker.py` | Create | `miss-alignment worker` subcommand logic | -| `src/miss_alignment/distributed/config.py` | Create | Read env vars, return `ClusterConfig` or raise | -| `src/miss_alignment/alignment/tilt_series.py` | Modify | Add optional `model` parameter to `evaluate_tilt_series` | -| `src/miss_alignment/alignment/parallel.py` | Modify | Replace `run_device_pool` call with manager; add `n_cluster_workers` param | -| `src/miss_alignment/train.py` | Modify | Add `--n-cluster-workers` option; pass to `run_alignment_parallel` | -| `src/miss_alignment/infer.py` | Modify | Add `--n-cluster-workers` option; pass to `run_alignment_parallel` | -| `src/miss_alignment/_cli.py` | Modify | Register `worker` subcommand | -| `src/miss_alignment/__init__.py` | Modify | Export `worker_miss_align` | -| `src/miss_alignment/_parallel.py` | Delete (Task 7) | Superseded by `LocalProvisioner` | -| `tests/distributed/__init__.py` | Create | Test package | -| `tests/distributed/test_queue.py` | Create | Queue layer unit tests | -| `tests/distributed/test_manager.py` | Create | Manager integration tests | -| `tests/distributed/test_worker.py` | Create | Worker claim loop unit tests | -| `tests/distributed/test_config.py` | Create | Config env-var parsing tests | -| `tests/distributed/test_provisioner.py` | Create | Provisioner unit tests | -| `tests/test_parallel.py` | Modify | Update to test new `LocalProvisioner` path | - ---- - -## Task 1: Queue layer (`distributed/queue.py`) - -**Files:** -- Create: `src/miss_alignment/__main__.py` -- Create: `src/miss_alignment/distributed/__init__.py` -- Create: `src/miss_alignment/distributed/queue.py` -- Create: `tests/distributed/__init__.py` -- Create: `tests/distributed/test_queue.py` - -**Interfaces:** -- Produces: - - `QueueLayout(root: Path)` — dataclass; `ensure_directories() -> None`; properties: `pending`, `running`, `done`, `failed`, `manager_hb`, `cluster` each returning `Path`; `worker_dir(worker_id: str) -> Path`. - - `TaskSpec` — dataclass with fields: `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`. - - `compute_fingerprint(model_checkpoint_path, setting, patch_size, patch_overlap, batch_size, apply_ctf, downsample) -> str` — SHA-256 hex. - - `write_pending(layout: QueueLayout, spec: TaskSpec) -> None` - - `claim_one(layout: QueueLayout, worker_id: str) -> TaskSpec | None` - - `mark_done(layout: QueueLayout, worker_id: str, spec: TaskSpec, final_loss: float, device: str) -> None` - - `mark_failed(layout: QueueLayout, worker_id: str, spec: TaskSpec, error: str) -> None` - - `clear_queue(layout: QueueLayout) -> None` — deletes `pending/done/failed` contents first, then recovers `running/` orphans into the now-empty `pending/`. - -- [ ] **Step 1: Create directory skeleton** - -```bash -mkdir -p tests/distributed -touch tests/distributed/__init__.py -``` - -- [ ] **Step 2: Write failing tests** - -```python -# tests/distributed/test_queue.py -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 -``` - -- [ ] **Step 3: Run tests to verify they fail** - -```bash -cd /Users/tegunovd/dev/miss-alignment -pytest tests/distributed/test_queue.py -v 2>&1 | head -20 -``` - -Expected: `ModuleNotFoundError: No module named 'miss_alignment.distributed'` - -- [ ] **Step 4: Create `__main__.py` and `distributed/__init__.py` skeleton** - -```python -# src/miss_alignment/__main__.py -from miss_alignment import cli - -cli() -``` - -```python -# src/miss_alignment/distributed/__init__.py -"""Disk-based distributed task queue for miss-alignment inference.""" -``` - -- [ ] **Step 5: Implement `queue.py`** - -```python -# src/miss_alignment/distributed/queue.py -"""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" - - def ensure_directories(self) -> None: - for d in ( - self.pending, - self.running, - self.done, - self.failed, - self.manager_hb, - self.cluster, - ): - 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 - - -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 -``` - -- [ ] **Step 6: Run tests to verify they pass** - -```bash -pytest tests/distributed/test_queue.py -v -``` - -Expected: all 10 tests PASS. - -- [ ] **Step 7: Commit** - -```bash -git add src/miss_alignment/__main__.py \ - src/miss_alignment/distributed/__init__.py \ - src/miss_alignment/distributed/queue.py \ - tests/distributed/__init__.py \ - tests/distributed/test_queue.py -git commit -m "feat: add distributed queue layer with atomic rename claim protocol" -``` - ---- - -## Task 2: Cluster configuration (`distributed/config.py`) - -**Files:** -- Create: `src/miss_alignment/distributed/config.py` -- Create: `tests/distributed/test_config.py` - -**Interfaces:** -- Produces: - - `ClusterConfig` — dataclass with fields: `submit: str`, `submit_job_id_regex: str`, `cancel: str`, `script_path: Path`. - - `load_cluster_config() -> ClusterConfig` — reads `MISS_CLUSTER_CONFIG` and `MISS_CLUSTER_SCRIPT`; raises `RuntimeError` if either is unset, `FileNotFoundError` if a path doesn't exist, `KeyError` if a required JSON key is missing. - -- [ ] **Step 1: Write failing tests** - -```python -# tests/distributed/test_config.py -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}}", - } - 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 - - -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() -``` - -- [ ] **Step 2: Run tests to verify they fail** - -```bash -pytest tests/distributed/test_config.py -v 2>&1 | head -20 -``` - -Expected: `ImportError` for `miss_alignment.distributed.config`. - -- [ ] **Step 3: Implement `config.py`** - -```python -# src/miss_alignment/distributed/config.py -"""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 -from pathlib import Path - - -@dataclass -class ClusterConfig: - submit: str - submit_job_id_regex: str - cancel: str - script_path: Path - - -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', and 'cancel' 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"], - script_path=script_path, - ) -``` - -- [ ] **Step 4: Run tests to verify they pass** - -```bash -pytest tests/distributed/test_config.py -v -``` - -Expected: all 5 tests PASS. - -- [ ] **Step 5: Commit** - -```bash -git add src/miss_alignment/distributed/config.py tests/distributed/test_config.py -git commit -m "feat: add cluster config reader (raises if env vars missing)" -``` - ---- - -## Task 3: Worker subcommand (`distributed/worker.py`) - -**Files:** -- Create: `src/miss_alignment/distributed/worker.py` -- Create: `tests/distributed/test_worker.py` - -**Interfaces:** -- Consumes: - - `QueueLayout`, `TaskSpec`, `claim_one`, `mark_done`, `mark_failed` from `distributed/queue.py` - - `MissAlignment` from `..models.models` - - `evaluate_tilt_series` from `..alignment.tilt_series` (after Task 4 adds the `model` param) -- Produces: - - `run_worker_loop(layout, worker_id, device, manager_hb_timeout_s) -> None` — testable loop body. - - `worker_miss_align(queue_dir, device, worker_id)` — Typer command. - -- [ ] **Step 1: Write failing tests** - -```python -# tests/distributed/test_worker.py -"""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 -``` - -- [ ] **Step 2: Run tests to verify they fail** - -```bash -pytest tests/distributed/test_worker.py -v 2>&1 | head -20 -``` - -Expected: `ImportError` for `miss_alignment.distributed.worker`. - -- [ ] **Step 3: Implement `worker.py`** - -```python -# src/miss_alignment/distributed/worker.py -"""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 time -import traceback -from pathlib import Path - -import torch -import typer - -from ..alignment.tilt_series import evaluate_tilt_series -from ..models.models import MissAlignment -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}" - new_hb.write_text("") - if seq > 0: - (worker_dir / f"hb-{seq - 1}").unlink(missing_ok=True) - - -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 run_worker_loop( - layout: QueueLayout, - worker_id: str, - device: str, - manager_hb_timeout_s: float = _MANAGER_HB_TIMEOUT_S, -) -> None: - """Main worker loop: claim → evaluate → write result. Repeat until queue empty.""" - worker_dir = layout.worker_dir(worker_id) - worker_dir.mkdir(parents=True, exist_ok=True) - - last_fingerprint: str | None = None - cached_model: MissAlignment | None = None - hb_seq = 0 - last_hb_time = 0.0 - - while True: - age = _manager_hb_age_s(layout) - if age > manager_hb_timeout_s: - print( - f"[{worker_id}] Manager heartbeat stale ({age:.0f}s). Exiting.", - file=sys.stderr, - ) - return - - now = time.time() - if now - last_hb_time >= _HB_INTERVAL_S: - _write_worker_hb(worker_dir, hb_seq) - hb_seq += 1 - last_hb_time = now - - spec = claim_one(layout, worker_id) - if spec is None: - return # queue empty, exit cleanly - - print(f"[{worker_id}] Claimed {spec.task_id}", file=sys.stderr) - - if spec.init_fingerprint != last_fingerprint: - cached_model = _load_model(spec.model_checkpoint_path) - last_fingerprint = spec.init_fingerprint - - # Convert setting back to tuple if it was serialized as a list. - setting = ( - tuple(spec.setting) if isinstance(spec.setting, list) else spec.setting - ) - - try: - _, 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, - ) - final_loss = float(loss_values[-1]) if loss_values else float("nan") - mark_done(layout, worker_id, spec, final_loss=final_loss, device=device) - print( - f"[{worker_id}] Done {spec.task_id} loss={final_loss:.4f}", - file=sys.stderr, - ) - except Exception: - error = traceback.format_exc() - mark_failed(layout, worker_id, spec, error=error) - print( - f"[{worker_id}] Failed {spec.task_id}:\n{error}", - file=sys.stderr, - ) - - -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: str | None = 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" - run_worker_loop(layout, worker_id, cuda_device) -``` - -- [ ] **Step 4: Run tests to verify they pass** - -```bash -pytest tests/distributed/test_worker.py -v -``` - -Expected: all 5 tests PASS. - -- [ ] **Step 5: Commit** - -```bash -git add src/miss_alignment/distributed/worker.py tests/distributed/test_worker.py -git commit -m "feat: add worker subcommand with model fingerprint reuse across tasks" -``` - ---- - -## Task 4: Add `model` parameter to `evaluate_tilt_series` - -**Files:** -- Modify: `src/miss_alignment/alignment/tilt_series.py` - -**Interfaces:** -- Produces: `evaluate_tilt_series(..., model: MissAlignment | None = None)` — when `model` is provided, uses it directly instead of loading from disk. All existing callers unaffected. - -- [ ] **Step 1: Read the current model-loading block** - -Open `src/miss_alignment/alignment/tilt_series.py` at lines 118–186. The relevant section is: - -```python -# line 131 ends the signature -) -> tuple[Path, list[float]]: - ... - # line 172 - model = MissAlignment.load_from_checkpoint( - model_checkpoint_path, - map_location="cpu", - ) - # line 184 - if hasattr(model.net, "_orig_mod"): - model.net = model.net._orig_mod -``` - -- [ ] **Step 2: Add the `model` parameter to the signature** - -In `src/miss_alignment/alignment/tilt_series.py`, add `model` as the last parameter before the closing `)`: - -```python -# Before: - n_control_points: int = 7, -) -> tuple[Path, list[float]]: - -# After: - n_control_points: int = 7, - model: "MissAlignment | None" = None, -) -> tuple[Path, list[float]]: -``` - -Use a string annotation to avoid a circular import (the type is already imported at the top of the file — verify with `grep "MissAlignment" src/miss_alignment/alignment/tilt_series.py` before committing; if already imported, use the bare type). - -- [ ] **Step 3: Replace the model-loading block** - -```python -# Before (lines ~172-185): - model = MissAlignment.load_from_checkpoint( - model_checkpoint_path, - map_location="cpu", - ) - if hasattr(model.net, "_orig_mod"): - model.net = model.net._orig_mod - -# After: - if model is None: - model = MissAlignment.load_from_checkpoint( - model_checkpoint_path, - map_location="cpu", - ) - if hasattr(model.net, "_orig_mod"): - model.net = model.net._orig_mod -``` - -- [ ] **Step 4: Run the existing alignment tests to verify nothing broke** - -```bash -pytest tests/alignment/ -v -``` - -Expected: all tests PASS. - -- [ ] **Step 5: Commit** - -```bash -git add src/miss_alignment/alignment/tilt_series.py -git commit -m "feat: add optional model param to evaluate_tilt_series for checkpoint reuse" -``` - ---- - -## Task 5: Provisioners (`distributed/provisioner.py`) - -**Files:** -- Create: `src/miss_alignment/distributed/provisioner.py` -- Create: `tests/distributed/test_provisioner.py` - -**Interfaces:** -- Consumes: `ClusterConfig` from `distributed/config.py`. -- Produces: - - `WorkerProvisioner` — ABC with `ensure_workers(n_workers: int) -> None` and `shutdown() -> None`. - - `LocalProvisioner(queue_dir: Path, devices: list[int])` — spawns `python -m miss_alignment worker` child processes, one per device. Ignores `n_workers`. - - `ClusterProvisioner(queue_dir: Path, config: ClusterConfig)` — submits exactly `n_workers` cluster jobs on the first `ensure_workers` call. - -- [ ] **Step 1: Write failing tests** - -```python -# tests/distributed/test_provisioner.py -import sys -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, LocalProvisioner - - -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 - # verify device args - all_args = [str(c) for c in mock_popen.call_args_list] - assert any("'0'" in s or '"0"' in s or "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) # should respawn because poll() != None - - 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() - - -def test_cluster_provisioner_submits_n_workers_jobs(tmp_path): - script = tmp_path / "worker.sh" - script.write_text("#!/bin/bash\n{{command}}\n") - cfg = ClusterConfig( - submit="sbatch {{script_path}}", - submit_job_id_regex=r"Submitted batch job (\d+)", - cancel="scancel {{job_id}}", - script_path=script, - ) - - submitted = [] - - def fake_run(cmd, **kwargs): - submitted.append(cmd) - result = MagicMock() - result.stdout = "Submitted batch job 12345\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(submitted) == 4 - - -def test_cluster_provisioner_cancels_on_shutdown(tmp_path): - script = tmp_path / "worker.sh" - script.write_text("#!/bin/bash\n{{command}}\n") - cfg = ClusterConfig( - submit="sbatch {{script_path}}", - submit_job_id_regex=r"Submitted batch job (\d+)", - cancel="scancel {{job_id}}", - script_path=script, - ) - - cancel_calls = [] - - def fake_run(cmd, **kwargs): - if "sbatch" in cmd: - result = MagicMock() - result.stdout = "Submitted batch job 99999\n" - 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) -``` - -- [ ] **Step 2: Run tests to verify they fail** - -```bash -pytest tests/distributed/test_provisioner.py -v 2>&1 | head -20 -``` - -Expected: `ImportError` for `miss_alignment.distributed.provisioner`. - -- [ ] **Step 3: Implement `provisioner.py`** - -```python -# src/miss_alignment/distributed/provisioner.py -"""Worker provisioners: spawn local child processes or submit cluster jobs.""" - -from __future__ import annotations - -import os -import re -import subprocess -import sys -from abc import ABC, abstractmethod -from pathlib import Path - -from .config import ClusterConfig - - -class WorkerProvisioner(ABC): - @abstractmethod - def ensure_workers(self, n_workers: int) -> None: - """Ensure workers are running. Called once at startup and each scheduler tick.""" - - @abstractmethod - def shutdown(self) -> None: - """Terminate all managed workers.""" - - -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] = {} - - def ensure_workers(self, n_workers: int) -> None: - 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() - - -class ClusterProvisioner(WorkerProvisioner): - """Submits exactly n_workers cluster jobs, each running until the queue drains.""" - - def __init__(self, queue_dir: Path, config: ClusterConfig) -> None: - self._queue_dir = queue_dir - self._config = config - self._job_ids: list[str] = [] - self._scripts_dir = queue_dir / "cluster" - self._scripts_dir.mkdir(parents=True, exist_ok=True) - - 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}"' - ) - rendered = template_text.replace("{{command}}", command) - 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 ensure_workers(self, n_workers: int) -> None: - already = len(self._job_ids) - for i in range(already, n_workers): - script_path = self._render_script(i) - 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_ids.append(match.group(1)) - - def shutdown(self) -> None: - for job_id in self._job_ids: - cancel_cmd = self._config.cancel.replace("{{job_id}}", job_id) - subprocess.run(cancel_cmd, shell=True, stderr=subprocess.DEVNULL) - self._job_ids.clear() -``` - -- [ ] **Step 4: Run tests to verify they pass** - -```bash -pytest tests/distributed/test_provisioner.py -v -``` - -Expected: all 6 tests PASS. - -- [ ] **Step 5: Commit** - -```bash -git add src/miss_alignment/distributed/provisioner.py tests/distributed/test_provisioner.py -git commit -m "feat: add LocalProvisioner and ClusterProvisioner" -``` - ---- - -## Task 6: Manager (`distributed/manager.py`) - -**Files:** -- Create: `src/miss_alignment/distributed/manager.py` -- Create: `tests/distributed/test_manager.py` - -**Interfaces:** -- Consumes: - - `QueueLayout`, `TaskSpec`, `write_pending`, `clear_queue`, `compute_fingerprint` from `distributed/queue.py` - - `WorkerProvisioner` from `distributed/provisioner.py` -- Produces: - - `run_distributed(tilt_series_list, model_checkpoint, output_directory, setting, patch_size, patch_overlap, batch_size, apply_ctf, downsample, devices, n_cluster_workers, queue_root) -> dict[str, float]` - -- [ ] **Step 1: Write failing tests** - -```python -# tests/distributed/test_manager.py -"""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 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() -``` - -- [ ] **Step 2: Run tests to verify they fail** - -```bash -pytest tests/distributed/test_manager.py -v 2>&1 | head -20 -``` - -Expected: `ImportError` for `miss_alignment.distributed.manager`. - -- [ ] **Step 3: Implement `manager.py`** - -```python -# src/miss_alignment/distributed/manager.py -"""Head-node coordinator: writes tasks, starts provisioner, blocks until done.""" - -from __future__ import annotations - -import json -import shutil -import sys -import threading -import time -from pathlib import Path - -import tqdm - -from .config import load_cluster_config -from .provisioner import ClusterProvisioner, 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-*")) - if ticks: - age = time.time() - max(ticks, key=lambda p: p.stat().st_mtime).stat().st_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"): - dest = layout.pending / task_file.name - try: - import os - os.rename(task_file, dest) - print( - f"[manager] Recovered stalled task {task_file.name} to pending", - file=sys.stderr, - ) - except FileNotFoundError: - pass - for hb in worker_dir.glob("hb-*"): - hb.unlink(missing_ok=True) - try: - worker_dir.rmdir() - except OSError: - pass - - -def _scheduler_thread( - layout: QueueLayout, - provisioner: WorkerProvisioner, - n_workers: int, - stop_event: threading.Event, -) -> None: - hb_seq = 1 # seq 0 written before thread starts - last_hb = time.time() - last_sweep = 0.0 - - 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) - provisioner.ensure_workers(n_workers) - last_sweep = now - - stop_event.wait(timeout=1.0) - - -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, -) -> 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. - """ - layout = QueueLayout(queue_root) - layout.ensure_directories() - clear_queue(layout) - - 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, - ) - - 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), - 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, - ) - write_pending(layout, spec) - - # Write the first manager heartbeat before starting workers so workers - # never see a missing heartbeat on startup. - _write_manager_hb(layout, seq=0) - - if n_cluster_workers is not None: - cluster_config = load_cluster_config() - provisioner: WorkerProvisioner = ClusterProvisioner( - queue_dir=queue_root, config=cluster_config - ) - n_workers = n_cluster_workers - else: - provisioner = LocalProvisioner(queue_dir=queue_root, devices=devices) - n_workers = len(devices) - - stop_event = threading.Event() - scheduler = threading.Thread( - target=_scheduler_thread, - args=(layout, provisioner, n_workers, stop_event), - daemon=True, - ) - scheduler.start() - provisioner.ensure_workers(n_workers) - - pending_ids = set(task_ids) - losses: dict[str, float] = {} - failed_series: list[str] = [] - - pbar = tqdm.tqdm(total=len(task_ids), desc="Tilt series alignment", file=sys.stdout) - try: - while pending_ids: - time.sleep(_POLL_INTERVAL_S) - - 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 -``` - -- [ ] **Step 4: Run all distributed tests** - -```bash -pytest tests/distributed/ -v -``` - -Expected: all tests PASS. - -- [ ] **Step 5: Commit** - -```bash -git add src/miss_alignment/distributed/manager.py tests/distributed/test_manager.py -git commit -m "feat: add distributed manager with scheduler thread, poll loop, and cleanup" -``` - ---- - -## Task 7: Wire up CLI, parallel, train, infer; delete `_parallel.py` - -**Files:** -- Modify: `src/miss_alignment/alignment/parallel.py` -- Modify: `src/miss_alignment/train.py` -- Modify: `src/miss_alignment/infer.py` -- Modify: `src/miss_alignment/_cli.py` -- Modify: `src/miss_alignment/__init__.py` -- Modify: `src/miss_alignment/distributed/__init__.py` -- Modify: `tests/test_parallel.py` -- Delete: `src/miss_alignment/_parallel.py` - -**Interfaces:** -- `run_alignment_parallel` gains `n_cluster_workers: int | None = None` parameter. -- `train_miss_align` and `infer_miss_align` each gain `n_cluster_workers: Optional[int] = typer.Option(None, ...)`. - -- [ ] **Step 1: Update `tests/test_parallel.py`** - -Replace the file entirely to test through the new `LocalProvisioner` path: - -```python -# tests/test_parallel.py -"""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.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() -``` - -- [ ] **Step 2: Run updated `test_parallel.py` to verify it passes now** - -```bash -pytest tests/test_parallel.py -v -``` - -Expected: all 3 tests PASS. - -- [ ] **Step 3: Update `alignment/parallel.py`** - -```python -# src/miss_alignment/alignment/parallel.py -from pathlib import Path - -from ..distributed.manager import run_distributed - - -def run_alignment_parallel( - model_checkpoint: Path, - tilt_series_list: list[Path], - output_directory: Path, - setting: str | tuple[int, int] | tuple[int, int, int, int], - patch_size: int, - patch_overlap: float, - batch_size: int, - apply_ctf: bool, - downsample: int, - devices_list: list[int], - n_cluster_workers: int | None = None, -) -> dict[str, float]: - """Distribute per-tilt-series alignment across local GPUs or a cluster. - - 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 mapping tilt-series stem names to their final loss values. - """ - 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, - n_cluster_workers=n_cluster_workers, - queue_root=queue_root, - ) -``` - -- [ ] **Step 4: Add `--n-cluster-workers` to `train.py`** - -Add the new option to `train_miss_align`'s signature. Find the `preprocess: bool` option (last existing option, around line 308) and add after it: - -```python - 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.", - ), -``` - -Then pass it through to `run_alignment_parallel` in the call at line ~487: - -```python - run_alignment_parallel( - model_checkpoint=str(training_model_path), - tilt_series_list=tilt_series_list, - output_directory=training_directory, - setting=iteration_settings["alignment"], - patch_size=alignment_config["patch_size"], - patch_overlap=alignment_config["patch_overlap"], - batch_size=alignment_config["batch_size"], - apply_ctf=general_config["apply_ctf"], - downsample=iteration_settings["downsample"], - devices_list=devices_alignment, - n_cluster_workers=n_cluster_workers, - ) -``` - -- [ ] **Step 5: Add `--n-cluster-workers` to `infer.py`** - -Add the same option to `infer_miss_align`'s signature after `preprocess: bool` (around line 41): - -```python - 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.", - ), -``` - -Then pass it through to the `run_alignment_parallel` call at line ~161: - -```python - run_alignment_parallel( - model_checkpoint=str(model_checkpoint), - tilt_series_list=tilt_series_list, - output_directory=data_directory, - setting=iteration_settings["alignment"], - patch_size=alignment_config["patch_size"], - patch_overlap=alignment_config["patch_overlap"], - batch_size=alignment_config["batch_size"], - apply_ctf=general_config["apply_ctf"], - downsample=iteration_settings["downsample"], - devices_list=devices_alignment, - n_cluster_workers=n_cluster_workers, - ) -``` - -- [ ] **Step 6: Register the `worker` subcommand in `_cli.py`** - -```python -# src/miss_alignment/_cli.py -from click import Context -import typer -from typer.core import TyperGroup - - -class OrderCommands(TyperGroup): - def list_commands(self, ctx: Context): - """Return list of commands in the order appear.""" - return list(self.commands) - - -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) -``` - -- [ ] **Step 7: Export `worker_miss_align` from `__init__.py`** - -```python -# src/miss_alignment/__init__.py -"""She has a chaotic good alignment for tilt-series.""" - -from importlib.metadata import PackageNotFoundError, version - -try: - __version__ = version("miss_alignment") -except PackageNotFoundError: - __version__ = "uninstalled" - -__author__ = "Marten Chaillet" -__email__ = "martenchaillet@gmail.com" -__all__ = [ - "__version__", - "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 -``` - -- [ ] **Step 8: Update `distributed/__init__.py`** - -```python -# src/miss_alignment/distributed/__init__.py -"""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, 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", - "LocalProvisioner", - "WorkerProvisioner", - "QueueLayout", - "TaskSpec", - "claim_one", - "clear_queue", - "compute_fingerprint", - "mark_done", - "mark_failed", - "write_pending", - "run_worker_loop", - "worker_miss_align", -] -``` - -- [ ] **Step 9: Delete `_parallel.py`** - -```bash -git rm src/miss_alignment/_parallel.py -``` - -- [ ] **Step 10: Run full test suite and linter** - -```bash -pytest --color=yes -v -ruff check --fix src/miss_alignment/ -ruff format src/miss_alignment/ -``` - -Expected: all tests PASS. If any test imports `miss_alignment._parallel` directly, fix that import (the module no longer exists). No ruff errors. - -- [ ] **Step 11: Verify the worker subcommand is registered** - -```bash -miss-alignment --help -miss-alignment worker --help -``` - -Expected: `worker` appears in the command list; `--queue-dir`, `--device`, `--worker-id` appear in its help. - -- [ ] **Step 12: Commit** - -```bash -git add \ - src/miss_alignment/alignment/parallel.py \ - src/miss_alignment/train.py \ - src/miss_alignment/infer.py \ - src/miss_alignment/_cli.py \ - src/miss_alignment/__init__.py \ - src/miss_alignment/distributed/__init__.py \ - tests/test_parallel.py -git commit -m "feat: wire distributed queue into run_alignment_parallel; add --n-cluster-workers; register worker subcommand; delete _parallel.py" -``` diff --git a/docs/superpowers/specs/2026-07-04-distributed-inference-design.md b/docs/superpowers/specs/2026-07-04-distributed-inference-design.md deleted file mode 100644 index d77214a..0000000 --- a/docs/superpowers/specs/2026-07-04-distributed-inference-design.md +++ /dev/null @@ -1,279 +0,0 @@ -# Distributed Inference Design - -**Date:** 2026-07-04 -**Status:** Approved (revised) -**Scope:** Cluster distribution of the per-tilt-series alignment/inference phase - ---- - -## Problem - -With 300 tilt-series, the alignment phase (inference) takes 3× longer than the preceding training phase. Each series is independently optimizable — no data exchange between series — but the current implementation is limited to GPUs on the local machine. The head node has to finish all series sequentially across its local GPU pool before the next macro-iteration can begin. - -## Goal - -Distribute the per-series alignment tasks across a compute cluster so the head node fans out work, blocks until all series are done, and then continues to the next macro-iteration. Cluster mode is opt-in via a new `--n-cluster-workers` CLI argument; without it the system runs exactly as today (local multi-GPU pool). - ---- - -## Architecture - -A new `miss_alignment/distributed/` module sits between the existing `alignment/parallel.py` public API and the underlying `evaluate_tilt_series` function. - -`run_alignment_parallel` accepts a new `n_cluster_workers: int | None` parameter. When it is `None`, `LocalProvisioner` is used (one subprocess per GPU, unchanged behavior). When it is set, `ClusterProvisioner` is used and submits exactly `n_cluster_workers` jobs; each worker runs the claim loop and processes as many series as it can grab until the queue drains. - -`evaluate_tilt_series` gains an optional `model: MissAlignment | None = None` parameter. When provided, the worker's resident model is used directly instead of loading from disk — this amortizes checkpoint loading across all series a worker processes in one macro-iteration. Callers that omit the parameter are unaffected. - -### Components - -| File | Role | -|---|---| -| `distributed/queue.py` | Queue directory layout, task JSON read/write, atomic rename claim | -| `distributed/manager.py` | Head-node coordinator: writes tasks, runs scheduler thread, blocks until done | -| `distributed/provisioner.py` | `WorkerProvisioner` interface + `LocalProvisioner` + `ClusterProvisioner` | -| `distributed/worker.py` | `miss-alignment worker` subcommand: claims tasks, runs inference, writes results | -| `distributed/config.py` | Reads env vars, returns `ClusterConfig` dataclass or raises if missing | - ---- - -## Task Format - -One JSON file per tilt-series, written to `/tasks/pending/-.json`: - -```json -{ - "task_id": "0000003-tilt_series_01", - "model_checkpoint_path": "/data/project/iter2/model.ckpt", - "tilt_series_path": "/data/project/tilt_series_01.xml", - "output_directory": "/data/project/", - "setting": "anchoring", - "patch_size": 96, - "patch_overlap": 0.1, - "batch_size": 32, - "apply_ctf": false, - "downsample": 2, - "init_fingerprint": "" -} -``` - -`init_fingerprint` covers the model checkpoint path and all alignment parameters. A worker skips reloading the model when consecutive tasks share the same fingerprint, amortizing checkpoint loading across the many series each worker processes in one macro-iteration. - -`setting` is serialized to JSON as a string or array. Workers convert array back to `tuple` before passing to `evaluate_tilt_series`. - -On completion, result fields are appended before writing to `done/`: -```json -{ "final_loss": 0.0312, "device": "cuda:0" } -``` - -On failure, error fields are appended before writing to `failed/`: -```json -{ "error": "CUDA out of memory...", "worker_id": "cluster-node01-12345-0" } -``` - ---- - -## Queue Directory Layout - -``` -/tasks/ -├── pending/ # one JSON per queued task -├── running/ -│ └── / # per-worker subdir -│ ├── .json # claimed task lives here during execution -│ └── hb- # worker heartbeat tick files (latest only) -├── done/ # completed task JSONs -├── failed/ # failed task JSONs -├── manager/ -│ └── hb- # manager heartbeat tick files (latest only) -└── cluster/ - └── worker-.sh # rendered cluster submission scripts -``` - -The queue directory is always `/tasks/`. Since the training directory must already be on a shared filesystem (workers need the XML/MRC files), no additional configuration is needed for cluster nodes to access it. - ---- - -## Claim Protocol - -1. Worker lists `pending/`, **shuffles randomly** (avoids thundering herd on the same lexicographically-first file) -2. Attempts `os.rename(pending/.json, running//.json)` -3. `FileNotFoundError` means another worker won the race — try the next candidate -4. On task completion: write `done/.json`, then delete `running//.json` (publish-before-delete: a crash mid-step leaves an orphan in `running/` that the scheduler sweeps, not a lost task) -5. On task failure: write `failed/.json` with error, delete running copy, continue - -No lock files. The OS rename is the coordination primitive. - ---- - -## Worker Lifecycle - -**Startup** (`miss-alignment worker --queue-dir --device [--worker-id ]`): -1. Derive `worker_id` from `--worker-id` or `local--gpu` -2. Create `tasks/running//` -3. Start background heartbeat thread: write `tasks/running//hb-` every 5s, keep only the latest tick file -4. Enter claim loop - -**Claim loop:** -- Check manager heartbeat (`tasks/manager/hb-*`): if latest tick is >120s old → exit cleanly (manager is dead) -- List `pending/`, shuffle, attempt rename -- On empty queue: exit cleanly -- On claim: check if `init_fingerprint` matches last loaded model; if not, load checkpoint from `model_checkpoint_path` and cache it -- Call `evaluate_tilt_series(..., model=cached_model, device=f"cuda:{device}")` — all internal LBFGS retries and optimization passes run unchanged -- Write result to `done/` or `failed/`, delete running copy, loop back - -Each worker processes **many series** per run. The checkpoint is loaded only once per macro-iteration (all tasks in a phase share the same `init_fingerprint`). - -**No task-level retries.** Failures are deterministic (bad data, corrupt file, config error); the manager hard-fails after all remaining tasks complete. - ---- - -## Manager Lifecycle - -**Startup:** -1. Clear stale queue state from any prior run: first delete `pending/`, `done/`, `failed/` contents; then recover any `running/` orphans back to `pending/` -2. Write all task JSONs to `pending/` -3. Write the first manager heartbeat tick immediately (before starting workers, to avoid a race where workers start and find no heartbeat) -4. Start scheduler background thread -5. Start `ClusterProvisioner` or `LocalProvisioner`, call `ensure_workers(n_workers)` -6. Block in poll loop (500ms interval) until all tasks are in `done/` or `failed/` - -**Scheduler thread** (runs every ~10s): -- Write manager heartbeat to `tasks/manager/hb-` -- Sweep stalled workers: for each `running//`, check age of latest `hb-` file; if >120s stale → move task back to `pending/`, clean up dir -- Call `provisioner.ensure_workers()` to respawn any dead local workers - -**Shutdown** (on completion or `KeyboardInterrupt`): -- Call `provisioner.shutdown()` (SIGTERM local children / cancel SLURM job IDs) -- Delete `tasks/` directory (clean state for next macro-iteration) -- Return `dict[series_name → final_loss]` on full success, or raise `RuntimeError` listing all failed series if any task is in `failed/` - -**Failure policy:** if any series ends in `failed/`, the manager raises after all other tasks complete. `train.py` and `infer.py` propagate this as a hard failure — training stops. - ---- - -## Provisioners - -### `LocalProvisioner` - -Activated when `n_cluster_workers` is `None` (default). - -- `ensure_workers(n_workers)`: spawns `miss-alignment worker --queue-dir --device ` via `subprocess.Popen`, one per GPU in `devices_alignment` (same list as today: `list(range(torch.cuda.device_count()))`, respecting `CUDA_VISIBLE_DEVICES`). `n_workers` is ignored; the number of local workers is always equal to the number of devices. -- Respawns any that have exited prematurely (checked each scheduler tick) -- `shutdown()`: SIGTERM all children, short timeout, SIGKILL if needed - -This replaces `_parallel.py`'s `run_device_pool` / `mp.spawn` internals with no behavioral change for the local case. - -### `ClusterProvisioner` - -Activated when `n_cluster_workers` is set. - -Requires both `MISS_CLUSTER_CONFIG` and `MISS_CLUSTER_SCRIPT` env vars to be set; raises `RuntimeError` immediately if either is absent, so the user gets a clear error rather than a silent fallback to local mode. - -**`MISS_CLUSTER_CONFIG`** — path to a JSON file: -```json -{ - "submit": "sbatch {{script_path}}", - "submit_job_id_regex": "Submitted batch job (\\d+)", - "cancel": "scancel {{job_id}}" -} -``` - -**`MISS_CLUSTER_SCRIPT`** — path to a `.sh` template: -```bash -#!/bin/bash -#SBATCH --nodes=1 -#SBATCH --gres=gpu:1 -#SBATCH --time=04:00:00 - -conda activate miss-alignment -{{command}} -``` - -- `{{command}}` is filled by the provisioner with the `miss-alignment worker` invocation -- All cluster-specific settings (partition, memory, time limit, environment setup) are the user's responsibility in the template -- Additional `{{custom_var}}` placeholders filled via `MISS_CLUSTER_VAR_=` env vars -- Rendered scripts are written to `tasks/cluster/worker-.sh` -- The injected `{{command}}` is: `miss-alignment worker --queue-dir --device 0 --worker-id "$(hostname)-$$-"` — `$(hostname)` and `$$` are expanded by the compute node's shell at runtime, guaranteeing unique, stable worker IDs -- `ensure_workers(n_workers)` submits exactly `n_workers` jobs; each worker pulls tasks from the shared queue until it drains -- `shutdown()`: runs configured `cancel` command for each stored job ID; registers SIGINT/SIGTERM handlers so Ctrl-C on the head node cancels the cluster pool - ---- - -## Configuration - -### CLI changes - -`miss-alignment train` and `miss-alignment infer` each gain one new optional argument: - -| Argument | Type | Default | Meaning | -|---|---|---|---| -| `--n-cluster-workers` | `int \| None` | `None` | Number of cluster jobs to submit. When set, activates cluster mode. When absent, local multi-GPU mode is used (unchanged). | - -### Environment variables (cluster mode only) - -| Env var | Purpose | -|---|---| -| `MISS_CLUSTER_CONFIG` | Path to cluster scheduler JSON config. **Required** when `--n-cluster-workers` is set. | -| `MISS_CLUSTER_SCRIPT` | Path to job submission shell script template. **Required** when `--n-cluster-workers` is set. | -| `MISS_CLUSTER_VAR_` | Additional template variables, e.g. `MISS_CLUSTER_VAR_partition=gpu`. | - -`config.py` is called only when `n_cluster_workers` is not `None`. It raises `RuntimeError` if either env var is absent, with a message naming which one is missing. - ---- - -## Changes to `evaluate_tilt_series` - -A single optional parameter is added: - -```python -def evaluate_tilt_series( - model_checkpoint_path: Path, - tilt_series_path: Path, - output_directory: Path, - setting: str | tuple[int, int] | tuple[int, int, int, int] = "anchoring", - patch_size: int = 96, - patch_overlap: float = 0.1, - batch_size: int = 16, - apply_ctf: bool = True, - downsample: int = 1, - device: str = "cpu", - initial_reliable_fraction: float = 1 / 2, - n_control_points: int = 7, - model: MissAlignment | None = None, # NEW -) -> tuple[Path, list[float]]: -``` - -When `model` is provided, the function uses it directly instead of loading from `model_checkpoint_path`. This is the only change to the function. All existing callers (which omit `model`) are unaffected. - ---- - -## What Does Not Change - -- `evaluate_tilt_series` internal logic — only the signature gains one optional parameter -- LBFGS retries, anchoring iterations, coarse-to-fine spline passes — all internal to `evaluate_tilt_series` -- `CUDA_VISIBLE_DEVICES` still controls which local GPUs are used -- Single-GPU local behavior is identical - ---- - -## File Changes Summary - -**New files:** -- `src/miss_alignment/distributed/__init__.py` -- `src/miss_alignment/distributed/queue.py` -- `src/miss_alignment/distributed/manager.py` -- `src/miss_alignment/distributed/provisioner.py` -- `src/miss_alignment/distributed/worker.py` -- `src/miss_alignment/distributed/config.py` -- `src/miss_alignment/__main__.py` — enables `python -m miss_alignment` for `LocalProvisioner` subprocess launch - -**Modified files:** -- `src/miss_alignment/alignment/tilt_series.py` — add optional `model` parameter to `evaluate_tilt_series` -- `src/miss_alignment/alignment/parallel.py` — replace `run_device_pool` call with distributed manager; add `n_cluster_workers` parameter -- `src/miss_alignment/train.py` — accept and pass through `--n-cluster-workers` -- `src/miss_alignment/infer.py` — accept and pass through `--n-cluster-workers` -- `src/miss_alignment/_cli.py` — register `worker` subcommand -- `src/miss_alignment/__init__.py` — export `worker_miss_align` - -**Deleted files:** -- `src/miss_alignment/_parallel.py` — superseded by `LocalProvisioner`; deleted once the new path is validated in tests From c941af31aa6b2096e706418eb3a970f2fe3187c2 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Mon, 6 Jul 2026 13:07:08 -0700 Subject: [PATCH 25/33] Get torch-projectors from PyPI in CI --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4bf3ace..5f07b0a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -46,7 +46,7 @@ jobs: - name: Install torch-projectors (CPU) run: | - python -m pip install torch-projectors --index-url https://warpem.github.io/torch-projectors/cpu/simple/ + python -m pip install torch-projectors - name: Install package run: python -m pip install -e .[test] From 2001f9db0975bcd00913c432cf0e4303eaed8f29 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Wed, 8 Jul 2026 11:29:31 -0700 Subject: [PATCH 26/33] feat: also launch local GPU workers when cluster mode is active Local GPUs were idle during cluster-distributed alignment phases because only cluster jobs were submitted. Since all workers share the same disk-based task queue regardless of whether they run locally or remotely, there is no coordination overhead in mixing them. When --n-cluster-workers is set and local GPU devices are available, a CompositeProvisioner now runs both ClusterProvisioner and LocalProvisioner simultaneously. Local workers claim tasks from the same queue as cluster workers and are respawned by the scheduler on exit, just like in local-only mode. Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/__init__.py | 8 ++++++- src/miss_alignment/distributed/manager.py | 18 ++++++++++++---- src/miss_alignment/distributed/provisioner.py | 19 +++++++++++++++++ tests/distributed/test_provisioner.py | 21 ++++++++++++++++++- 4 files changed, 60 insertions(+), 6 deletions(-) diff --git a/src/miss_alignment/distributed/__init__.py b/src/miss_alignment/distributed/__init__.py index 9d2d21d..1b05278 100644 --- a/src/miss_alignment/distributed/__init__.py +++ b/src/miss_alignment/distributed/__init__.py @@ -2,7 +2,12 @@ from .config import ClusterConfig, load_cluster_config from .manager import run_distributed -from .provisioner import ClusterProvisioner, LocalProvisioner, WorkerProvisioner +from .provisioner import ( + ClusterProvisioner, + CompositeProvisioner, + LocalProvisioner, + WorkerProvisioner, +) from .queue import ( QueueLayout, TaskSpec, @@ -20,6 +25,7 @@ "load_cluster_config", "run_distributed", "ClusterProvisioner", + "CompositeProvisioner", "LocalProvisioner", "WorkerProvisioner", "QueueLayout", diff --git a/src/miss_alignment/distributed/manager.py b/src/miss_alignment/distributed/manager.py index af3194d..660a7e2 100644 --- a/src/miss_alignment/distributed/manager.py +++ b/src/miss_alignment/distributed/manager.py @@ -13,7 +13,12 @@ import tqdm from .config import load_cluster_config -from .provisioner import ClusterProvisioner, LocalProvisioner, WorkerProvisioner +from .provisioner import ( + ClusterProvisioner, + CompositeProvisioner, + LocalProvisioner, + WorkerProvisioner, +) from .queue import ( QueueLayout, TaskSpec, @@ -189,9 +194,14 @@ def run_distributed( if n_cluster_workers is not None: cluster_config = load_cluster_config() - provisioner: WorkerProvisioner = ClusterProvisioner( - queue_dir=queue_root, config=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) diff --git a/src/miss_alignment/distributed/provisioner.py b/src/miss_alignment/distributed/provisioner.py index 04c42cc..5ad0eed 100644 --- a/src/miss_alignment/distributed/provisioner.py +++ b/src/miss_alignment/distributed/provisioner.py @@ -122,3 +122,22 @@ def shutdown(self) -> None: cancel_cmd = self._config.cancel.replace("{{job_id}}", job_id) subprocess.run(cancel_cmd, shell=True, stderr=subprocess.DEVNULL) self._job_ids.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 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/tests/distributed/test_provisioner.py b/tests/distributed/test_provisioner.py index f06f574..3c91aac 100644 --- a/tests/distributed/test_provisioner.py +++ b/tests/distributed/test_provisioner.py @@ -5,7 +5,11 @@ import pytest from miss_alignment.distributed.config import ClusterConfig -from miss_alignment.distributed.provisioner import ClusterProvisioner, LocalProvisioner +from miss_alignment.distributed.provisioner import ( + ClusterProvisioner, + CompositeProvisioner, + LocalProvisioner, +) def test_local_provisioner_spawns_one_process_per_device(tmp_path): @@ -118,3 +122,18 @@ def fake_run(cmd, **kwargs): assert len(cancel_calls) == 3 assert all("scancel" in c for c in cancel_calls) + + +def test_composite_provisioner_delegates_to_all(tmp_path): + """CompositeProvisioner calls ensure_workers and shutdown on every child.""" + 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() From 2d3db1892f98cda1bc6d85a97fcf4b8751b4ff17 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Wed, 8 Jul 2026 11:58:47 -0700 Subject: [PATCH 27/33] feat: replenish cluster worker pool on preemption; show live worker counts in progress bar ClusterProvisioner.ensure_workers() now submits against the live worker count (running/ dirs with a fresh heartbeat) rather than the submitted- job-ids list, so preempted jobs are automatically resubmitted each scheduler tick. A startup grace period (60s) prevents newly-created worker dirs from being misread as stale. The scheduler gates replenishment on n_pending > 0 and caps at min(n_workers, n_pending). This prevents spurious resubmissions when workers exit cleanly on an empty queue: a clean exit leaves nothing in pending/, so ensure_workers is never called. Only if the stall sweep re-pends an orphaned task does n_pending go positive again. _live_worker_dirs() iterates heartbeat tick mtimes with per-file FileNotFoundError handling, fixing the TOCTOU race where the heartbeat thread rotates a tick between glob and stat. All provisioners expose live_worker_count() and worker_counts_by_type(). CompositeProvisioner sums across its children and merges the type dicts. The manager poll loop updates the tqdm postfix each scheduler tick with live counts e.g. 'cluster=87 local=2', so pool depletion is visible without checking squeue. Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/manager.py | 20 +++- src/miss_alignment/distributed/provisioner.py | 112 +++++++++++++++--- tests/distributed/test_manager.py | 6 + tests/distributed/test_provisioner.py | 4 +- 4 files changed, 121 insertions(+), 21 deletions(-) diff --git a/src/miss_alignment/distributed/manager.py b/src/miss_alignment/distributed/manager.py index 660a7e2..f2c25b3 100644 --- a/src/miss_alignment/distributed/manager.py +++ b/src/miss_alignment/distributed/manager.py @@ -86,12 +86,17 @@ def _sweep_stalled_workers(layout: QueueLayout) -> None: 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() @@ -110,7 +115,13 @@ def _scheduler_thread( if now - last_sweep >= _SCHEDULER_INTERVAL_S: _sweep_stalled_workers(layout) - provisioner.ensure_workers(n_workers) + 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.update(provisioner.worker_counts_by_type()) last_sweep = now stop_event.wait(timeout=1.0) @@ -209,9 +220,12 @@ def run_distributed( stop_event = threading.Event() scheduler_errors: list = [] + worker_counts: dict = provisioner.worker_counts_by_type() scheduler = threading.Thread( target=_scheduler_thread, - args=(layout, provisioner, n_workers, stop_event, scheduler_errors), + args=( + layout, provisioner, n_workers, stop_event, scheduler_errors, worker_counts + ), daemon=True, ) scheduler.start() @@ -234,6 +248,8 @@ def run_distributed( 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"] diff --git a/src/miss_alignment/distributed/provisioner.py b/src/miss_alignment/distributed/provisioner.py index 5ad0eed..f52c1b1 100644 --- a/src/miss_alignment/distributed/provisioner.py +++ b/src/miss_alignment/distributed/provisioner.py @@ -6,11 +6,49 @@ import re import subprocess import sys +import time from abc import ABC, abstractmethod 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 @@ -24,6 +62,14 @@ def ensure_workers(self, n_workers: int) -> None: 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.""" @@ -33,6 +79,12 @@ def __init__(self, queue_dir: Path, devices: list[int]) -> None: self._devices = devices self._procs: dict[int, subprocess.Popen] = {} + def live_worker_count(self) -> int: + 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: for device in self._devices: proc = self._procs.get(device) @@ -67,15 +119,24 @@ def shutdown(self) -> None: class ClusterProvisioner(WorkerProvisioner): - """Submits exactly n_workers cluster jobs, each running until the queue drains.""" + """Submits cluster jobs; resubmits when alive workers fall below target.""" def __init__(self, queue_dir: Path, config: ClusterConfig) -> None: self._queue_dir = queue_dir self._config = config self._job_ids: list[str] = [] + # Monotonically increasing index for unique script names. Separate + # from _job_ids so we never reuse a script index even after replenishment. + self._next_index: int = 0 self._scripts_dir = queue_dir / "cluster" self._scripts_dir.mkdir(parents=True, exist_ok=True) + def live_worker_count(self) -> int: + return _live_worker_dirs(self._queue_dir / "running") + + def worker_counts_by_type(self) -> dict[str, int]: + return {"cluster": self.live_worker_count()} + 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. @@ -99,23 +160,30 @@ def _render_script(self, index: int) -> Path: 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_ids.append(match.group(1)) + def ensure_workers(self, n_workers: int) -> None: - already = len(self._job_ids) - for i in range(already, n_workers): - script_path = self._render_script(i) - 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_ids.append(match.group(1)) + """Submit new jobs until live worker count reaches n_workers.""" + alive = self.live_worker_count() + deficit = n_workers - alive + for _ in range(deficit): + self._submit_one() def shutdown(self) -> None: for job_id in self._job_ids: @@ -134,6 +202,16 @@ class CompositeProvisioner(WorkerProvisioner): 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) diff --git a/tests/distributed/test_manager.py b/tests/distributed/test_manager.py index cdbdb42..50a70c0 100644 --- a/tests/distributed/test_manager.py +++ b/tests/distributed/test_manager.py @@ -63,6 +63,12 @@ def ensure_workers(self, n_workers): 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") diff --git a/tests/distributed/test_provisioner.py b/tests/distributed/test_provisioner.py index 3c91aac..0a20b37 100644 --- a/tests/distributed/test_provisioner.py +++ b/tests/distributed/test_provisioner.py @@ -86,7 +86,7 @@ def fake_run(cmd, **kwargs): with patch( "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run - ): + ), patch.object(ClusterProvisioner, "live_worker_count", return_value=0): p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) p.ensure_workers(n_workers=4) @@ -115,7 +115,7 @@ def fake_run(cmd, **kwargs): with patch( "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run - ): + ), patch.object(ClusterProvisioner, "live_worker_count", return_value=0): p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) p.ensure_workers(n_workers=3) p.shutdown() From 2cddbaf1d561001639480ec1a9c5d21dde46731d Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Wed, 8 Jul 2026 13:20:07 -0700 Subject: [PATCH 28/33] fix: parallelize cluster job cancellation on shutdown ClusterProvisioner.shutdown() was cancelling jobs serially, one scancel subprocess call per job ID. With preemption-triggered resubmissions accumulating over a long run this could mean dozens of sequential 1-2s round-trips to the scheduler, blocking the manager after the progress bar hit 100%. Use ThreadPoolExecutor to fire all cancel commands concurrently. Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/provisioner.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/miss_alignment/distributed/provisioner.py b/src/miss_alignment/distributed/provisioner.py index f52c1b1..3d91c56 100644 --- a/src/miss_alignment/distributed/provisioner.py +++ b/src/miss_alignment/distributed/provisioner.py @@ -8,6 +8,7 @@ import sys import time from abc import ABC, abstractmethod +from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path from .config import ClusterConfig @@ -186,9 +187,12 @@ def ensure_workers(self, n_workers: int) -> None: self._submit_one() def shutdown(self) -> None: - for job_id in self._job_ids: + 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) + + with ThreadPoolExecutor(max_workers=min(32, len(self._job_ids) or 1)) as pool: + list(as_completed([pool.submit(_cancel, jid) for jid in self._job_ids])) self._job_ids.clear() From 4858828f2b9e5610f58d371f95897e57590e167d Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Wed, 8 Jul 2026 13:44:57 -0700 Subject: [PATCH 29/33] feat: scheduler-aware job liveness via status_list command (SLURM/LSF/PBS/SGE) Replaces heartbeat-based liveness with a proper scheduler status query. ClusterConfig gains a 'status_list' command (replacing 'status' + 'status_job_id_regex') that lists all active user jobs in 'id,STATUS' format. The recommended SLURM command is: squeue -u $USER -h -o "%i,%T" _parse_status_output() maps status tokens to alive/terminal across four schedulers: SLURM, LSF, PBS, SGE. The scheduler is auto-detected from output tokens unless 'scheduler' is set explicitly in the config. A 'custom' mode accepts user-defined alive status strings. This correctly handles jobs in PENDING state (queued but not yet started), which the heartbeat approach could not see -- preventing the runaway resubmission that produced >3000 cluster jobs for a pool of 160. cluster_example/cluster_config.json updated to use status_list. cluster_example/README.md adds per-scheduler examples for SLURM/LSF/PBS/custom. Co-Authored-By: Claude Fable 5 --- cluster_example/README.md | 70 ++++++- cluster_example/cluster_config.json | 3 +- src/miss_alignment/distributed/config.py | 15 +- src/miss_alignment/distributed/provisioner.py | 142 ++++++++++++- tests/distributed/test_config.py | 19 ++ tests/distributed/test_provisioner.py | 187 +++++++++++++++--- 6 files changed, 394 insertions(+), 42 deletions(-) diff --git a/cluster_example/README.md b/cluster_example/README.md index 091c9e0..db3cb18 100644 --- a/cluster_example/README.md +++ b/cluster_example/README.md @@ -54,13 +54,63 @@ automatically. ## Adapting for other schedulers -Replace the three fields in `cluster_config.json`: - -| Field | Purpose | PBS/Torque example | -|---|---|---| -| `submit` | Command to submit `{{script_path}}` | `"qsub {{script_path}}"` | -| `submit_job_id_regex` | Regex capturing the job ID from submit stdout | `"(\\d+)\\..*"` | -| `cancel` | Command to cancel `{{job_id}}` | `"qdel {{job_id}}"` | - -Update `worker.sh` with the corresponding scheduler directives (`#PBS` instead of -`#SBATCH`, etc.). +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 index 0383e97..15fe8c7 100644 --- a/cluster_example/cluster_config.json +++ b/cluster_example/cluster_config.json @@ -1,5 +1,6 @@ { "submit": "sbatch {{script_path}}", "submit_job_id_regex": "Submitted batch job (\\d+)", - "cancel": "scancel {{job_id}}" + "cancel": "scancel {{job_id}}", + "status_list": "squeue -u $USER -h -o \"%i,%T\"" } diff --git a/src/miss_alignment/distributed/config.py b/src/miss_alignment/distributed/config.py index 219ebd0..fe48341 100644 --- a/src/miss_alignment/distributed/config.py +++ b/src/miss_alignment/distributed/config.py @@ -9,7 +9,7 @@ import json import os -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path @@ -18,7 +18,14 @@ 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: @@ -28,7 +35,7 @@ def load_cluster_config() -> ClusterConfig: 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', and 'cancel' keys." + "'submit', 'submit_job_id_regex', 'cancel', and 'status_list' keys." ) script_path_str = os.environ.get("MISS_CLUSTER_SCRIPT") @@ -52,5 +59,9 @@ def load_cluster_config() -> 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/provisioner.py b/src/miss_alignment/distributed/provisioner.py index 3d91c56..a9ff4f8 100644 --- a/src/miss_alignment/distributed/provisioner.py +++ b/src/miss_alignment/distributed/provisioner.py @@ -119,21 +119,148 @@ def shutdown(self) -> None: 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 that mean the job is still occupying a slot (queued or running). +_ALIVE_STATUSES: dict[str, set[str]] = { + "slurm": {"PENDING", "PD", "RUNNING", "R", "COMPLETING", "CG", + "RESIZING", "SUSPENDED", "S"}, + "lsf": {"PEND", "RUN", "SSUSP", "USUSP", "PSUSP"}, + "pbs": {"Q", "R", "E", "H", "W", "T", "S"}, + "sge": {"qw", "r", "t", "Rr", "Rq", "hqw", "hRwq"}, +} + + +def _parse_status_output( + output: str, + our_job_ids: set[str], + scheduler: str, + custom_alive_statuses: list[str], + custom_status_regex: str, +) -> set[str]: + """Parse status_list output and return the subset of our_job_ids still alive. + + Each line is expected to contain a job ID and a status token separated by + a comma (matching the recommended status_list format "%i,%T" for SLURM). + Falls back to just checking if our job ID appears anywhere on a line for + schedulers whose status format varies. + """ + alive: set[str] = set() + + 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: + # No comma — treat the whole line as a job ID with unknown status. + job_id = line + status_token = "" + + if job_id not in our_job_ids: + continue + + if scheduler == "custom": + if status_token in custom_alive_statuses: + alive.add(job_id) + elif not status_token: + alive.add(job_id) # presence without status → treat as alive + continue + + if not status_token: + # Present in output but no status parsed → assume alive. + alive.add(job_id) + continue + + for sched in schedulers_to_try: + if status_token in _ALIVE_STATUSES[sched]: + alive.add(job_id) + break + else: + # Status token found but not in any alive set → job is terminal. + pass + + return alive + + class ClusterProvisioner(WorkerProvisioner): - """Submits cluster jobs; resubmits when alive workers fall below target.""" + """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. + + Recommended status_list format (one job per line, id,STATUS): + SLURM: squeue -u $USER -h -o "%i,%T" + LSF: bjobs -noheader -o "jobid stat" (with custom_status_regex) + PBS: qstat -u $USER (with custom parser) + SGE: qstat -u $USER (with custom parser) + """ def __init__(self, queue_dir: Path, config: ClusterConfig) -> None: self._queue_dir = queue_dir self._config = config self._job_ids: list[str] = [] - # Monotonically increasing index for unique script names. Separate - # from _job_ids so we never reuse a script index even after replenishment. + # 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) + def _alive_job_ids(self) -> set[str]: + """Query the scheduler and return the subset of our job IDs still alive. + + On scheduler error, returns the full submitted-ID set to avoid + resubmission storms during transient scheduler outages. + """ + if not self._job_ids: + return set() + our_ids = set(self._job_ids) + 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, + ) + alive = _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, + ) + # Prune terminated jobs from our tracking list. + self._job_ids = [jid for jid in self._job_ids if jid in alive] + return alive + except Exception: + # Scheduler unavailable — assume all submitted jobs are still alive. + return our_ids + def live_worker_count(self) -> int: - return _live_worker_dirs(self._queue_dir / "running") + return len(self._alive_job_ids()) def worker_counts_by_type(self) -> dict[str, int]: return {"cluster": self.live_worker_count()} @@ -180,8 +307,11 @@ def _submit_one(self) -> None: self._job_ids.append(match.group(1)) def ensure_workers(self, n_workers: int) -> None: - """Submit new jobs until live worker count reaches n_workers.""" - alive = self.live_worker_count() + """Submit new jobs until alive job count reaches n_workers. + + Alive means queued or running according to the cluster scheduler. + """ + alive = len(self._alive_job_ids()) deficit = n_workers - alive for _ in range(deficit): self._submit_one() diff --git a/tests/distributed/test_config.py b/tests/distributed/test_config.py index db033ab..1fa005a 100644 --- a/tests/distributed/test_config.py +++ b/tests/distributed/test_config.py @@ -10,6 +10,7 @@ def cluster_json(tmp_path): "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)) @@ -45,6 +46,24 @@ def test_load_cluster_config_returns_config(monkeypatch, cluster_json, cluster_s 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): diff --git a/tests/distributed/test_provisioner.py b/tests/distributed/test_provisioner.py index 0a20b37..a43935c 100644 --- a/tests/distributed/test_provisioner.py +++ b/tests/distributed/test_provisioner.py @@ -9,9 +9,94 @@ 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_alive(): + output = "12345,PENDING\n12346,RUNNING\n12347,COMPLETING\n" + our_ids = {"12345", "12346", "12347", "99999"} + alive = _parse_status_output(output, our_ids, "slurm", [], "") + assert alive == {"12345", "12346", "12347"} + + +def test_parse_status_output_slurm_terminal(): + output = "12345,COMPLETED\n12346,FAILED\n12347,CANCELLED\n" + our_ids = {"12345", "12346", "12347"} + alive = _parse_status_output(output, our_ids, "slurm", [], "") + assert alive == set() + + +def test_parse_status_output_lsf_alive(): + output = "12345,PEND\n12346,RUN\n" + our_ids = {"12345", "12346"} + alive = _parse_status_output(output, our_ids, "lsf", [], "") + assert alive == {"12345", "12346"} + + +def test_parse_status_output_pbs_alive(): + output = "12345,Q\n12346,R\n12347,C\n" + our_ids = {"12345", "12346", "12347"} + alive = _parse_status_output(output, our_ids, "pbs", [], "") + assert alive == {"12345", "12346"} + + +def test_parse_status_output_sge_alive(): + output = "12345,qw\n12346,r\n" + our_ids = {"12345", "12346"} + alive = _parse_status_output(output, our_ids, "sge", [], "") + assert alive == {"12345", "12346"} + + +def test_parse_status_output_auto_detects_slurm(): + output = "12345,PENDING\n12346,RUNNING\n" + our_ids = {"12345", "12346"} + alive = _parse_status_output(output, our_ids, "auto", [], "") + assert alive == {"12345", "12346"} + + +def test_parse_status_output_ignores_unknown_ids(): + output = "99999,RUNNING\n" + our_ids = {"12345"} + alive = _parse_status_output(output, our_ids, "slurm", [], "") + assert alive == set() + + +def test_parse_status_output_no_status_token_treated_as_alive(): + """A job present in output with no comma is assumed alive.""" + output = "12345\n" + our_ids = {"12345"} + alive = _parse_status_output(output, our_ids, "slurm", [], "") + assert alive == {"12345"} + + +def test_parse_status_output_custom(): + output = "12345,INPROGRESS\n12346,DONE\n" + our_ids = {"12345", "12346"} + alive = _parse_status_output(output, our_ids, "custom", ["INPROGRESS"], "") + assert alive == {"12345"} + + +# --------------------------------------------------------------------------- +# 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() @@ -22,7 +107,6 @@ def test_local_provisioner_spawns_one_process_per_device(tmp_path): p.ensure_workers(n_workers=10) assert mock_popen.call_count == 3 - # verify device args appear in calls all_args = [str(c) for c in mock_popen.call_args_list] assert any("0" in s for s in all_args) @@ -48,7 +132,7 @@ def test_local_provisioner_respawns_dead_worker(tmp_path): p = LocalProvisioner(queue_dir=tmp_path, devices=[0]) p.ensure_workers(n_workers=5) - p.ensure_workers(n_workers=5) # should respawn because poll() != None + p.ensure_workers(n_workers=5) assert mock_popen.call_count == 2 @@ -66,16 +150,12 @@ def test_local_provisioner_shutdown_terminates(tmp_path): mock_proc.terminate.assert_called() -def test_cluster_provisioner_submits_n_workers_jobs(tmp_path): - script = tmp_path / "worker.sh" - script.write_text("#!/bin/bash\n{{command}}\n") - cfg = ClusterConfig( - submit="sbatch {{script_path}}", - submit_job_id_regex=r"Submitted batch job (\d+)", - cancel="scancel {{job_id}}", - script_path=script, - ) +# --------------------------------------------------------------------------- +# ClusterProvisioner tests +# --------------------------------------------------------------------------- +def test_cluster_provisioner_submits_n_workers_jobs(tmp_path): + cfg = _cluster_cfg(tmp_path) submitted = [] def fake_run(cmd, **kwargs): @@ -86,23 +166,57 @@ def fake_run(cmd, **kwargs): with patch( "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run - ), patch.object(ClusterProvisioner, "live_worker_count", return_value=0): + ), patch.object(ClusterProvisioner, "_alive_job_ids", return_value=set()): p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) p.ensure_workers(n_workers=4) assert len(submitted) == 4 -def test_cluster_provisioner_cancels_on_shutdown(tmp_path): - script = tmp_path / "worker.sh" - script.write_text("#!/bin/bash\n{{command}}\n") - cfg = ClusterConfig( - submit="sbatch {{script_path}}", - submit_job_id_regex=r"Submitted batch job (\d+)", - cancel="scancel {{job_id}}", - script_path=script, - ) +def test_cluster_provisioner_does_not_resubmit_alive_jobs(tmp_path): + cfg = _cluster_cfg(tmp_path) + submitted = [] + + def fake_run(cmd, **kwargs): + submitted.append(cmd) + result = MagicMock() + result.stdout = "Submitted batch job 12345\n" + return result + + with patch( + "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run + ), patch.object( + ClusterProvisioner, "_alive_job_ids", return_value={"1", "2", "3", "4"} + ): + p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) + p.ensure_workers(n_workers=4) + + assert len(submitted) == 0 + + +def test_cluster_provisioner_replenishes_preempted_jobs(tmp_path): + cfg = _cluster_cfg(tmp_path) + submitted = [] + + def fake_run(cmd, **kwargs): + submitted.append(cmd) + result = MagicMock() + result.stdout = "Submitted batch job 99999\n" + return result + + with patch( + "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run + ), patch.object( + ClusterProvisioner, "_alive_job_ids", return_value={"1", "2"} + ): + p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) + p.ensure_workers(n_workers=4) + assert len(submitted) == 2 + + +def test_cluster_provisioner_cancels_on_shutdown(tmp_path): + cfg = _cluster_cfg(tmp_path) cancel_calls = [] def fake_run(cmd, **kwargs): @@ -115,7 +229,7 @@ def fake_run(cmd, **kwargs): with patch( "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run - ), patch.object(ClusterProvisioner, "live_worker_count", return_value=0): + ), patch.object(ClusterProvisioner, "_alive_job_ids", return_value=set()): p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) p.ensure_workers(n_workers=3) p.shutdown() @@ -124,8 +238,35 @@ def fake_run(cmd, **kwargs): assert all("scancel" in c for c in cancel_calls) +def test_cluster_provisioner_prunes_terminated_jobs(tmp_path): + """_alive_job_ids prunes job IDs no longer in scheduler output.""" + cfg = _cluster_cfg(tmp_path) + + # Simulate: 3 jobs submitted, only 2 alive in scheduler + status_output = "11111,PENDING\n22222,RUNNING\n" + + def fake_run(cmd, **kwargs): + result = MagicMock() + result.stdout = status_output + result.returncode = 0 + return result + + with patch( + "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run + ): + p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) + p._job_ids = ["11111", "22222", "33333"] + alive = p._alive_job_ids() + + assert alive == {"11111", "22222"} + assert p._job_ids == ["11111", "22222"] + + +# --------------------------------------------------------------------------- +# CompositeProvisioner tests +# --------------------------------------------------------------------------- + def test_composite_provisioner_delegates_to_all(tmp_path): - """CompositeProvisioner calls ensure_workers and shutdown on every child.""" a = MagicMock(spec=LocalProvisioner) b = MagicMock(spec=LocalProvisioner) p = CompositeProvisioner([a, b]) From c6e7c46c4fc4497ee26efec51a94e3ea7f80e3b0 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Wed, 8 Jul 2026 14:40:51 -0700 Subject: [PATCH 30/33] fix: split cluster worker count into running vs pending in progress bar worker_counts_by_type() now returns 'cluster-running' and 'cluster-pending' separately so the progress bar shows e.g. 'cluster-running=3 cluster-pending=157' instead of a misleading 'cluster=160'. _parse_status_output() returns dict[job_id -> 'running'|'pending'] instead of set[job_id]. _RUNNING_STATUSES and _PENDING_STATUSES replace _ALIVE_STATUSES. _query_job_states() is the new primary method; _alive_job_ids() is a thin wrapper over it for ensure_workers(). Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/provisioner.py | 90 ++++++++++++------- tests/distributed/test_provisioner.py | 79 +++++++++------- 2 files changed, 101 insertions(+), 68 deletions(-) diff --git a/src/miss_alignment/distributed/provisioner.py b/src/miss_alignment/distributed/provisioner.py index a9ff4f8..f618a30 100644 --- a/src/miss_alignment/distributed/provisioner.py +++ b/src/miss_alignment/distributed/provisioner.py @@ -126,13 +126,20 @@ def shutdown(self) -> None: # (e.g. SLURM: squeue -u $USER -h -o "%i,%T"). # --------------------------------------------------------------------------- -# Statuses that mean the job is still occupying a slot (queued or running). -_ALIVE_STATUSES: dict[str, set[str]] = { - "slurm": {"PENDING", "PD", "RUNNING", "R", "COMPLETING", "CG", - "RESIZING", "SUSPENDED", "S"}, - "lsf": {"PEND", "RUN", "SSUSP", "USUSP", "PSUSP"}, - "pbs": {"Q", "R", "E", "H", "W", "T", "S"}, - "sge": {"qw", "r", "t", "Rr", "Rq", "hqw", "hRwq"}, +# 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"}, } @@ -142,15 +149,15 @@ def _parse_status_output( scheduler: str, custom_alive_statuses: list[str], custom_status_regex: str, -) -> set[str]: - """Parse status_list output and return the subset of our_job_ids still alive. +) -> dict[str, str]: + """Parse status_list output; return dict[job_id → 'running'|'pending']. Each line is expected to contain a job ID and a status token separated by a comma (matching the recommended status_list format "%i,%T" for SLURM). - Falls back to just checking if our job ID appears anywhere on a line for - schedulers whose status format varies. + Jobs absent from the output are considered terminal and not returned. + Jobs present without a recognisable status are treated as 'pending'. """ - alive: set[str] = set() + result: dict[str, str] = {} if scheduler == "auto": schedulers_to_try = ["slurm", "lsf", "pbs", "sge"] @@ -169,7 +176,6 @@ def _parse_status_output( job_id = job_id.strip() status_token = status_token.strip() else: - # No comma — treat the whole line as a job ID with unknown status. job_id = line status_token = "" @@ -177,26 +183,29 @@ def _parse_status_output( continue if scheduler == "custom": - if status_token in custom_alive_statuses: - alive.add(job_id) - elif not status_token: - alive.add(job_id) # presence without status → treat as alive + if status_token in custom_alive_statuses or not status_token: + result[job_id] = "pending" # custom mode doesn't distinguish continue if not status_token: - # Present in output but no status parsed → assume alive. - alive.add(job_id) + result[job_id] = "pending" continue + state = "pending" for sched in schedulers_to_try: - if status_token in _ALIVE_STATUSES[sched]: - alive.add(job_id) + if status_token in _RUNNING_STATUSES[sched]: + state = "running" + break + if status_token in _PENDING_STATUSES[sched]: + state = "pending" break else: - # Status token found but not in any alive set → job is terminal. - pass + # Not in any known set → treat as terminal, don't include. + continue + + result[job_id] = state - return alive + return result class ClusterProvisioner(WorkerProvisioner): @@ -224,14 +233,14 @@ def __init__(self, queue_dir: Path, config: ClusterConfig) -> None: self._scripts_dir = queue_dir / "cluster" self._scripts_dir.mkdir(parents=True, exist_ok=True) - def _alive_job_ids(self) -> set[str]: - """Query the scheduler and return the subset of our job IDs still alive. + def _query_job_states(self) -> dict[str, str]: + """Query the scheduler; return dict[job_id → 'running'|'pending']. - On scheduler error, returns the full submitted-ID set to avoid + On scheduler error, returns all submitted IDs as 'pending' to avoid resubmission storms during transient scheduler outages. """ if not self._job_ids: - return set() + return {} our_ids = set(self._job_ids) status_cmd = self._config.status_list.replace( "{{user}}", os.environ.get("USER", os.environ.get("USERNAME", "")) @@ -245,7 +254,7 @@ def _alive_job_ids(self) -> set[str]: text=True, timeout=30, ) - alive = _parse_status_output( + states = _parse_status_output( output=result.stdout, our_job_ids=our_ids, scheduler=self._config.scheduler, @@ -253,17 +262,30 @@ def _alive_job_ids(self) -> set[str]: custom_status_regex=self._config.custom_status_regex, ) # Prune terminated jobs from our tracking list. - self._job_ids = [jid for jid in self._job_ids if jid in alive] - return alive + self._job_ids = [jid for jid in self._job_ids if jid in states] + return states except Exception: - # Scheduler unavailable — assume all submitted jobs are still alive. - return our_ids + # Scheduler unavailable — assume all submitted jobs are pending. + return {jid: "pending" for jid in self._job_ids} + + def _alive_job_ids(self) -> set[str]: + return set(self._query_job_states().keys()) def live_worker_count(self) -> int: return len(self._alive_job_ids()) def worker_counts_by_type(self) -> dict[str, int]: - return {"cluster": self.live_worker_count()} + states = self._query_job_states() + running = sum(1 for s in states.values() if s == "running") + pending = sum(1 for s in states.values() if s == "pending") + counts = {} + if running: + counts["cluster-running"] = running + if pending: + counts["cluster-pending"] = pending + if not counts: + counts["cluster"] = 0 + return counts def _render_script(self, index: int) -> Path: template_text = self._config.script_path.read_text() diff --git a/tests/distributed/test_provisioner.py b/tests/distributed/test_provisioner.py index a43935c..ed736b3 100644 --- a/tests/distributed/test_provisioner.py +++ b/tests/distributed/test_provisioner.py @@ -29,68 +29,68 @@ def _cluster_cfg(tmp_path) -> ClusterConfig: # Status parser unit tests # --------------------------------------------------------------------------- -def test_parse_status_output_slurm_alive(): +def test_parse_status_output_slurm_running_and_pending(): output = "12345,PENDING\n12346,RUNNING\n12347,COMPLETING\n" our_ids = {"12345", "12346", "12347", "99999"} - alive = _parse_status_output(output, our_ids, "slurm", [], "") - assert alive == {"12345", "12346", "12347"} + states = _parse_status_output(output, our_ids, "slurm", [], "") + assert states == {"12345": "pending", "12346": "running", "12347": "running"} def test_parse_status_output_slurm_terminal(): output = "12345,COMPLETED\n12346,FAILED\n12347,CANCELLED\n" our_ids = {"12345", "12346", "12347"} - alive = _parse_status_output(output, our_ids, "slurm", [], "") - assert alive == set() + states = _parse_status_output(output, our_ids, "slurm", [], "") + assert states == {} -def test_parse_status_output_lsf_alive(): +def test_parse_status_output_lsf(): output = "12345,PEND\n12346,RUN\n" our_ids = {"12345", "12346"} - alive = _parse_status_output(output, our_ids, "lsf", [], "") - assert alive == {"12345", "12346"} + states = _parse_status_output(output, our_ids, "lsf", [], "") + assert states == {"12345": "pending", "12346": "running"} -def test_parse_status_output_pbs_alive(): +def test_parse_status_output_pbs(): output = "12345,Q\n12346,R\n12347,C\n" our_ids = {"12345", "12346", "12347"} - alive = _parse_status_output(output, our_ids, "pbs", [], "") - assert alive == {"12345", "12346"} + states = _parse_status_output(output, our_ids, "pbs", [], "") + assert states == {"12345": "pending", "12346": "running"} -def test_parse_status_output_sge_alive(): +def test_parse_status_output_sge(): output = "12345,qw\n12346,r\n" our_ids = {"12345", "12346"} - alive = _parse_status_output(output, our_ids, "sge", [], "") - assert alive == {"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"} - alive = _parse_status_output(output, our_ids, "auto", [], "") - assert alive == {"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"} - alive = _parse_status_output(output, our_ids, "slurm", [], "") - assert alive == set() + states = _parse_status_output(output, our_ids, "slurm", [], "") + assert states == {} -def test_parse_status_output_no_status_token_treated_as_alive(): - """A job present in output with no comma is assumed alive.""" +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"} - alive = _parse_status_output(output, our_ids, "slurm", [], "") - assert alive == {"12345"} + states = _parse_status_output(output, our_ids, "slurm", [], "") + assert states == {"12345": "pending"} def test_parse_status_output_custom(): output = "12345,INPROGRESS\n12346,DONE\n" our_ids = {"12345", "12346"} - alive = _parse_status_output(output, our_ids, "custom", ["INPROGRESS"], "") - assert alive == {"12345"} + states = _parse_status_output(output, our_ids, "custom", ["INPROGRESS"], "") + assert states == {"12345": "pending"} # --------------------------------------------------------------------------- @@ -166,7 +166,7 @@ def fake_run(cmd, **kwargs): with patch( "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run - ), patch.object(ClusterProvisioner, "_alive_job_ids", return_value=set()): + ), patch.object(ClusterProvisioner, "_query_job_states", return_value={}): p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) p.ensure_workers(n_workers=4) @@ -183,11 +183,10 @@ def fake_run(cmd, **kwargs): result.stdout = "Submitted batch job 12345\n" return result + alive = {"1": "running", "2": "running", "3": "pending", "4": "pending"} with patch( "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run - ), patch.object( - ClusterProvisioner, "_alive_job_ids", return_value={"1", "2", "3", "4"} - ): + ), patch.object(ClusterProvisioner, "_query_job_states", return_value=alive): p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) p.ensure_workers(n_workers=4) @@ -207,7 +206,9 @@ def fake_run(cmd, **kwargs): with patch( "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run ), patch.object( - ClusterProvisioner, "_alive_job_ids", return_value={"1", "2"} + ClusterProvisioner, + "_query_job_states", + return_value={"1": "running", "2": "pending"}, ): p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) p.ensure_workers(n_workers=4) @@ -215,6 +216,15 @@ def fake_run(cmd, **kwargs): assert len(submitted) == 2 +def test_cluster_provisioner_worker_counts_split_running_pending(tmp_path): + cfg = _cluster_cfg(tmp_path) + states = {"1": "running", "2": "running", "3": "pending", "4": "pending", "5": "pending"} + with patch.object(ClusterProvisioner, "_query_job_states", return_value=states): + p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) + counts = p.worker_counts_by_type() + assert counts == {"cluster-running": 2, "cluster-pending": 3} + + def test_cluster_provisioner_cancels_on_shutdown(tmp_path): cfg = _cluster_cfg(tmp_path) cancel_calls = [] @@ -229,7 +239,7 @@ def fake_run(cmd, **kwargs): with patch( "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run - ), patch.object(ClusterProvisioner, "_alive_job_ids", return_value=set()): + ), patch.object(ClusterProvisioner, "_query_job_states", return_value={}): p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) p.ensure_workers(n_workers=3) p.shutdown() @@ -239,10 +249,9 @@ def fake_run(cmd, **kwargs): def test_cluster_provisioner_prunes_terminated_jobs(tmp_path): - """_alive_job_ids prunes job IDs no longer in scheduler output.""" + """_query_job_states prunes job IDs no longer in scheduler output.""" cfg = _cluster_cfg(tmp_path) - # Simulate: 3 jobs submitted, only 2 alive in scheduler status_output = "11111,PENDING\n22222,RUNNING\n" def fake_run(cmd, **kwargs): @@ -256,9 +265,11 @@ def fake_run(cmd, **kwargs): ): p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) p._job_ids = ["11111", "22222", "33333"] - alive = p._alive_job_ids() + states = p._query_job_states() - assert alive == {"11111", "22222"} + assert set(states.keys()) == {"11111", "22222"} + assert states["11111"] == "pending" + assert states["22222"] == "running" assert p._job_ids == ["11111", "22222"] From 2aa1fd159d281aecc2627acdfe4c1208bf85a697 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Wed, 8 Jul 2026 15:43:11 -0700 Subject: [PATCH 31/33] fix: remove spurious cluster=0 from progress bar postfix MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two causes: 1. worker_counts_by_type() returned {'cluster': 0} as a fallback when no jobs were alive yet, which got merged into worker_counts and stuck there even after real keys appeared. Removed the fallback — empty dict is fine. 2. worker_counts was initialised by calling worker_counts_by_type() before any jobs were submitted (always empty). Changed to start as {}. 3. The scheduler tick used dict.update() which merges, preserving stale keys from previous ticks. Added worker_counts.clear() before the update. Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/manager.py | 3 ++- src/miss_alignment/distributed/provisioner.py | 2 -- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/src/miss_alignment/distributed/manager.py b/src/miss_alignment/distributed/manager.py index f2c25b3..497463d 100644 --- a/src/miss_alignment/distributed/manager.py +++ b/src/miss_alignment/distributed/manager.py @@ -121,6 +121,7 @@ def _scheduler_thread( # 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 @@ -220,7 +221,7 @@ def run_distributed( stop_event = threading.Event() scheduler_errors: list = [] - worker_counts: dict = provisioner.worker_counts_by_type() + worker_counts: dict = {} scheduler = threading.Thread( target=_scheduler_thread, args=( diff --git a/src/miss_alignment/distributed/provisioner.py b/src/miss_alignment/distributed/provisioner.py index f618a30..f70bdf3 100644 --- a/src/miss_alignment/distributed/provisioner.py +++ b/src/miss_alignment/distributed/provisioner.py @@ -283,8 +283,6 @@ def worker_counts_by_type(self) -> dict[str, int]: counts["cluster-running"] = running if pending: counts["cluster-pending"] = pending - if not counts: - counts["cluster"] = 0 return counts def _render_script(self, index: int) -> Path: From aed3c1c27ecc468d814459962660024b8e857ed1 Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Wed, 8 Jul 2026 18:45:12 -0700 Subject: [PATCH 32/33] fix: grace period for newly-submitted jobs missing from squeue output Jobs absent from squeue are kept for _JOB_GRACE_S (120s) before being pruned, rather than being dropped immediately or requiring an explicit terminal status. This correctly handles the window between job submission and scheduler registration (typically seconds, but up to minutes on busy clusters), which was causing runaway resubmission (>3000 jobs for a pool of 160). _job_submit_time replaces _job_ids to track each job's submission timestamp. ensure_workers uses len(_job_submit_time) as the authoritative count; _query_job_states only prunes jobs that are both absent from squeue AND have exceeded the grace period. worker_counts_by_type reports not-yet-visible jobs as cluster-pending so the progress bar count is accurate from the moment of submission. Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/provisioner.py | 99 +++++++----- tests/distributed/test_provisioner.py | 151 ++++++++++++------ 2 files changed, 156 insertions(+), 94 deletions(-) diff --git a/src/miss_alignment/distributed/provisioner.py b/src/miss_alignment/distributed/provisioner.py index f70bdf3..c3e1c19 100644 --- a/src/miss_alignment/distributed/provisioner.py +++ b/src/miss_alignment/distributed/provisioner.py @@ -143,6 +143,12 @@ def shutdown(self) -> None: } +# 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], @@ -152,10 +158,9 @@ def _parse_status_output( ) -> dict[str, str]: """Parse status_list output; return dict[job_id → 'running'|'pending']. - Each line is expected to contain a job ID and a status token separated by - a comma (matching the recommended status_list format "%i,%T" for SLURM). - Jobs absent from the output are considered terminal and not returned. - Jobs present without a recognisable status are treated as '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] = {} @@ -184,26 +189,22 @@ def _parse_status_output( if scheduler == "custom": if status_token in custom_alive_statuses or not status_token: - result[job_id] = "pending" # custom mode doesn't distinguish + result[job_id] = "pending" + # Otherwise not recognised as alive — leave absent from result. continue if not status_token: result[job_id] = "pending" continue - state = "pending" for sched in schedulers_to_try: if status_token in _RUNNING_STATUSES[sched]: - state = "running" + result[job_id] = "running" break if status_token in _PENDING_STATUSES[sched]: - state = "pending" + result[job_id] = "pending" break - else: - # Not in any known set → treat as terminal, don't include. - continue - - result[job_id] = state + # Jobs with unrecognised status are simply absent from the result. return result @@ -216,32 +217,39 @@ class ClusterProvisioner(WorkerProvisioner): 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" - LSF: bjobs -noheader -o "jobid stat" (with custom_status_regex) - PBS: qstat -u $USER (with custom parser) - SGE: qstat -u $USER (with custom parser) """ def __init__(self, queue_dir: Path, config: ClusterConfig) -> None: self._queue_dir = queue_dir self._config = config - self._job_ids: list[str] = [] + # 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']. - On scheduler error, returns all submitted IDs as 'pending' to avoid - resubmission storms during transient scheduler outages. + 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_ids: + if not self._job_submit_time: return {} - our_ids = set(self._job_ids) + our_ids = set(self._job_submit_time.keys()) status_cmd = self._config.status_list.replace( "{{user}}", os.environ.get("USER", os.environ.get("USERNAME", "")) ) @@ -261,23 +269,27 @@ def _query_job_states(self) -> dict[str, str]: custom_alive_statuses=self._config.custom_alive_statuses, custom_status_regex=self._config.custom_status_regex, ) - # Prune terminated jobs from our tracking list. - self._job_ids = [jid for jid in self._job_ids if jid in states] + 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: - # Scheduler unavailable — assume all submitted jobs are pending. - return {jid: "pending" for jid in self._job_ids} - - def _alive_job_ids(self) -> set[str]: - return set(self._query_job_states().keys()) + return {jid: "pending" for jid in self._job_submit_time} def live_worker_count(self) -> int: - return len(self._alive_job_ids()) + 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 @@ -285,6 +297,16 @@ def worker_counts_by_type(self) -> dict[str, int]: 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. @@ -324,26 +346,17 @@ def _submit_one(self) -> None: ) match = re.search(self._config.submit_job_id_regex, result.stdout) if match: - self._job_ids.append(match.group(1)) - - def ensure_workers(self, n_workers: int) -> None: - """Submit new jobs until alive job count reaches n_workers. - - Alive means queued or running according to the cluster scheduler. - """ - alive = len(self._alive_job_ids()) - deficit = n_workers - alive - for _ in range(deficit): - self._submit_one() + 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) - with ThreadPoolExecutor(max_workers=min(32, len(self._job_ids) or 1)) as pool: - list(as_completed([pool.submit(_cancel, jid) for jid in self._job_ids])) - self._job_ids.clear() + 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): diff --git a/tests/distributed/test_provisioner.py b/tests/distributed/test_provisioner.py index ed736b3..c0eddbc 100644 --- a/tests/distributed/test_provisioner.py +++ b/tests/distributed/test_provisioner.py @@ -1,4 +1,5 @@ import sys +import time from pathlib import Path from unittest.mock import MagicMock, patch @@ -36,7 +37,8 @@ def test_parse_status_output_slurm_running_and_pending(): assert states == {"12345": "pending", "12346": "running", "12347": "running"} -def test_parse_status_output_slurm_terminal(): +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", [], "") @@ -87,6 +89,7 @@ def test_parse_status_output_no_status_token_treated_as_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"], "") @@ -155,91 +158,162 @@ def test_local_provisioner_shutdown_terminates(tmp_path): # --------------------------------------------------------------------------- def test_cluster_provisioner_submits_n_workers_jobs(tmp_path): + """ensure_workers submits until _job_submit_time reaches target.""" cfg = _cluster_cfg(tmp_path) - submitted = [] + call_num = [0] def fake_run(cmd, **kwargs): - submitted.append(cmd) + call_num[0] += 1 result = MagicMock() - result.stdout = "Submitted batch job 12345\n" + 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 - ), patch.object(ClusterProvisioner, "_query_job_states", return_value={}): + ): p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) p.ensure_workers(n_workers=4) - assert len(submitted) == 4 + assert len(p._job_submit_time) == 4 -def test_cluster_provisioner_does_not_resubmit_alive_jobs(tmp_path): +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) - submitted = [] def fake_run(cmd, **kwargs): - submitted.append(cmd) result = MagicMock() - result.stdout = "Submitted batch job 12345\n" + result.stdout = "" # squeue returns nothing (jobs not yet registered) return result - alive = {"1": "running", "2": "running", "3": "pending", "4": "pending"} with patch( "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run - ), patch.object(ClusterProvisioner, "_query_job_states", return_value=alive): + ): 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(submitted) == 0 + 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_replenishes_preempted_jobs(tmp_path): +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): - submitted.append(cmd) result = MagicMock() - result.stdout = "Submitted batch job 99999\n" + 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 - ), patch.object( - ClusterProvisioner, - "_query_job_states", - return_value={"1": "running", "2": "pending"}, ): 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 + 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) - states = {"1": "running", "2": "running", "3": "pending", "4": "pending", "5": "pending"} - with patch.object(ClusterProvisioner, "_query_job_states", return_value=states): + + 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, "cluster-pending": 3} + + 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 = "Submitted batch job 99999\n" + 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 - ), patch.object(ClusterProvisioner, "_query_job_states", return_value={}): + ): p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) p.ensure_workers(n_workers=3) p.shutdown() @@ -248,31 +322,6 @@ def fake_run(cmd, **kwargs): assert all("scancel" in c for c in cancel_calls) -def test_cluster_provisioner_prunes_terminated_jobs(tmp_path): - """_query_job_states prunes job IDs no longer in scheduler output.""" - cfg = _cluster_cfg(tmp_path) - - status_output = "11111,PENDING\n22222,RUNNING\n" - - def fake_run(cmd, **kwargs): - result = MagicMock() - result.stdout = status_output - result.returncode = 0 - return result - - with patch( - "miss_alignment.distributed.provisioner.subprocess.run", side_effect=fake_run - ): - p = ClusterProvisioner(queue_dir=tmp_path, config=cfg) - p._job_ids = ["11111", "22222", "33333"] - states = p._query_job_states() - - assert set(states.keys()) == {"11111", "22222"} - assert states["11111"] == "pending" - assert states["22222"] == "running" - assert p._job_ids == ["11111", "22222"] - - # --------------------------------------------------------------------------- # CompositeProvisioner tests # --------------------------------------------------------------------------- From 824c0fbc555680a88d737211a312561e6031113a Mon Sep 17 00:00:00 2001 From: Dimitry Tegunov Date: Wed, 8 Jul 2026 18:47:42 -0700 Subject: [PATCH 33/33] fix: lock LocalProvisioner.ensure_workers to prevent concurrent double-spawn The scheduler thread and the startup ensure_workers() call in run_distributed could execute concurrently. If both checked _procs[device] before either had stored the new Popen, both would see None and each spawn a separate process for the same GPU, explaining the 2 processes per GPU observed in running/. Added threading.Lock to LocalProvisioner; ensure_workers and live_worker_count both acquire it. Co-Authored-By: Claude Fable 5 --- src/miss_alignment/distributed/provisioner.py | 44 ++++++++++--------- 1 file changed, 24 insertions(+), 20 deletions(-) diff --git a/src/miss_alignment/distributed/provisioner.py b/src/miss_alignment/distributed/provisioner.py index c3e1c19..93ae8fa 100644 --- a/src/miss_alignment/distributed/provisioner.py +++ b/src/miss_alignment/distributed/provisioner.py @@ -6,6 +6,7 @@ import re import subprocess import sys +import threading import time from abc import ABC, abstractmethod from concurrent.futures import ThreadPoolExecutor, as_completed @@ -79,33 +80,36 @@ 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: - return sum(1 for p in self._procs.values() if p.poll() is None) + 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: - 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 + 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():