#!/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))