Thomas Ales
9a7cc396e7
Three defects found by testing the parallel paths end to end.
Nested parallelism was unbounded. DcPrecomputeWorker and the alignment
driver both parallelise over angles, and each angle's compute_dc_image
then parallelised over rows again: 14 x 14 threads, and — worse — every
angle sized its chunk against the *whole* memory budget, so worst-case
live buffers were ~13.9 GB against a 1 GB budget. Capping the inner
worker count alone does not fix it; the chunk is still sized against the
full budget. plan_angle_level now hands each angle both max_workers=1
and its share of the budget, which bounds the real file at 1002 MB.
Chunk planning could not produce concurrency at all. _plan_chunks sized
chunk_rows first, and _chunk_rows_for spends whatever budget it is given
— so a single chunk consumed the lot and budget//chunk_bytes came back
as one worker. It now picks the worker count first and sizes the chunk to
it, giving 7-16 workers on the real file instead of 1.
Cancellation was too coarse to shut down. stop() was only checked between
angles, and one angle of the real scan is ~40 s, so closing the window
sat through it and returned with threads still running. should_stop is
now polled per row chunk and closeEvent signals every worker before
waiting: close went from 4.0 s with two live threads to 1.04 s with both
finished. A cancelled ComputeWorker emits None so a half-filled image is
never cached.
Also: BatchCacheWorker no longer reports every file as failed when the
process pool itself dies (unguarded __main__, frozen build, sandbox) —
it retries those files in-process. Batches below 512 MB skip the pool
entirely, since interpreter startup would otherwise make small jobs
slower. cache_file rejects an unknown mode instead of silently treating
it as "fft".
Measured on the real 496 GB scan, angle 3, 32 rows:
DC 2.94s -> 0.75s RF 6.10s -> 1.99s peak RSS 1.40 -> 1.93 GB
Batch convert, 24 files / 1 GB: 4.03s -> 2.36s.
Memory budget defaults to 1024 MB (SRAS_MEM_BUDGET_MB), the measured knee.
Testing: tools/test_refactor.py (70 checks: cache round-trip and block
carry-forward, partial-cache handling, parallel-vs-serial identity,
no-mask path, ROI mask vs full-grid, v2/v3/v4 parsing, sras_average
round-trip incl. remainder handling) and tools/test_gui.py (46 checks
driving the real widgets and workers headlessly). check_equivalence.py
still reports all 167 outputs identical to ed0eba4.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
307 lines
12 KiB
Python
307 lines
12 KiB
Python
#!/usr/bin/env python3
|
||
"""Background workers for the SRAS viewer.
|
||
|
||
Every worker is a plain QObject moved onto its own QThread by
|
||
SrasViewerWindow._run_worker, exposing signals only. Workers must never touch
|
||
GUI-thread-owned state (the display caches in particular) — they take
|
||
everything they need through their constructor and hand results back by signal.
|
||
"""
|
||
|
||
import os
|
||
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
|
||
from concurrent.futures.process import BrokenProcessPool
|
||
|
||
import numpy as np
|
||
from PyQt6.QtCore import QObject, pyqtSignal
|
||
|
||
import sras_compute as compute
|
||
from sras_compute import (
|
||
cache_file, compute_angle_alignment, compute_rf_image, dc_image_mv,
|
||
)
|
||
from sras_format import CH3_IDX, CH4_IDX, SrasFile
|
||
|
||
# Concurrency caps. Batch conversion runs one process per file, and each of
|
||
# those processes threads internally, so the two must be divided rather than
|
||
# both set to the core count. Files also commonly sit on one external drive,
|
||
# where a dozen concurrent readers is slower than a few — hence the low
|
||
# default, overridable from the environment.
|
||
_BATCH_MAX_PROCS = int(os.environ.get("SRAS_BATCH_PROCS", 0)) or min(
|
||
4, os.cpu_count() or 2)
|
||
|
||
# Spawning a pool costs roughly a second of interpreter startup (each child
|
||
# re-imports the entry module). That is noise against a multi-GB scan but
|
||
# dominates a batch of small files, where it would make the job *slower* —
|
||
# so below this total size the batch just runs in the worker thread.
|
||
_BATCH_POOL_MIN_BYTES = int(os.environ.get("SRAS_BATCH_POOL_MIN_MB", 512)) * 1024 * 1024
|
||
|
||
|
||
class CancellableWorker(QObject):
|
||
"""A worker whose compute polls stop() between row chunks.
|
||
|
||
Without this a shutdown has to wait out whatever is in flight, and on a
|
||
large scan a single angle is ~40 s — far too long to block closing the
|
||
window. Chunk-level polling bounds the wait to one chunk instead.
|
||
"""
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self._stop = False
|
||
|
||
def stop(self):
|
||
self._stop = True
|
||
|
||
def _stopped(self) -> bool:
|
||
return self._stop
|
||
|
||
|
||
class LoadWorker(QObject):
|
||
finished = pyqtSignal(object) # SrasFile | None
|
||
error = pyqtSignal(str)
|
||
|
||
def __init__(self, path: str):
|
||
super().__init__()
|
||
self._path = path
|
||
|
||
def run(self):
|
||
try:
|
||
self.finished.emit(SrasFile(self._path))
|
||
except Exception as exc:
|
||
self.error.emit(str(exc))
|
||
self.finished.emit(None)
|
||
|
||
|
||
class ComputeWorker(CancellableWorker):
|
||
"""Computes one displayable image for (angle, channel).
|
||
|
||
For CH1/Velocity (FFT-derived) channels, the FFT is only run for pixels
|
||
whose DC4 (Bias B) mean is at or above dc_threshold_mv — masked pixels are
|
||
left at 0 MHz without ever being FFT'd, since that's the expensive part of
|
||
a scan. If the DC4 image for this angle is already known, pass it in as
|
||
*dc4_mv* to skip re-reading the CH4 channel from disk entirely.
|
||
|
||
Emits a plain ``np.ndarray`` already in display units.
|
||
"""
|
||
finished = pyqtSignal(object)
|
||
error = pyqtSignal(str)
|
||
|
||
def __init__(self, sras: SrasFile, angle_idx: int, ch_idx: int,
|
||
apply_bg_sub: bool = True, n_fft: int | None = None,
|
||
dc_threshold_mv: float = 0.0,
|
||
dc4_mv: np.ndarray | None = None,
|
||
is_fft_mode: bool = False):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._angle = angle_idx
|
||
self._ch = ch_idx
|
||
self._apply_bg_sub = apply_bg_sub
|
||
self._n_fft = n_fft
|
||
self._dc_threshold = dc_threshold_mv
|
||
self._dc4_mv = dc4_mv
|
||
self._is_fft_mode = is_fft_mode
|
||
|
||
def run(self):
|
||
try:
|
||
if self._is_fft_mode:
|
||
img = compute_rf_image(
|
||
self._sras, self._angle, dc_threshold_mv=self._dc_threshold,
|
||
apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft,
|
||
dc4_mv=self._dc4_mv, should_stop=self._stopped)
|
||
else:
|
||
img = dc_image_mv(self._sras, self._angle, self._ch,
|
||
should_stop=self._stopped)
|
||
# On cancellation the image is only partly filled, so hand back
|
||
# None rather than something that would be cached as real. The
|
||
# signal still fires either way — it is what quits the thread.
|
||
self.finished.emit(None if self._stop else img)
|
||
except Exception as exc:
|
||
self.error.emit(str(exc))
|
||
|
||
|
||
class DcPrecomputeWorker(CancellableWorker):
|
||
"""Computes CH3/CH4 DC images for every angle in the background.
|
||
|
||
DC images are cheap (a per-waveform mean, no FFT) compared to the
|
||
CH1/Velocity FFT, so precomputing them for the whole file right after load
|
||
makes switching angles instant while on a DC channel, and also means the
|
||
FFT masking step (which needs a DC4 image) rarely has to wait on anything.
|
||
|
||
Angles are computed on a thread pool — the work is a pure mean over the
|
||
waveform block, so it is I/O- and bandwidth-bound and embarrassingly
|
||
parallel. Results are emitted one at a time as they land (out of angle
|
||
order), and always from this worker's own thread: nothing emits a Qt
|
||
signal from a pool thread.
|
||
"""
|
||
angle_done = pyqtSignal(int, np.ndarray, np.ndarray) # angle_idx, dc3_mv, dc4_mv
|
||
finished = pyqtSignal()
|
||
error = pyqtSignal(str)
|
||
|
||
def __init__(self, sras: SrasFile):
|
||
super().__init__()
|
||
self._sras = sras
|
||
|
||
def _one_angle(self, a: int) -> tuple[int, np.ndarray, np.ndarray]:
|
||
# max_workers=1 *and* a budget share: this call is one of several
|
||
# concurrent angles, and both the thread count and the buffer size
|
||
# have to be divided (see compute.plan_angle_level).
|
||
kw = dict(max_workers=1, budget=self._angle_budget,
|
||
should_stop=self._stopped)
|
||
return (a,
|
||
dc_image_mv(self._sras, a, CH3_IDX, **kw),
|
||
dc_image_mv(self._sras, a, CH4_IDX, **kw))
|
||
|
||
def run(self):
|
||
try:
|
||
n = self._sras.n_angles
|
||
n_workers, self._angle_budget = compute.plan_angle_level(self._sras)
|
||
pool = ThreadPoolExecutor(max_workers=n_workers)
|
||
try:
|
||
futures = {pool.submit(self._one_angle, a): a for a in range(n)}
|
||
for fut in as_completed(futures):
|
||
if self._stop:
|
||
break
|
||
a, dc3, dc4 = fut.result()
|
||
self.angle_done.emit(a, dc3, dc4)
|
||
finally:
|
||
# cancel_futures drops the queued angles; should_stop lets the
|
||
# in-flight ones bail within a chunk. Not waiting here is what
|
||
# keeps closing the window responsive on a large scan.
|
||
pool.shutdown(wait=not self._stop, cancel_futures=True)
|
||
self.finished.emit()
|
||
except Exception as exc:
|
||
self.error.emit(str(exc))
|
||
|
||
|
||
class BatchCacheWorker(QObject):
|
||
"""Batch-computes and stores DC or FFT images into each of *paths*'s v7
|
||
CACH tail, in place — converting v6 sources to v7 on first use, or updating
|
||
an existing v7 file's cache blocks without disturbing whatever the other
|
||
block already holds.
|
||
|
||
*mode* is ``"dc"`` (CH3/CH4 mean images) or ``"fft"`` (CH1 peak-frequency
|
||
images, unmasked — masking is applied at display time, same as v5's PREC
|
||
convention).
|
||
|
||
Files are processed one per subprocess: they are fully independent, each
|
||
opens its own memmap and writes only its own bytes, and only path strings
|
||
and scalars cross the process boundary. Emits ``progress(int)`` (0–100 by
|
||
files completed), ``file_done(str, str)`` (path, error message or "") so
|
||
one file's failure doesn't abort the batch, and ``finished()``.
|
||
"""
|
||
progress = pyqtSignal(int)
|
||
file_done = pyqtSignal(str, str)
|
||
finished = pyqtSignal()
|
||
|
||
def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool):
|
||
super().__init__()
|
||
self._paths = paths
|
||
self._mode = mode
|
||
self._apply_bg_sub = apply_bg_sub
|
||
|
||
def _report(self, path: str, err: str, done: int, total: int):
|
||
self.file_done.emit(path, err)
|
||
self.progress.emit(int(done / max(1, total) * 100))
|
||
|
||
def _run_pooled(self, paths: list[str], n_procs: int) -> list[str]:
|
||
"""Process the batch across *n_procs* subprocesses. Returns the paths
|
||
that never got a real answer because the pool itself died, so the
|
||
caller can retry them in-process.
|
||
|
||
Under spawn each child re-imports the entry module, so the batch must
|
||
survive that going wrong (an unguarded __main__, a frozen build, a
|
||
sandbox that forbids subprocesses) rather than reporting every file as
|
||
failed — hence the retry list instead of a per-file error.
|
||
"""
|
||
# Each child threads internally; divide the machine rather than
|
||
# letting every process claim every core.
|
||
per_proc_workers = max(1, (os.cpu_count() or 4) // n_procs)
|
||
unresolved: list[str] = []
|
||
done = 0
|
||
|
||
with ProcessPoolExecutor(max_workers=n_procs) as executor:
|
||
futures = {
|
||
executor.submit(cache_file, p, self._mode, self._apply_bg_sub,
|
||
compute.get_fft_backend(), per_proc_workers): p
|
||
for p in paths
|
||
}
|
||
for fut in as_completed(futures):
|
||
path = futures[fut]
|
||
try:
|
||
err = fut.result()
|
||
except BrokenProcessPool:
|
||
unresolved.append(path)
|
||
continue
|
||
except Exception as exc:
|
||
err = str(exc)
|
||
done += 1
|
||
self._report(path, err, done, len(paths))
|
||
|
||
return unresolved
|
||
|
||
def _run_inline(self, paths: list[str], done: int, total: int):
|
||
"""Fallback / single-file path: compute in this thread. Still uses the
|
||
full core count internally, since nothing else is competing."""
|
||
for path in paths:
|
||
try:
|
||
err = cache_file(path, self._mode, self._apply_bg_sub,
|
||
compute.get_fft_backend(), compute._MAX_WORKERS)
|
||
except Exception as exc:
|
||
err = str(exc)
|
||
done += 1
|
||
self._report(path, err, done, total)
|
||
|
||
def _worth_pooling(self, paths: list[str]) -> bool:
|
||
if len(paths) < 2:
|
||
return False
|
||
total = 0
|
||
for p in paths:
|
||
try:
|
||
total += os.path.getsize(p)
|
||
except OSError:
|
||
pass # unreadable files are reported by cache_file
|
||
return total >= _BATCH_POOL_MIN_BYTES
|
||
|
||
def run(self):
|
||
paths = self._paths
|
||
n_procs = max(1, min(_BATCH_MAX_PROCS, len(paths)))
|
||
|
||
if not self._worth_pooling(paths):
|
||
self._run_inline(paths, 0, len(paths))
|
||
self.finished.emit()
|
||
return
|
||
|
||
try:
|
||
unresolved = self._run_pooled(paths, n_procs)
|
||
except Exception:
|
||
# The pool could not be created or collapsed wholesale.
|
||
unresolved = list(paths)
|
||
|
||
if unresolved:
|
||
self._run_inline(unresolved, len(paths) - len(unresolved), len(paths))
|
||
|
||
self.finished.emit()
|
||
|
||
|
||
class AngleAlignmentWorker(QObject):
|
||
"""Computes rotation+translation alignment for every angle in *sras*,
|
||
referenced to *ref_angle_idx*, from each angle's binarized CH4 mask.
|
||
Rotation is analytic (from sras.angles_deg); only translation is found by
|
||
phase correlation.
|
||
"""
|
||
progress = pyqtSignal(int) # 0–100
|
||
finished = pyqtSignal(object, str) # AlignmentResult|None, error ("" = success)
|
||
|
||
def __init__(self, sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._ref = ref_angle_idx
|
||
self._threshold = dc_threshold_mv
|
||
|
||
def run(self):
|
||
try:
|
||
result = compute_angle_alignment(
|
||
self._sras, self._ref, self._threshold,
|
||
progress_cb=self.progress.emit)
|
||
self.finished.emit(result, "")
|
||
except Exception as exc:
|
||
self.finished.emit(None, str(exc))
|