Fix crash on angle switch with large v6 scans: bound chunk memory, harden thread teardown

Root cause: compute_dc_image/compute_rf_image chunked processing by a fixed
32-row count, sized for old small-format scans. A real v6 file with
spf=2500 and up to ~7500 frames/angle needed ~6-7 GB for a single chunk's
core buffers, on a 17 GB machine — pushing past available memory on every
angle and aborting the process (not a clean MemoryError) when switching to
a new angle piled more allocations on top. Chunk size is now computed from
actual scan dimensions to hit a fixed ~128 MB budget instead of a fixed row
count, cutting peak footprint by roughly 16-18x (verified against the real
532 GB / 17-angle dataset: ~400 MB peak, no crash).

Also hardens QThread lifecycle handling, found while chasing this crash:
- _on_compute_thread_finished/_on_load_thread_finished/
  _on_preprocess_thread_finished now call wait() before dropping the last
  reference to a finished QThread, avoiding "QThread: Destroyed while
  thread is still running" aborts if the OS thread hasn't fully joined
  when finished() fires.
- _start_compute/_load_file/_on_preprocess now claim their QThread
  immediately after the busy-guard check, before any call that can pump
  the Qt event loop (e.g. first QProgressDialog.show()), closing a
  reentrancy window where a second call could start a thread that the
  first call's own assignment would then clobber mid-run.
- faulthandler is enabled at startup so any future native crash prints a
  real stack trace instead of a bare "Aborted".

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales
2026-07-30 20:04:39 -05:00
parent 4a4e588097
commit 9f92d1826e
+83 -22
View File
@@ -25,10 +25,13 @@ those arrays simply repeat the same value n_angles times.
import re import re
import sys import sys
import struct import struct
import faulthandler
import numpy as np import numpy as np
from pathlib import Path from pathlib import Path
import os import os
faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc.
from PyQt6.QtWidgets import ( from PyQt6.QtWidgets import (
QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
QGroupBox, QLabel, QPushButton, QComboBox, QSpinBox, QDoubleSpinBox, QGroupBox, QLabel, QPushButton, QComboBox, QSpinBox, QDoubleSpinBox,
@@ -576,21 +579,38 @@ class SrasFile:
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
_CHUNK_ROWS = 32 # rows processed per FFT/DC chunk; tune for RAM vs. throughput # Rows are batched so the float32 working buffer for one channel's chunk
# (chunk_rows × n_frames × spf × 4 bytes) stays under this budget. A fixed
# row count (the original design) works fine for small legacy scans but is
# catastrophic for a v6 scan with a large per-angle frame/sample count —
# e.g. a 7500-frame × 2500-sample angle needs ~2.4 GB for a single 32-row
# chunk, times several such buffers alive at once for the FFT step, which
# can exceed physical RAM entirely on its own. Sizing the chunk to the
# actual dimensions keeps peak RAM bounded regardless of scan size.
_CHUNK_BYTES_BUDGET = 128 * 1024 * 1024 # ~128 MB per channel-buffer chunk
_CHUNK_ROWS_MAX = 32 # cap for small scans (old behavior)
def _chunk_rows_for(n_frames: int, samples_per_frame: int) -> int:
bytes_per_row = max(1, n_frames * samples_per_frame * 4) # float32
rows = _CHUNK_BYTES_BUDGET // bytes_per_row
return int(max(1, min(_CHUNK_ROWS_MAX, rows)))
def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray: def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray:
"""Mean of each waveform → (n_rows, n_frames) float32. """Mean of each waveform → (n_rows, n_frames) float32.
Processes in row chunks so the float32 working buffer is bounded to Processes in row chunks sized to a fixed memory budget (see
``_CHUNK_ROWS × n_frames × spf × 4`` bytes regardless of scan size. ``_chunk_rows_for``) so the float32 working buffer stays bounded
regardless of scan size.
""" """
n_rows = int(sras.n_rows[angle_idx]) n_rows = int(sras.n_rows[angle_idx])
n_frames = int(sras.n_frames[angle_idx]) n_frames = int(sras.n_frames[angle_idx])
data = sras.data[angle_idx] data = sras.data[angle_idx]
chunk_rows = _chunk_rows_for(n_frames, sras.samples_per_frame)
img = np.empty((n_rows, n_frames), dtype=np.float32) img = np.empty((n_rows, n_frames), dtype=np.float32)
for r0 in range(0, n_rows, _CHUNK_ROWS): for r0 in range(0, n_rows, chunk_rows):
r1 = min(r0 + _CHUNK_ROWS, n_rows) r1 = min(r0 + chunk_rows, n_rows)
img[r0:r1] = ( img[r0:r1] = (
data[r0:r1, ch_idx, :, :] data[r0:r1, ch_idx, :, :]
.astype(np.float32) .astype(np.float32)
@@ -611,9 +631,9 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
and zero-padding is not active, and the bg-sub flag matches, the stored and zero-padding is not active, and the bg-sub flag matches, the stored
images are used directly — no FFT is run. images are used directly — no FFT is run.
Otherwise, data is processed in row chunks of ``_CHUNK_ROWS`` to bound Otherwise, data is processed in row chunks sized to a fixed memory
peak RAM to roughly ``_CHUNK_ROWS × n_frames × max(spf, n_fft) × 24`` budget (see ``_chunk_rows_for``) to bound peak RAM regardless of scan
bytes regardless of scan size. size.
""" """
n_rows = int(sras.n_rows[angle_idx]) n_rows = int(sras.n_rows[angle_idx])
n_frames = int(sras.n_frames[angle_idx]) n_frames = int(sras.n_frames[angle_idx])
@@ -635,9 +655,11 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
freq_axis = sras.freq_axis_mhz(n_fft) freq_axis = sras.freq_axis_mhz(n_fft)
img = np.zeros((n_rows, n_frames), dtype=np.float32) img = np.zeros((n_rows, n_frames), dtype=np.float32)
_n_workers = os.cpu_count() or 4 _n_workers = os.cpu_count() or 4
n_fft_bins = n_fft if n_fft is not None else sras.samples_per_frame
chunk_rows = _chunk_rows_for(n_frames, max(sras.samples_per_frame, n_fft_bins))
for r0 in range(0, n_rows, _CHUNK_ROWS): for r0 in range(0, n_rows, chunk_rows):
r1 = min(r0 + _CHUNK_ROWS, n_rows) r1 = min(r0 + chunk_rows, n_rows)
# DC mask for this chunk (float32 expansion is only chunk-sized) # DC mask for this chunk (float32 expansion is only chunk-sized)
dc4_raw = data[r0:r1, CH4_IDX, :, :].astype(np.float32) dc4_raw = data[r0:r1, CH4_IDX, :, :].astype(np.float32)
@@ -1720,10 +1742,8 @@ class SrasViewerWindow(QMainWindow):
def _load_file(self, path: str): def _load_file(self, path: str):
if self._load_thread is not None: if self._load_thread is not None:
return return
self.btn_open.setEnabled(False) # Claim self._load_thread before any call below that can pump the
self.statusBar().showMessage(f"Loading {Path(path).name}") # Qt event loop — see the comment in _start_compute for why.
self._show_progress(f"Loading {Path(path).name}")
self._load_worker = LoadWorker(path) self._load_worker = LoadWorker(path)
self._load_thread = QThread() self._load_thread = QThread()
self._load_worker.moveToThread(self._load_thread) self._load_worker.moveToThread(self._load_thread)
@@ -1733,9 +1753,22 @@ class SrasViewerWindow(QMainWindow):
lambda msg: self.statusBar().showMessage(f"Error: {msg}") lambda msg: self.statusBar().showMessage(f"Error: {msg}")
) )
self._load_worker.finished.connect(self._load_thread.quit) self._load_worker.finished.connect(self._load_thread.quit)
self._load_thread.finished.connect(lambda: setattr(self, "_load_thread", None)) self._load_thread.finished.connect(self._on_load_thread_finished)
self.btn_open.setEnabled(False)
self.statusBar().showMessage(f"Loading {Path(path).name}")
self._show_progress(f"Loading {Path(path).name}")
self._load_thread.start() self._load_thread.start()
def _on_load_thread_finished(self):
# See the comment in _on_compute_thread_finished: wait() before
# releasing our reference to avoid destroying a QThread whose OS
# thread hasn't fully joined yet.
if self._load_thread is not None:
self._load_thread.wait()
self._load_thread = None
def _on_load_done(self, sras): def _on_load_done(self, sras):
self._close_progress() self._close_progress()
self.btn_open.setEnabled(True) self.btn_open.setEnabled(True)
@@ -2044,9 +2077,13 @@ class SrasViewerWindow(QMainWindow):
self._pending_bg_sub = apply_bg_sub self._pending_bg_sub = apply_bg_sub
self._pending_fft_pad_factor = pad_factor self._pending_fft_pad_factor = pad_factor
self.statusBar().showMessage("Computing image…") # Claim self._compute_thread *before* anything below that can pump
self._show_progress("Computing image…") # the Qt event loop (e.g. QProgressDialog.show() on first display).
# If that happened first, a re-entrant editingFinished/signal could
# slip past the guard above, start a second thread, and then have
# this call's own assignment clobber (and destroy while still
# running) that second thread's QThread object — which aborts the
# process. Assigning immediately closes that window.
self._compute_worker = ComputeWorker( self._compute_worker = ComputeWorker(
self._sras, angle_idx, ch_idx, threshold_mv, grating_um, apply_bg_sub, self._sras, angle_idx, ch_idx, threshold_mv, grating_um, apply_bg_sub,
n_fft=n_fft, n_fft=n_fft,
@@ -2060,9 +2097,20 @@ class SrasViewerWindow(QMainWindow):
) )
self._compute_worker.finished.connect(self._compute_thread.quit) self._compute_worker.finished.connect(self._compute_thread.quit)
self._compute_thread.finished.connect(self._on_compute_thread_finished) self._compute_thread.finished.connect(self._on_compute_thread_finished)
self.statusBar().showMessage("Computing image…")
self._show_progress("Computing image…")
self._compute_thread.start() self._compute_thread.start()
def _on_compute_thread_finished(self): def _on_compute_thread_finished(self):
# Block until the OS thread has actually joined before dropping our
# last reference — deallocating a QThread whose thread hasn't fully
# terminated yet logs "QThread: Destroyed while thread is still
# running" and aborts the process. The finished() signal fires as
# the thread is winding down but does not guarantee it has joined.
if self._compute_thread is not None:
self._compute_thread.wait()
self._compute_thread = None self._compute_thread = None
angle_idx = self.spin_angle.value() angle_idx = self.spin_angle.value()
ch_idx = self.combo_channel.currentIndex() ch_idx = self.combo_channel.currentIndex()
@@ -2192,6 +2240,8 @@ class SrasViewerWindow(QMainWindow):
def _on_preprocess(self): def _on_preprocess(self):
if self._sras is None: if self._sras is None:
return return
if self._preprocess_thread is not None:
return
s = self._sras s = self._sras
if s.version == 6: if s.version == 6:
self.statusBar().showMessage( self.statusBar().showMessage(
@@ -2205,17 +2255,17 @@ class SrasViewerWindow(QMainWindow):
) )
if not dest: if not dest:
return return
if self._preprocess_thread is not None:
return # a second trigger snuck in while the save dialog was open
apply_bg = self.chk_bg_sub.isChecked() and s.background is not None apply_bg = self.chk_bg_sub.isChecked() and s.background is not None
n_total = s.n_angles n_total = s.n_angles
n_px = int(s.n_rows[0]) * int(s.n_frames[0]) n_px = int(s.n_rows[0]) * int(s.n_frames[0])
approx_mb = n_total * n_px * 3 * 4 / 1e6 approx_mb = n_total * n_px * 3 * 4 / 1e6
self._show_progress(
f"Pre-processing {n_total} angle(s) "
f"({int(s.n_rows[0])}×{int(s.n_frames[0])} px each, ~{approx_mb:.0f} MB output)…"
)
# Claim self._preprocess_thread before any call below that can pump
# the Qt event loop — see the comment in _start_compute for why.
self._preprocess_worker = PreprocessWorker(s, dest, apply_bg) self._preprocess_worker = PreprocessWorker(s, dest, apply_bg)
self._preprocess_thread = QThread() self._preprocess_thread = QThread()
self._preprocess_worker.moveToThread(self._preprocess_thread) self._preprocess_worker.moveToThread(self._preprocess_thread)
@@ -2224,7 +2274,13 @@ class SrasViewerWindow(QMainWindow):
self._preprocess_worker.finished.connect(self._on_preprocess_done) self._preprocess_worker.finished.connect(self._on_preprocess_done)
self._preprocess_worker.finished.connect(self._preprocess_thread.quit) self._preprocess_worker.finished.connect(self._preprocess_thread.quit)
self._preprocess_thread.finished.connect(self._on_preprocess_thread_finished) self._preprocess_thread.finished.connect(self._on_preprocess_thread_finished)
self._preprocess_act.setEnabled(False) self._preprocess_act.setEnabled(False)
self._show_progress(
f"Pre-processing {n_total} angle(s) "
f"({int(s.n_rows[0])}×{int(s.n_frames[0])} px each, ~{approx_mb:.0f} MB output)…"
)
self._preprocess_thread.start() self._preprocess_thread.start()
def _on_preprocess_progress(self, pct: int): def _on_preprocess_progress(self, pct: int):
@@ -2239,6 +2295,11 @@ class SrasViewerWindow(QMainWindow):
self.statusBar().showMessage("v5 file written — re-open it for instant display.") self.statusBar().showMessage("v5 file written — re-open it for instant display.")
def _on_preprocess_thread_finished(self): def _on_preprocess_thread_finished(self):
# See the comment in _on_compute_thread_finished: wait() before
# releasing our reference to avoid destroying a QThread whose OS
# thread hasn't fully joined yet.
if self._preprocess_thread is not None:
self._preprocess_thread.wait()
self._preprocess_thread = None self._preprocess_thread = None
self._preprocess_act.setEnabled( self._preprocess_act.setEnabled(
self._sras is not None and self._sras.version != 6) self._sras is not None and self._sras.version != 6)