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

Bound nested parallelism, make compute cancellable, harden batch pool

Browse Source 
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>
This commit is contained in:
Thomas Ales
2026-07-31 00:09:46 -05:00
parent 55a4c5e42f
commit 9a7cc396e7
6 changed files with 909 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sras_compute.py
+98 -25
View File
@@ -66,7 +66,11 @@ def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1,
# per-chunk size cannot buy more concurrency — the worker count must be derived
# from the budget instead. See _plan_chunks.
_TOTAL_BYTES_BUDGET = int(os.environ.get("SRAS_MEM_BUDGET_MB", 1536)) * 1024 * 1024
# 1024 MB is the measured knee on a 16-core machine against a 7507-frame x
# 2500-sample angle: 512 MB leaves ~20% of the FFT speedup on the table, and
# 1536+ MB costs ~0.4 GB more resident for no further gain. Override with
# SRAS_MEM_BUDGET_MB on a smaller machine.
_TOTAL_BYTES_BUDGET = int(os.environ.get("SRAS_MEM_BUDGET_MB", 1024)) * 1024 * 1024
_CHUNK_ROWS_MAX = 32 # cap for small scans (original behavior)
_MAX_WORKERS = int(os.environ.get("SRAS_MAX_WORKERS", 0)) or (os.cpu_count() or 4)
@@ -83,21 +87,38 @@ def _chunk_rows_for(n_frames: int, samples_per_frame: int,
def _plan_chunks(n_rows: int, n_frames: int, samples_per_frame: int,
live_multiplier: int = 1,
max_workers: int | None = None) -> tuple[int, int]:
"""(chunk_rows, n_workers) sized so live_multiplier x chunk_rows x
n_frames x samples x 4 bytes x n_workers stays within the budget."""
per_chunk_budget = max(1, _TOTAL_BYTES_BUDGET // live_multiplier)
chunk_rows = _chunk_rows_for(n_frames, samples_per_frame, per_chunk_budget)
max_workers: int | None = None,
budget: int | None = None) -> tuple[int, int]:
"""(chunk_rows, n_workers) such that all concurrent chunks together fit
the budget: n_workers x live_multiplier x chunk_rows x n_frames x
samples x 4 bytes <= budget.
chunk_bytes = max(1, live_multiplier * chunk_rows * n_frames * samples_per_frame * 4)
The worker count is chosen *first* and the chunk sized to it. Sizing the
chunk first is the trap: _chunk_rows_for spends whatever budget it is
given, so a single chunk would always consume the lot and leave room for
exactly one worker — no concurrency, on precisely the large scans that
need it most.
A caller that is itself running several of these concurrently must pass
*both* max_workers=1 and its share of the budget. Capping the workers
alone is not enough: the chunk would still be sized against the whole
budget, and N concurrent callers would each allocate all of it.
"""
total = _TOTAL_BYTES_BUDGET if budget is None else max(1, budget)
cap = max_workers if max_workers is not None else _MAX_WORKERS
n_workers = int(max(1, min(cap,
_TOTAL_BYTES_BUDGET // chunk_bytes,
-(-n_rows // chunk_rows))))
bytes_per_row = max(1, live_multiplier * n_frames * samples_per_frame * 4)
# A chunk is at least one row, so that alone caps how many can be live.
n_workers = int(max(1, min(cap, n_rows, total // bytes_per_row)))
chunk_rows = _chunk_rows_for(n_frames, samples_per_frame,
max(1, total // (n_workers * live_multiplier)))
# With chunk_rows known, more workers than chunks buys nothing.
n_workers = int(max(1, min(n_workers, -(-n_rows // chunk_rows))))
return chunk_rows, n_workers
def _map_row_chunks(n_rows: int, chunk_rows: int, n_workers: int, fn):
def _map_row_chunks(n_rows: int, chunk_rows: int, n_workers: int, fn,
should_stop=None):
"""Apply fn(r0, r1) over row chunks, in parallel when it pays.
Chunks write to disjoint output slices, so no locking is needed. numpy
@@ -105,48 +126,94 @@ def _map_row_chunks(n_rows: int, chunk_rows: int, n_workers: int, fn):
threads give real parallelism here — and on a very large file they also
keep many more page-fault requests in flight, which is what the I/O path
wants.
*should_stop* is polled per chunk so a cancelled job abandons work within
one chunk rather than one whole angle — on a large scan an angle is
~40 s, which is far too long to block application shutdown.
"""
bounds = [(r0, min(r0 + chunk_rows, n_rows))
for r0 in range(0, n_rows, chunk_rows)]
def run(b):
if should_stop is not None and should_stop():
return
fn(*b)
if n_workers <= 1 or len(bounds) == 1:
for r0, r1 in bounds:
fn(r0, r1)
for b in bounds:
run(b)
return
with ThreadPoolExecutor(max_workers=min(n_workers, len(bounds))) as pool:
list(pool.map(lambda b: fn(*b), bounds))
list(pool.map(run, bounds))
# ---------------------------------------------------------------------------
# Image computation
# ---------------------------------------------------------------------------
def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray:
"""Mean of each waveform → (n_rows, n_frames) float32, in ADC counts."""
def plan_angle_level(sras: SrasFile,
live_multiplier: int = 1) -> tuple[int, int]:
"""(n_angle_workers, per_angle_budget) for a caller that parallelises
over angles instead of over rows.
Callers that parallelise over angles must not also let each angle
parallelise over rows: the two levels multiply, in threads and in memory.
Each angle is therefore given max_workers=1 and the returned budget share,
which together keep total live buffers within _TOTAL_BYTES_BUDGET.
"""
biggest = max(range(sras.n_angles), key=lambda a: int(sras.n_frames[a]))
_, n_workers = _plan_chunks(int(sras.n_rows[biggest]),
int(sras.n_frames[biggest]),
sras.samples_per_frame,
live_multiplier=live_multiplier)
n_workers = max(1, min(n_workers, sras.n_angles))
return n_workers, max(1, _TOTAL_BYTES_BUDGET // n_workers)
def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int,
max_workers: int | None = None,
budget: int | None = None,
should_stop=None) -> np.ndarray:
"""Mean of each waveform → (n_rows, n_frames) float32, in ADC counts.
Pass max_workers=1 when the caller is already parallelising over angles.
If *should_stop* ever returns True the result is incomplete — the caller
is expected to be abandoning it.
"""
n_rows, n_frames = sras.image_shape(angle_idx)
data = sras.data[angle_idx]
chunk_rows, n_workers = _plan_chunks(n_rows, n_frames, sras.samples_per_frame)
chunk_rows, n_workers = _plan_chunks(n_rows, n_frames, sras.samples_per_frame,
max_workers=max_workers, budget=budget)
img = np.empty((n_rows, n_frames), dtype=np.float32)
def chunk(r0: int, r1: int):
img[r0:r1] = data[r0:r1, ch_idx, :, :].astype(np.float32).mean(axis=-1)
_map_row_chunks(n_rows, chunk_rows, n_workers, chunk)
_map_row_chunks(n_rows, chunk_rows, n_workers, chunk, should_stop=should_stop)
return img
def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray:
def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int,
max_workers: int | None = None, budget: int | None = None,
should_stop=None) -> np.ndarray:
"""DC image for (angle, channel) in mV, preferring a stored v5/v7 cache."""
cached = sras.cached_dc_mv(angle_idx, ch_idx)
if cached is not None:
return cached
return adc_to_mv(compute_dc_image(sras, angle_idx, ch_idx), *sras.cal(ch_idx))
return adc_to_mv(
compute_dc_image(sras, angle_idx, ch_idx, max_workers=max_workers,
budget=budget, should_stop=should_stop),
*sras.cal(ch_idx))
def compute_rf_image(sras: SrasFile, angle_idx: int,
dc_threshold_mv: float | None,
apply_bg_sub: bool = True,
n_fft: int | None = None,
dc4_mv: np.ndarray | None = None) -> np.ndarray:
dc4_mv: np.ndarray | None = None,
max_workers: int | None = None,
budget: int | None = None,
should_stop=None) -> np.ndarray:
"""FFT of each CH1 waveform; pixel = peak frequency in MHz.
Pixels where CH4_dc < dc_threshold_mv are set to 0 — and the FFT is
@@ -193,7 +260,8 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
n_fft_bins = n_fft if n_fft is not None else sras.samples_per_frame
chunk_rows, n_workers = _plan_chunks(
n_rows, n_frames, max(sras.samples_per_frame, n_fft_bins),
live_multiplier=_FFT_LIVE_MULTIPLIER)
live_multiplier=_FFT_LIVE_MULTIPLIER, max_workers=max_workers,
budget=budget)
background = sras.background if (apply_bg_sub and sras.background is not None) else None
cal4 = sras.cal(CH4_IDX)
@@ -239,7 +307,7 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
else:
img[r0:r1] = freq_axis[peak_bins]
_map_row_chunks(n_rows, chunk_rows, n_workers, chunk)
_map_row_chunks(n_rows, chunk_rows, n_workers, chunk, should_stop=should_stop)
return img
@@ -258,6 +326,8 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool,
"""
global _MAX_WORKERS
try:
if mode not in ("dc", "fft"):
return f"unknown cache mode {mode!r} (expected 'dc' or 'fft')"
set_fft_backend(fft_backend)
if max_workers:
_MAX_WORKERS = max_workers
@@ -493,7 +563,8 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int,
background-thread workers must not touch GUI-thread-owned caches."""
n = sras.n_angles # already the *complete*-angle count for aborted v6 scans
dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx)
n_workers = max(1, min(_MAX_WORKERS, n))
# Parallel over angles, serial within each — see plan_angle_level.
n_workers, angle_budget = plan_angle_level(sras)
# progress_cb fires from pool threads, so the counter behind it must be
# atomic. list.append is, and len() of a list is a consistent read.
@@ -506,7 +577,9 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int,
# ---- Step 1: binarized CH4 mask per angle, native per-angle grid -----
def mask_for(a: int) -> np.ndarray:
dc4 = adc_to_mv(compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX))
dc4 = adc_to_mv(
compute_dc_image(sras, a, CH4_IDX, max_workers=1, budget=angle_budget),
*sras.cal(CH4_IDX))
m = (dc4 >= dc_threshold_mv).astype(np.float32)
tick(0, 25)
return m