Worker pool with disk-based communication for distributed inference

[dtegunov]

What this does

The alignment phase — where the trained model scores reconstruction quality across all tilt series — previously ran sequentially on local GPUs. With 300+ tilt series this took hours per macro-iteration, making it the dominant bottleneck in a full training run. This PR adds a worker pool with disk-based communication that distributes inference across a SLURM cluster, bringing that phase from hours to minutes.

How it works

A filesystem queue under <training_dir>/tasks/ coordinates work with no network dependency — atomic file rename is the claim mutex (the same protocol used by WarpTools). The head node writes one task JSON per tilt series into tasks/pending/, submits N cluster jobs, and blocks until all tasks reach tasks/done/ or tasks/failed/. Workers claim tasks by rename, process them, and loop until the queue drains.

The same mechanism covers the stack preparation (--prepare-stacks) and cross-correlation pre-alignment (--preprocess) phases.

Usage

Set two environment variables pointing at a cluster config and submission script template, then add --n-cluster-workers N to your train or infer command:

export MISS_CLUSTER_CONFIG=/path/to/cluster_config.json
export MISS_CLUSTER_SCRIPT=/path/to/worker.sh
miss-alignment train --config-file config.yaml --n-cluster-workers 100

SLURM config and a ready-to-adapt submission script are provided in cluster_example/. Without --n-cluster-workers, behaviour is identical to before (local multi-GPU pool).

New CLI command

miss-alignment worker --queue-dir <path> --device <int> [--worker-id <str>]

Workers can also be launched manually against a running queue, useful for topping up a depleted pool mid-run.

Key design details

• Model checkpoint reuse: workers load the checkpoint once and reuse it across all series they process in a macro-iteration — for 300 series split across 10 workers that is 30 loads instead of 300
• Heartbeat in background thread: the worker heartbeat runs in a daemon thread independent of task execution, so long-running series (>2 min) do not appear stale to the manager's sweep
• Stall recovery: the manager's scheduler thread detects stalled workers (missing heartbeat >120s) and re-pends their tasks; a stalled worker's running/ directory is cleaned up so it can recreate it if it comes back
• Hard fail on any series failure: if any series lands in failed/, the manager raises after all remaining tasks complete — training stops rather than producing silently incomplete alignment
• Exit reason files: each worker writes tasks/logs/<worker_id>.exit on shutdown, recording why it stopped (queue empty / manager stale / unhandled exception with traceback) and how many series it processed
• SLURM logs in tasks/logs/: slurm-<jobid>.out/err land alongside the exit files, not in the working directory

Bugs fixed during testing

• Scheduler thread crashing silently (unhandled exception now surfaces to main thread immediately)
• TOCTOU race reading heartbeat mtime (heartbeat thread could delete the tick file between max() and the second .stat() call)
• Double worker submission at startup (scheduler thread and explicit startup call both fired ensure_workers on the same empty job list)
• Stall sweep re-pending tasks already in done/ (visible as ls tasks/done | wc -l exceeding total series count)

🤖 Generated with Claude Code

[McHaillet]

The tests are failing. It seems to have something to do with torch-projectors, might be a combination of the new version with a pytorch mismatch in the CI.

[github-actions]

Coverage report

Click to see where and how coverage changed

File	Statements	Missing	Coverage	Coverage (new stmts)	Lines missing
src/miss_alignment
__init__.py
__main__.py					1-3
_cli.py
infer.py					120
prepare_stacks.py					89-92, 138-167
preprocessing.py
train.py					418-424
src/miss_alignment/alignment
parallel.py					28-30
tilt_series.py					174-175, 187-191
src/miss_alignment/data
training_datamodule.py					215-217
src/miss_alignment/distributed
__init__.py
config.py					55
manager.py					44, 48-58, 60-79, 81-85, 112-114, 117-126, 129-131
provisioner.py					26-41, 43-50, 52, 73, 86-87, 90, 121-122, 181, 245, 284-285, 288, 330-331, 377, 380-384
queue.py					148-150, 196, 201, 204, 207
worker.py					41-49, 64-65, 76, 92, 139
Project Total

This report was generated by python-coverage-comment-action

[dtegunov]

The tests are failing. It seems to have something to do with torch-projectors, might be a combination of the new version with a pytorch mismatch in the CI.

Fixed.

[McHaillet]

Thanks for the CI fix!

I had some time to read thought it now. I like the functionality, but have a concern about the added complexity this adds to miss-alignment while cluster job spawning might be better handled via manager programs. Would it not be easier to run the inference script on a single GPU with a single tilt-series, and let manager programs handle launching jobs over GPU's ? Such as Warp or the Relay GUI, or RELION/Doppio if they want to integrate it. That would keep the code here more minimal. Since that is not supported yet, it might be nice for now, and I also like that it is backwards compatible. What are your thoughts on this?

Also question about the cluster workers in the training script, does it submit to the gpu's on the training node as well?

[dtegunov]

I agree that a single-item-per-invocation model would help replicate the behavior in inference when paired with a scheduler that's aware of individual tilt series. However, such schedulers are hard to come by (Relay doesn't support this; not sure about Doppio), and it wouldn't solve the training problem because it alternates between training and inference. The latter is certainly not impossible, but it would require a bespoke communication mechanism between MA and the scheduler that I doubt anyone would want to implement (I briefly considered using Relay's native worker pools, but it was too much of a mess). And then you'd still want to support scheduler-free inference, where MA has to handle the parallelization internally. I think using the same parallelization mechanism everywhere is just easier.

[McHaillet]

I get what you mean.

It does seem like the original GPU's for training are left idle, might be a bit wasteful. Although it would be difficult to submit to those of course

[McHaillet]

FYI: I merged my changes to main with this PR -- it removed pre-tilt estimation with tiltxcorr which was not working very reliably. I found WarpTools --auto_zero does a very decent job and finding the pretilt.

[dtegunov]

It does seem like the original GPU's for training are left idle, might be a bit wasteful. Although it would be difficult to submit to those of course

That's actually easy since we already have separate logic for spawning local and remote workers – they just need to be combined. Good point! Implemented.