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 sys
import struct
import faulthandler
import numpy as np
from pathlib import Path
import os
faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc.
from PyQt6.QtWidgets import (
QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
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:
"""Mean of each waveform → (n_rows, n_frames) float32.
Processes in row chunks so the float32 working buffer is bounded to
``_CHUNK_ROWS × n_frames × spf × 4`` bytes regardless of scan size.
Processes in row chunks sized to a fixed memory budget (see
``_chunk_rows_for``) so the float32 working buffer stays bounded
regardless of scan size.
"""
n_rows = int(sras.n_rows[angle_idx])
n_frames = int(sras.n_frames[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)
for r0 in range(0, n_rows, _CHUNK_ROWS):
r1 = min(r0 + _CHUNK_ROWS, n_rows)
for r0 in range(0, n_rows, chunk_rows):
r1 = min(r0 + chunk_rows, n_rows)
img[r0:r1] = (
data[r0:r1, ch_idx, :, :]
.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
images are used directly — no FFT is run.
Otherwise, data is processed in row chunks of ``_CHUNK_ROWS`` to bound
peak RAM to roughly ``_CHUNK_ROWS × n_frames × max(spf, n_fft) × 24``
bytes regardless of scan size.
Otherwise, data is processed in row chunks sized to a fixed memory
budget (see ``_chunk_rows_for``) to bound peak RAM regardless of scan
size.
"""
n_rows = int(sras.n_rows[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)
img = np.zeros((n_rows, n_frames), dtype=np.float32)
_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):
r1 = min(r0 + _CHUNK_ROWS, n_rows)
for r0 in range(0, n_rows, chunk_rows):
r1 = min(r0 + chunk_rows, n_rows)
# DC mask for this chunk (float32 expansion is only chunk-sized)
dc4_raw = data[r0:r1, CH4_IDX, :, :].astype(np.float32)
@@ -1720,10 +1742,8 @@ class SrasViewerWindow(QMainWindow):
def _load_file(self, path: str):
if self._load_thread is not None:
return
self.btn_open.setEnabled(False)
self.statusBar().showMessage(f"Loading {Path(path).name}")
self._show_progress(f"Loading {Path(path).name}")
# Claim self._load_thread before any call below that can pump the
# Qt event loop — see the comment in _start_compute for why.
self._load_worker = LoadWorker(path)
self._load_thread = QThread()
self._load_worker.moveToThread(self._load_thread)
@@ -1733,9 +1753,22 @@ class SrasViewerWindow(QMainWindow):
lambda msg: self.statusBar().showMessage(f"Error: {msg}")
)
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()
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):
self._close_progress()
self.btn_open.setEnabled(True)
@@ -2044,9 +2077,13 @@ class SrasViewerWindow(QMainWindow):
self._pending_bg_sub = apply_bg_sub
self._pending_fft_pad_factor = pad_factor
self.statusBar().showMessage("Computing image…")
self._show_progress("Computing image…")
# Claim self._compute_thread *before* anything below that can pump
# 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._sras, angle_idx, ch_idx, threshold_mv, grating_um, apply_bg_sub,
n_fft=n_fft,
@@ -2060,9 +2097,20 @@ class SrasViewerWindow(QMainWindow):
)
self._compute_worker.finished.connect(self._compute_thread.quit)
self._compute_thread.finished.connect(self._on_compute_thread_finished)
self.statusBar().showMessage("Computing image…")
self._show_progress("Computing image…")
self._compute_thread.start()
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
angle_idx = self.spin_angle.value()
ch_idx = self.combo_channel.currentIndex()
@@ -2192,6 +2240,8 @@ class SrasViewerWindow(QMainWindow):
def _on_preprocess(self):
if self._sras is None:
return
if self._preprocess_thread is not None:
return
s = self._sras
if s.version == 6:
self.statusBar().showMessage(
@@ -2205,17 +2255,17 @@ class SrasViewerWindow(QMainWindow):
)
if not dest:
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
n_total = s.n_angles
n_px = int(s.n_rows[0]) * int(s.n_frames[0])
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_thread = QThread()
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._preprocess_thread.quit)
self._preprocess_thread.finished.connect(self._on_preprocess_thread_finished)
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()
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.")
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_act.setEnabled(
self._sras is not None and self._sras.version != 6)