tka/sras-viewer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Dedup workers: one _PooledWorker base for the three pool fan-outs

Browse Source 
DcPrecomputeWorker, Ch4MaskWorker and CrossCorrelateWorker shared the
same submit/as_completed/emit/shutdown skeleton; they are now
constructors plus _plan/_items/_one/_emit hooks. All three inherit the
cancellation-aware shutdown (wait=not stopped, cancel_futures) that
only DcPrecomputeWorker had before, and the per-run budget attributes
are initialised in __init__ instead of appearing mid-run.

Also removes the viewer's write-only _dc_precompute_worker plumbing
(job tracking already owns worker lifetime via _jobs).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales
2026-08-06 10:32:06 -05:00
parent 9a9a2557d6
commit e80717d5c5
2 changed files with 67 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sras_workers.py
+66 -70
View File
@@ -54,6 +54,34 @@ class CancellableWorker(QObject):
return self._stop
class _PooledWorker(CancellableWorker):
"""Fans a per-item computation across a thread pool, emitting each result
from this worker's own thread as it lands (never from a pool thread).
Subclasses provide _plan() -> n_workers (stashing whatever per-run
context they need), _items(), _one(item) -> result, and _emit(result).
On stop(): queued items are dropped, in-flight ones are not waited for —
that is what keeps closing the window responsive on a large scan.
"""
finished = pyqtSignal()
error = pyqtSignal(str)
def run(self):
try:
pool = ThreadPoolExecutor(max_workers=max(1, self._plan()))
try:
futures = [pool.submit(self._one, it) for it in self._items()]
for fut in as_completed(futures):
if self._stop:
break
self._emit(fut.result())
finally:
pool.shutdown(wait=not self._stop, cancel_futures=True)
self.finished.emit()
except Exception as exc:
self.error.emit(str(exc))
class LoadWorker(QObject):
finished = pyqtSignal(object) # SrasFile | None
error = pyqtSignal(str)
@@ -117,29 +145,29 @@ class ComputeWorker(CancellableWorker):
self.error.emit(str(exc))
class DcPrecomputeWorker(CancellableWorker):
class DcPrecomputeWorker(_PooledWorker):
"""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
self._angle_budget = 0
def _one_angle(self, a: int) -> tuple[int, np.ndarray, np.ndarray]:
def _plan(self) -> int:
n_workers, self._angle_budget = compute.plan_angle_level(self._sras)
return n_workers
def _items(self):
return range(self._sras.n_angles)
def _one(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).
@@ -149,26 +177,8 @@ class DcPrecomputeWorker(CancellableWorker):
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))
def _emit(self, result):
self.angle_done.emit(*result)
class BatchCacheWorker(QObject):
@@ -306,7 +316,7 @@ class AngleAlignmentWorker(QObject):
self.finished.emit(None, str(exc))
class Ch4MaskWorker(QObject):
class Ch4MaskWorker(_PooledWorker):
"""Fetches each requested angle's CH4 (Bias B) DC image in mV, for
ManualAlignmentDialog's initial threshold-mask overlay.
@@ -320,35 +330,29 @@ class Ch4MaskWorker(QObject):
ManualAlignmentDialog._start_mask_prep).
"""
angle_done = pyqtSignal(int, np.ndarray) # angle_idx, dc4_mv
finished = pyqtSignal()
error = pyqtSignal(str)
def __init__(self, sras: SrasFile, angle_indices: list[int]):
super().__init__()
self._sras = sras
self._angles = angle_indices
self._budget = 0
def run(self):
try:
n_workers, budget = compute.plan_angle_level(self._sras)
pool = ThreadPoolExecutor(max_workers=n_workers)
try:
futures = {
pool.submit(dc_image_mv, self._sras, a, CH4_IDX,
max_workers=1, budget=budget): a
for a in self._angles
}
for fut in as_completed(futures):
a = futures[fut]
self.angle_done.emit(a, fut.result())
finally:
pool.shutdown(wait=True)
self.finished.emit()
except Exception as exc:
self.error.emit(str(exc))
def _plan(self) -> int:
n_workers, self._budget = compute.plan_angle_level(self._sras)
return n_workers
def _items(self):
return self._angles
def _one(self, a: int) -> tuple[int, np.ndarray]:
return a, dc_image_mv(self._sras, a, CH4_IDX,
max_workers=1, budget=self._budget)
def _emit(self, result):
self.angle_done.emit(*result)
class CrossCorrelateWorker(QObject):
class CrossCorrelateWorker(_PooledWorker):
"""Rigid registration (rotation + translation, never scale) of each of
*angle_indices* against *ref_angle_idx*, for ManualAlignmentDialog's Auto
Cross-Correlate button.
@@ -363,8 +367,6 @@ class CrossCorrelateWorker(QObject):
"""
# angle_idx, rotation_deg, shift_x_mm, shift_y_mm, score, source
angle_done = pyqtSignal(int, float, float, float, float, str)
finished = pyqtSignal()
error = pyqtSignal(str)
def __init__(self, sras: SrasFile, ref_angle_idx: int, angle_indices: list[int],
dc4_mv: dict[int, np.ndarray], *,
@@ -379,25 +381,19 @@ class CrossCorrelateWorker(QObject):
self._threshold = dc_threshold_mv
self._search_deg = search_deg
def _plan(self) -> int:
return compute._registration_workers(self._sras, compute._DEFAULT_FINE_DIM)
def _items(self):
return self._angles
def _one(self, a: int) -> tuple[int, compute.RigidFit]:
return a, compute.register_angle_to_reference(
self._sras, a, self._ref, self._dc4_mv,
dc_threshold_mv=self._threshold, sources=self._sources,
search_deg=self._search_deg)
def run(self):
try:
n_workers = compute._registration_workers(
self._sras, compute._DEFAULT_FINE_DIM)
pool = ThreadPoolExecutor(max_workers=max(1, n_workers))
try:
futures = [pool.submit(self._one, a) for a in self._angles]
for fut in as_completed(futures):
a, fit = fut.result()
self.angle_done.emit(a, fit.rotation_deg, fit.shift_mm[0],
fit.shift_mm[1], fit.score, fit.source)
finally:
pool.shutdown(wait=True)
self.finished.emit()
except Exception as exc:
self.error.emit(str(exc))
def _emit(self, result):
a, fit = result
self.angle_done.emit(a, fit.rotation_deg, fit.shift_mm[0],
fit.shift_mm[1], fit.score, fit.source)