diff --git a/CHANGELOG.md b/CHANGELOG.md index 08abd440..0cbace71 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,69 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed + +- **`torch.compile` has never taken effect in `predict`.** `run_inference` + assigned the compiled module to `model_wrapper.model`, but `forward_batch` + calls `self.forward_module` (separate attributes — see + `ModelInferenceWrapper.__init__`), so every "compiled" run went on executing + the original eager module. Measured on the 20-class production bundle at + batch 1024 on an A30: 16,556 chunks/s assigning `.model` against 16,571 + chunks/s with no compile at all, i.e. 1.00x. + + Two further problems sat behind it. Compilation was skipped outright + whenever a CL-regression repr-capture hook was installed, which is every + production multiclass bundle — but only *CUDA graphs* cannot carry the + hook's Python side effect; plain inductor graph-breaks at the hook and still + fuses around it, worth 1.81x against the 1.83x no-hook ceiling (30,019 vs + 30,227 vs 16,571 chunks/s eager). And every mega-batch ends in a short + batch, so an unpadded run presents ~19 distinct batch shapes and recompiles + for each; `_stack_to_device` now takes `pad_to` and the batch runners slice + the padding rows back off. The model is per-sample throughout (convolutions, + GroupNorm/LayerNorm, BatchNorm in eval), so padding cannot change a real + row's result. + + The auto-enable threshold moves from 5,000 reads to 2,000,000. Compiling + this model costs ~60s wall and buys ~10% at the size `predict` is deployed + at, because extraction and not the GPU is that configuration's bottleneck — + so the old threshold made compilation a large net loss on ordinary runs: a + 176k-read sample measured 45.8s uncompiled against 101.5s compiled. + +- **Extraction sub-batching was dead code.** `_SUB_BATCH_SIZE` was fixed at + 50,000 reads against a default `read_batch_size` of 10,000, so the loop that + exists to feed the GPU continuously always ran exactly once and the consumer + saw nothing until a whole mega-batch had been extracted — the failure mode + escapepod-rs#361 documents on its own GPU pipeline. It is now a real, + configurable size (`LEECH_PREDICT_EXTRACT_CHUNK`), defaulting to the whole + mega-batch because splitting only pays on an allocation with idle cores. + +### Changed + +- **BAM writes go to a dedicated writer thread behind a bounded queue** rather + than a single future the consumer blocked on before it could submit the next + one. pysam is still touched by exactly one thread and writes still land in + mega-batch order; what changed is that the consumer no longer waits. Writing + a mega-batch takes longer than producing one (pysam tagging runs at ~7k + reads/s), so that wait was 22.9s of a 46s run on a 176k-read sample, and it + is when the GPU went idle — which is what made utilization arrive in bursts. + BGZF compression also moves into an htslib thread pool + (`LEECH_PREDICT_BAM_THREADS`) to get deflate off the GIL. + +- **GPU batches queue `LEECH_PREDICT_GPU_IN_FLIGHT` deep** instead of a + one-deep handshake that let the thread filling batches run at most one batch + ahead of the GPU. The executor is still `max_workers=1`, which is + load-bearing: it keeps GPU calls sequential, keeps `_stack_to_device`'s + thread-local pinned staging buffer single-owner, and keeps `pending[read_id]` + in extraction order. + +- **`cap_rayon_threads_for_slurm` reserves CPUs proportionally, not a flat + six.** `avail - 6` left 1 extraction thread at `--cpus-per-task 4` and 2 at + 8, on the pipeline stage that is ~70% of the job's CPU. It now reserves a + quarter, capped at 4. Measured on a 176k-read sample, extraction scales + cleanly with this number up to the allocation's *physical* core count and + only ~12% further across the hyperthread siblings: 149.1s at 2 threads, + 76.5s at 4, 40.1s at 8, 35.8s at 14, 34.3s at 16 — on 8 physical cores. + ## [0.11.1] - 2026-09-12 ### Changed diff --git a/src/leech/inference/helpers.py b/src/leech/inference/helpers.py index baaab0e2..387df1be 100644 --- a/src/leech/inference/helpers.py +++ b/src/leech/inference/helpers.py @@ -626,7 +626,9 @@ def prepare_signal_channels(chunk: dict, signal_len: int) -> np.ndarray: _pinned_staging = threading.local() -def _stack_to_device(arrays: list[np.ndarray], device: str, slot: Hashable) -> torch.Tensor: +def _stack_to_device( + arrays: list[np.ndarray], device: str, slot: Hashable, pad_to: int | None = None +) -> torch.Tensor: """Stack ``arrays`` into one tensor on ``device``. On CUDA the stack lands in a per-thread pinned staging buffer (keyed by @@ -638,22 +640,38 @@ def _stack_to_device(arrays: list[np.ndarray], device: str, slot: Hashable) -> t have the row shape and dtype of ``arrays[0]`` — which is what every caller produces (float32 throughout) and what plain ``np.stack`` requires for the shape anyway. + + ``pad_to`` zero-pads the batch dimension up to a fixed size. Only the + compiled path passes it, and it is what makes compiling worth anything + here: every mega-batch ends in a short batch, so an uncompiled-shape run + presents ~19 distinct batch sizes and ``torch.compile`` recompiles for each + one. Measured on the production bundle, that turned a 1.8x model speedup + into a 2.4x end-to-end *slowdown* (35s -> 83s). The model is per-sample + throughout (convolutions, GroupNorm/LayerNorm, BatchNorm in eval), so the + padding rows cannot affect the real ones; callers slice the output back. """ + n = len(arrays) + total = n if pad_to is None else max(n, pad_to) + if not device.startswith("cuda"): - return torch.from_numpy(np.stack(arrays)).to(device) + stacked = np.stack(arrays) + if total > n: + stacked = np.concatenate( + [stacked, np.zeros((total - n,) + stacked.shape[1:], dtype=stacked.dtype)] + ) + return torch.from_numpy(stacked).to(device) first = arrays[0] - n = len(arrays) key = (slot, first.shape, first.dtype.str) cache = getattr(_pinned_staging, "buffers", None) if cache is None: cache = {} _pinned_staging.buffers = cache entry = cache.get(key) - if entry is None or entry[0].shape[0] < n: + if entry is None or entry[0].shape[0] < total: entry = ( torch.empty( - (n,) + first.shape, + (total,) + first.shape, dtype=torch.from_numpy(first).dtype, pin_memory=True, ), @@ -664,8 +682,11 @@ def _stack_to_device(arrays: list[np.ndarray], device: str, slot: Hashable) -> t # Do not overwrite a buffer whose last copy is still in flight. entry[1].synchronize() buf, copied = entry - view = buf[:n] - np.stack(arrays, out=view.numpy()) + view = buf[:total] + np_view = view.numpy() + np.stack(arrays, out=np_view[:n]) + if total > n: + np_view[n:] = 0 out = view.to(device, non_blocking=True) copied.record() return out @@ -682,19 +703,23 @@ def _run_batch_multiclass( pending: dict[str, list[tuple[int, int, float, list[float], float | None]]], calibration: dict | None = None, cl_regression_head: "torch.nn.Module | None" = None, + pad_to: int | None = None, ) -> None: """Run a multi-class batch: store (base_idx, class_idx, confidence, all_probs, cl_pred) per read.""" - signal_t = _stack_to_device(signals, device, "signal") - seq_t = _stack_to_device(sequences, device, "sequence") + n = len(meta) + signal_t = _stack_to_device(signals, device, "signal", pad_to) + seq_t = _stack_to_device(sequences, device, "sequence", pad_to) batch = {"signal": signal_t, "sequence": seq_t} if requires_features: valid_feats = [f for f in features if f is not None] if valid_feats: - batch["features"] = _stack_to_device(valid_feats, device, "features") + batch["features"] = _stack_to_device(valid_feats, device, "features", pad_to) with torch.inference_mode(): logits = model_wrapper.forward_batch(batch, device) + # Drop the padding rows before anything reads the batch dimension. + logits = logits[:n] if calibration is not None: from leech.calibration import apply_calibration @@ -710,7 +735,7 @@ def _run_batch_multiclass( and isinstance(model_wrapper, ModelInferenceWrapper) and model_wrapper.captured_repr is not None ): - cl_preds = cl_regression_head(model_wrapper.captured_repr).cpu().numpy() + cl_preds = cl_regression_head(model_wrapper.captured_repr[:n]).cpu().numpy() # One `tolist()` for the whole batch, not `float(p)` per class per chunk. # numpy promotes float32 -> Python float identically either way. @@ -733,19 +758,21 @@ def _run_batch( requires_features: bool, device: str, pending: dict[str, list[tuple[int, float]]], + pad_to: int | None = None, ) -> None: """Run a batch through the model and accumulate results into pending.""" - signal_t = _stack_to_device(signals, device, "signal") - seq_t = _stack_to_device(sequences, device, "sequence") + n = len(meta) + signal_t = _stack_to_device(signals, device, "signal", pad_to) + seq_t = _stack_to_device(sequences, device, "sequence", pad_to) batch = {"signal": signal_t, "sequence": seq_t} if requires_features: valid_feats = [f for f in features if f is not None] if valid_feats: - batch["features"] = _stack_to_device(valid_feats, device, "features") + batch["features"] = _stack_to_device(valid_feats, device, "features", pad_to) with torch.inference_mode(): - logits = model_wrapper.forward_batch(batch, device) + logits = model_wrapper.forward_batch(batch, device)[:n] probs = torch.sigmoid(logits).cpu().numpy().flatten() for (read_id, base_idx), prob in zip(meta, probs.tolist(), strict=True): @@ -868,12 +895,114 @@ def cap_rayon_threads_for_slurm(max_cap: int | None = None) -> int: if max_cap is not None: avail = min(avail, max_cap) if "RAYON_NUM_THREADS" not in os.environ: - rayon_threads = max(1, avail - 6) # reserve headroom for main + GPU I/O + # Reserve headroom for the threads that are *not* rayon (consumer, GPU + # worker, BAM writer, POD5 prefetch), but proportionally rather than as + # a flat subtraction. `avail - 6` reserved the same six CPUs whatever + # the allocation, which is most of a small one: it left 1 extraction + # thread at `--cpus-per-task 4` and 2 at 8, and extraction is where + # this pipeline spends ~70% of its CPU. Measured on a 176k-read sample + # (A30, one GPU), Rust extraction scales cleanly with this number to + # the allocation's *physical* core count and only ~12% further across + # the hyperthread siblings: 149.1s at 2 threads, 76.5s at 4, 40.1s at + # 8, 35.8s at 14, 34.3s at 16 on 8 physical cores. Reserving a quarter + # (min 1, max 4) keeps the other threads fed without capping extraction + # at a fraction of the job. + rayon_threads = max(1, avail - min(4, max(1, avail // 4))) os.environ["RAYON_NUM_THREADS"] = str(rayon_threads) logger.info(f"Set RAYON_NUM_THREADS={rayon_threads} (from {avail} available CPUs)") return avail +def _htslib_write_threads() -> int: + """Threads for htslib's BGZF compression on the output BAM. + + Small on purpose: these threads compete with Rust extraction for the same + cores, and compression is not the pipeline's bottleneck -- getting it off + the GIL is. Scales with the allocation and is overridable with + ``LEECH_PREDICT_BAM_THREADS``. + """ + import os + + raw = os.environ.get("LEECH_PREDICT_BAM_THREADS") + if raw is not None: + try: + return max(1, int(raw)) + except ValueError: + logger.warning(f"LEECH_PREDICT_BAM_THREADS={raw!r} is not an integer; ignoring") + # One extra thread per 8 allocated CPUs, capped at 2. The deployment this + # is sized for gives each job 16 logical CPUs -- 8 physical cores, because + # four such jobs share one 64-core/4-GPU node -- and at that size the + # process already runs at ~90% of the box's effective core throughput. + # Extra compression threads there take cores from extraction, which is 70% + # of the job's CPU; the win is getting deflate off the GIL, not running + # more of it. + slurm_cpus = int(os.environ.get("SLURM_CPUS_PER_TASK", 0)) or (os.cpu_count() or 4) + return max(1, min(2, slurm_cpus // 8)) + + +def resolve_pipeline_depths(*, read_batch_size: int, batch_size: int) -> tuple[int, int, int, int]: + """Sizes for the four stages of ``run_inference``'s streaming pipeline. + + Returns ``(extract_chunk_reads, extract_queue_depth, gpu_in_flight, + write_queue_depth)``. Each is overridable by the matching + ``LEECH_PREDICT_*`` environment variable, which is how the values below + were swept. + + The pipeline is a chain of four stages on their own threads -- Rust + extraction, batch accumulation, one GPU worker, one BAM writer -- and on a + real workload each stage costs the same order of magnitude, so wall time is + set by how well they overlap rather than by any one of them. Measured on a + 176,210-read sample (A30, one GPU): extraction 17-38s depending on cores, + GPU 20-28s, BAM write 14-23s. + + ``extract_chunk_reads`` splits one mega-batch into several extraction calls + so the consumer can start on the first piece while rayon is still building + the rest. It defaults to the whole mega-batch, i.e. off, and that is a + measured choice rather than caution. Splitting helps only when there are + spare cores for the consumer to use: at this pipeline's deployment size -- + 16 logical CPUs, which is 8 physical cores, four jobs to a 64-core/4-GPU + node -- the process already runs at ~90% of the box's effective core + throughput, so interleaving consumer work with extraction just adds GIL + contention to a saturated machine and cost ~4 s on a 176k-read sample. + Lower it (``LEECH_PREDICT_EXTRACT_CHUNK``) on an allocation with idle + cores, where escapepod-rs#361's finding applies: a prep chunk as large as + the whole superbatch leaves "the GPU consumer... fully idle for the entire + duration of that one giant prep call". + + The three depths follow escapepod-rs#361's other finding: a buffer has to + be deep enough to absorb one producer burst, or the producer blocks on a + full channel exactly as a one-deep handshake would. One prep cycle yields + about ``extract_chunk_reads / batch_size`` GPU batches, which is what + ``gpu_in_flight`` is sized against. + """ + import os + + def _env(name: str, default: int) -> int: + raw = os.environ.get(name) + if raw is None: + return default + try: + value = int(raw) + except ValueError: + logger.warning(f"{name}={raw!r} is not an integer; using {default}") + return default + if value < 1: + logger.warning(f"{name}={value} must be >= 1; using {default}") + return default + return value + + extract_chunk_reads = _env("LEECH_PREDICT_EXTRACT_CHUNK", read_batch_size) + extract_queue_depth = _env("LEECH_PREDICT_EXTRACT_QUEUE", 4) + gpu_in_flight = _env("LEECH_PREDICT_GPU_IN_FLIGHT", 4) + write_queue_depth = _env("LEECH_PREDICT_WRITE_QUEUE", 2) + logger.info( + f"Pipeline depths: extract_chunk={extract_chunk_reads} reads, " + f"extract_queue={extract_queue_depth}, gpu_in_flight={gpu_in_flight}, " + f"write_queue={write_queue_depth}" + ) + return extract_chunk_reads, extract_queue_depth, gpu_in_flight, write_queue_depth + + def build_rust_extraction_kwargs( *, signal_context: tuple[int, int], diff --git a/src/leech/inference/single.py b/src/leech/inference/single.py index e5e6d9ff..d41cc5b5 100644 --- a/src/leech/inference/single.py +++ b/src/leech/inference/single.py @@ -18,6 +18,7 @@ BatchAccumulator, _check_config_consistency, _encode_sequence_for_inference, + _htslib_write_threads, _run_batch, _run_batch_multiclass, _write_mega_batch_predictions, @@ -28,6 +29,7 @@ load_model_auto, prepare_inference_features, prepare_signal_channels, + resolve_pipeline_depths, validate_inference_shapes, ) from leech.io.bam_reader import count_bam_reads, iter_bam_batches @@ -499,7 +501,16 @@ def run_inference( ) logger.info(f"TSV output: {output_path} ({len(_tsv_class_names)} classes)") else: - bam_out = pysam.AlignmentFile(str(output_path), "wb", template=bam_in) + # BGZF compression in htslib's own thread pool rather than on the + # writer thread. The writer is the one stage here that is pure Python + # plus a blocking C call: every record pays `set_tag` under the GIL and + # then a synchronous deflate. Handing the deflate to htslib threads + # takes it off the GIL entirely, which matters because the GPU worker + # needs the GIL for every eager kernel launch and the consumer holds it + # in a per-chunk loop. + bam_out = pysam.AlignmentFile( + str(output_path), "wb", template=bam_in, threads=_htslib_write_threads() + ) bam_in.close() total_reads = 0 @@ -575,11 +586,22 @@ def run_inference( ) mega_batch_idx = 0 - # torch.compile the model for faster inference (CUDA graph + kernel fusion). - # Auto-skip for small runs (<5000 reads) where compilation overhead (~15-30s) - # outweighs the speedup. Also skip when repr capture hooks are active or - # when --no-compile is set. - _COMPILE_THRESHOLD = 5000 + # torch.compile the model for faster inference (kernel fusion, and CUDA + # graphs when nothing blocks them). Auto-skip for runs too short to amortize + # compilation; also skip when --no-compile is set. + # + # Compiling this model costs ~60s wall -- measured on the 20-class + # production bundle on an A30, from the compile call to the first + # mega-batch landing -- and buys ~10% end to end at the size this pipeline + # is deployed at (16 logical CPUs, i.e. 8 physical cores, four jobs sharing + # a 64-core/4-GPU node), because extraction rather than the GPU is that + # configuration's bottleneck. Break-even is therefore around ten minutes of + # inference, ~2M reads at the ~3,800 reads/s it sustains there. The old + # 5,000-read threshold turned it on for essentially every run: a 176k-read + # sample measured 45.8s uncompiled against 101.5s compiled, 60s of + # compilation against a ~4s gain. + _COMPILE_THRESHOLD = 2_000_000 + _compiled = False _has_repr_hook = ( isinstance(model_wrapper, ModelInferenceWrapper) and model_wrapper._repr_hook is not None ) @@ -593,15 +615,34 @@ def run_inference( isinstance(model_wrapper, ModelInferenceWrapper) and device.startswith("cuda") and hasattr(torch, "compile") - and not _has_repr_hook ): + # Assign `forward_module`, not `model`. `forward_batch` calls + # `self.forward_module` (see ModelInferenceWrapper.__init__); replacing + # `self.model` leaves `forward_module` bound to the original eager + # module, so compiling was a no-op on every inference path that went + # through this branch. Measured on the production 20-class bundle, + # batch 1024: 16,556 chunks/s assigning `.model` against 16,571 + # chunks/s with no compile at all -- 1.00x, i.e. nothing. + # + # Only CUDA graphs are incompatible with a repr-capture hook, because + # the hook's Python side effect cannot be replayed by a graph. Plain + # inductor just breaks the graph there and still fuses everything + # around it, so a bundle with a CL-regression head gets essentially the + # whole win rather than none of it: same benchmark, 30,019 chunks/s + # compiled-with-hook against 30,227 chunks/s reduce-overhead-no-hook + # and 16,571 eager -- 1.81x vs 1.83x. The old branch skipped compile + # outright whenever the hook was present, which is every production + # multiclass bundle. + _mode = "default" if _has_repr_hook else "reduce-overhead" try: - model_wrapper.model = torch.compile(model_wrapper.model, mode="reduce-overhead") # ty: ignore[invalid-assignment] - logger.info("torch.compile enabled (mode=reduce-overhead)") + model_wrapper.forward_module = torch.compile(model_wrapper.model, mode=_mode) # ty: ignore[invalid-assignment] + _compiled = True + logger.info( + f"torch.compile enabled (mode={_mode}" + + (", repr capture hook present)" if _has_repr_hook else ")") + ) except Exception as e: logger.warning(f"torch.compile failed, using eager mode: {e}") - elif _has_repr_hook: - logger.info("torch.compile skipped (repr capture hooks incompatible with CUDA graphs)") if num_workers > 0: # ---- Parallel path (mega-batched) ---- @@ -749,59 +790,88 @@ def _run_worker_batch(sigs, seqs, feats, meta) -> None: ) else: - # ---- Sequential path (mega-batched, double-buffered GPU) ---- + # ---- Sequential path (mega-batched, pipelined GPU) ---- + from collections import deque from concurrent.futures import Future, ThreadPoolExecutor + ( + _extract_chunk_reads, + _extract_queue_depth, + _gpu_in_flight, + _write_queue_depth, + ) = resolve_pipeline_depths(read_batch_size=read_batch_size, batch_size=batch_size) + pending: dict[str, list] = {} _shape_validated = False calibration = config.get("calibration") if is_multiclass else None + # Only pad when compiled: a fixed batch shape is what keeps + # `torch.compile` from recompiling on every mega-batch's short tail + # (see `_stack_to_device`). Uncompiled, padding would just be wasted + # forward work on rows nobody reads. + _pad_to = batch_size if _compiled else None _batch_fn = ( functools.partial( _run_batch_multiclass, calibration=calibration, cl_regression_head=_cl_head, + pad_to=_pad_to, ) if is_multiclass - else _run_batch + else functools.partial(_run_batch, pad_to=_pad_to) ) + # One GPU worker, several batches queued behind it. `max_workers=1` is + # load-bearing, not conservatism: it keeps GPU calls sequential, keeps + # `_stack_to_device`'s thread-local pinned staging buffer single-owner, + # and makes the order in which batches append to `pending[read_id]` + # exactly the order they were extracted in. The depth is what changed -- + # see `_gpu_in_flight` below. _gpu_executor = ThreadPoolExecutor(max_workers=1) - _gpu_future: Future | None = None - _bam_write_executor = ThreadPoolExecutor(max_workers=1) - _bam_write_future: Future | None = None + _gpu_futures: deque[Future] = deque() def _submit_gpu_batch(sigs, seqs, feats, meta) -> None: - """Flush callback: hand one batch to the GPU thread (double-buffered). + """Flush callback: hand one batch to the GPU thread. The accumulator has already detached these buffers, so the GPU thread owns them and extraction can keep filling the next batch. + + This used to wait for the *immediately* preceding batch before + submitting, which is a one-deep handshake: the thread that fills + batches could never be more than one batch ahead of the GPU, so + every GPU batch stalled extraction for its full duration -- and a + "GPU batch" here is mostly host work (stack, D2H, scatter), not + kernel time. Queuing `_gpu_in_flight` batches instead lets + extraction run a whole prep cycle ahead; the executor still runs + them one at a time, in order. """ - nonlocal _gpu_future - # Wait for previous GPU batch before submitting next - if _gpu_future is not None: - _gpu_future.result() - # Submit GPU work -- runs while main thread continues extraction - _gpu_future = _gpu_executor.submit( - _batch_fn, - sigs, - seqs, - feats, - meta, - model_wrapper, - requires_features, - device, - pending, + while len(_gpu_futures) >= _gpu_in_flight: + _gpu_futures.popleft().result() + _gpu_futures.append( + _gpu_executor.submit( + _batch_fn, + sigs, + seqs, + feats, + meta, + model_wrapper, + requires_features, + device, + pending, + ) ) accumulator = BatchAccumulator(batch_size, _submit_gpu_batch) def _drain_gpu() -> None: - """Wait for any in-flight GPU batch to complete.""" - nonlocal _gpu_future - if _gpu_future is not None: - _gpu_future.result() - _gpu_future = None + """Wait for every in-flight GPU batch to complete. + + Called at a mega-batch boundary, where `pending` is about to be + snapshotted for writing, so every batch of this mega-batch must + have landed in it. + """ + while _gpu_futures: + _gpu_futures.popleft().result() seq_signal_config = SignalConfig( reverse_signal=reverse_signal, @@ -973,14 +1043,16 @@ def _collect_bam_metadata(aln_batch: list) -> tuple: reference_sequences=reference_sequences, ) - _SUB_BATCH_SIZE = 50_000 # Sub-batch Rust extraction for continuous GPU feeding - def _extract_chunks_from_preloaded(preloaded, rs_meta): - """Yield chunks in sub-batches for continuous GPU feeding. - - Instead of extracting all reads at once (blocking GPU for ~2 min), - process ~25K reads at a time (~15-20s each) so chunks flow to GPU - after each sub-batch completes. + """Yield one *list* of chunks per extraction sub-batch. + + Yielding per sub-batch rather than per chunk is what lets the + producer hand each piece to the consumer as it is built, instead of + accumulating a whole mega-batch first. The sub-batch size is + `_extract_chunk_reads` (see `resolve_pipeline_depths`); it was a + fixed 50,000, which is five times the default `read_batch_size`, so + this loop always ran exactly once and the "continuous GPU feeding" + it was written for never happened. """ ( rs_rids, @@ -995,15 +1067,9 @@ def _extract_chunks_from_preloaded(preloaded, rs_meta): ) = rs_meta assert _rs_extract_chunks_from_preloaded is not None n = len(rs_rids) - for start in range(0, n, _SUB_BATCH_SIZE): - end = min(start + _SUB_BATCH_SIZE, n) - if n > _SUB_BATCH_SIZE: - logger.info( - f" Sub-batch {start // _SUB_BATCH_SIZE + 1}/" - f"{(n + _SUB_BATCH_SIZE - 1) // _SUB_BATCH_SIZE}: " - f"reads {start}-{end} of {n}" - ) - sub_chunks = _rs_extract_chunks_from_preloaded( + for start in range(0, n, _extract_chunk_reads): + end = min(start + _extract_chunk_reads, n) + yield _rs_extract_chunks_from_preloaded( preloaded, read_ids=rs_rids[start:end], sequences=rs_seqs[start:end], @@ -1016,7 +1082,6 @@ def _extract_chunks_from_preloaded(preloaded, rs_meta): reference_sequences=rs_refs[start:end] if anchor == "reference" else None, **_rs_kwargs, ) - yield from sub_chunks def _consume_rust_chunks(chunks): """Iterate Rust chunks into batch buffers, flushing to GPU as needed.""" @@ -1041,50 +1106,96 @@ def _consume_rust_chunks(chunks): _shape_validated = True accumulator.add(sig, seq_arr, feat, (read_id, base_idx)) - def _wait_for_bam_write(): - """Wait for any in-flight async BAM write to complete.""" - nonlocal _bam_write_future - if _bam_write_future is not None: - _bam_write_future.result() - _bam_write_future = None + import queue as _queue + import threading as _threading + import time as _time - def _finalize_mega_batch(aln_batch_to_write): - """Flush GPU, submit async BAM write, update counters. + # One writer thread behind a bounded queue, rather than one write + # behind a single future the consumer blocks on. + # + # pysam is not thread-safe, so there is still exactly one thread + # touching `bam_out` and writes still happen in mega-batch order -- + # that part is unchanged and must stay. What changed is who waits: the + # consumer used to call `.result()` on the previous write before it + # could submit the next one, so whenever writing a mega-batch took + # longer than producing one (it does -- pysam tagging is ~7k reads/s), + # the consumer sat blocked with the GPU queue empty. That stall was + # half the wall clock on a 176k-read sample, and it is what made GPU + # utilization arrive in bursts. + _write_queue: _queue.Queue = _queue.Queue(maxsize=_write_queue_depth) + _WRITE_SENTINEL = object() + _writer_error: BaseException | None = None + + def _bam_writer_loop() -> None: + nonlocal _writer_error + while True: + item = _write_queue.get() + if item is _WRITE_SENTINEL: + return + aln_batch_to_write, write_pending = item + try: + if tsv_writer is not None: + tsv_writer.write_predictions( + aln_batch_to_write, write_pending, int_to_label + ) + else: + _write_mega_batch_predictions( + aln_batch_to_write, + write_pending, + bam_out, + is_multiclass, + int_to_label, + class_names_str, + raw, + min_confidence, + min_margin, + ) + except BaseException as exc: # surfaced by _wait_for_bam_write + _writer_error = exc + return + + _writer_thread = _threading.Thread(target=_bam_writer_loop, daemon=True) + _writer_thread.start() + + def _queue_bam_write(item) -> None: + """Hand one mega-batch to the writer, without risking a deadlock. - BAM writes are overlapped with the next mega-batch's extraction. - We serialize writes (wait for previous) since pysam is not thread-safe. + A plain blocking `put` on a bounded queue hangs forever if the + writer has already died: nothing will ever drain it again. Time the + put out and re-check instead, so a writer failure surfaces as its + own exception rather than as a stalled run. """ + while True: + if _writer_error is not None: + raise RuntimeError("BAM writer thread failed") from _writer_error + try: + _write_queue.put(item, timeout=5.0) + return + except _queue.Full: + continue + + def _wait_for_bam_write(): + """Drain every queued write and stop the writer thread.""" + if _writer_thread.is_alive(): + _queue_bam_write(_WRITE_SENTINEL) + _writer_thread.join() + if _writer_error is not None: + raise RuntimeError("BAM writer thread failed") from _writer_error + + def _finalize_mega_batch(aln_batch_to_write): + """Flush GPU, queue the BAM write, update counters.""" nonlocal total_reads, total_predictions, mega_batch_idx - nonlocal pending, _bam_write_future + nonlocal pending accumulator.flush() _drain_gpu() - # Wait for any previous BAM write (serializes bam_out access) - _wait_for_bam_write() # Swap pending -> snapshot; next mega-batch gets a fresh dict write_pending = pending pending = {} batch_preds = len(write_pending) - # Submit write to background thread - if tsv_writer is not None: - _bam_write_future = _bam_write_executor.submit( - tsv_writer.write_predictions, - aln_batch_to_write, - write_pending, - int_to_label, - ) - else: - _bam_write_future = _bam_write_executor.submit( - _write_mega_batch_predictions, - aln_batch_to_write, - write_pending, - bam_out, - is_multiclass, - int_to_label, - class_names_str, - raw, - min_confidence, - min_margin, - ) + # Blocks only once `_write_queue_depth` mega-batches are already + # queued, which bounds memory without putting the writer on the + # consumer's critical path. + _queue_bam_write((aln_batch_to_write, write_pending)) total_reads += len(aln_batch_to_write) total_predictions += batch_preds mega_batch_idx += 1 @@ -1093,10 +1204,6 @@ def _finalize_mega_batch(aln_batch_to_write): f"wrote {batch_preds} predictions for {len(aln_batch_to_write)} reads" ) - import queue as _queue - import threading as _threading - import time as _time - _t_total_start = _time.perf_counter() with Progress() as progress: @@ -1122,9 +1229,39 @@ def _finalize_mega_batch(aln_batch_to_write): assert _rs_extract_chunks_from_preloaded is not None _SENTINEL = object() - _extraction_queue: _queue.Queue = _queue.Queue(maxsize=2) + _extraction_queue: _queue.Queue = _queue.Queue(maxsize=_extract_queue_depth) _producer_error: BaseException | None = None + def _emit_mega_batch(preloaded, p_meta, p_aln) -> None: + """Push one mega-batch to the consumer, a sub-batch at a time. + + The consumer needs `aln_batch` only to finalize, so it + rides on the last item; every earlier item carries just + chunks. Emitting per sub-batch is the point: the GPU + starts on the first sub-batch while rayon is still + extracting the rest, instead of waiting for the whole + mega-batch to be materialized. + """ + p_rids = p_meta[0] + sub_batches = ( + list(_extract_chunks_from_preloaded(preloaded, p_meta)) + if p_rids + else [] + ) + if not sub_batches: + _extraction_queue.put(([], p_aln, len(p_rids), True)) + return + last = len(sub_batches) - 1 + for i, chunks in enumerate(sub_batches): + _extraction_queue.put( + ( + chunks, + p_aln if i == last else None, + len(p_rids) if i == last else 0, + i == last, + ) + ) + def _extraction_producer(): """Background thread: reads BAM -> metadata -> prefetch -> extract -> queue.""" nonlocal _producer_error @@ -1142,7 +1279,6 @@ def _extraction_producer(): if _prev is not None: p_future, p_meta, p_aln = _prev preloaded = p_future.result() - p_rids = p_meta[0] # Overlap: start metadata for CURRENT batch # while extracting PREVIOUS (Rust releases GIL) @@ -1150,15 +1286,7 @@ def _extraction_producer(): _collect_bam_metadata, aln_batch ) - chunk_list: list = [] - if p_rids: - for chunk in _extract_chunks_from_preloaded( - preloaded, p_meta - ): - chunk_list.append(chunk) - - # Push to queue (blocks if queue full -- backpressure) - _extraction_queue.put((p_aln, chunk_list, len(p_rids))) + _emit_mega_batch(preloaded, p_meta, p_aln) # Get metadata result (should be done by now) rs_meta = _meta_future.result() @@ -1178,13 +1306,7 @@ def _extraction_producer(): # Process final batch if _prev is not None: p_future, p_meta, p_aln = _prev - preloaded = p_future.result() - p_rids = p_meta[0] - chunk_list = [] - if p_rids: - for chunk in _extract_chunks_from_preloaded(preloaded, p_meta): - chunk_list.append(chunk) - _extraction_queue.put((p_aln, chunk_list, len(p_rids))) + _emit_mega_batch(p_future.result(), p_meta, p_aln) _meta_exec.shutdown(wait=True) _prefetch_exec.shutdown(wait=True) @@ -1197,22 +1319,25 @@ def _extraction_producer(): _producer_thread.start() logger.info("Queue-based extraction pipeline started (producer thread)") - # Consumer loop: pull from queue -> GPU -> finalize + # Consumer loop: pull sub-batches -> GPU; finalize on the + # sub-batch flagged as its mega-batch's last. + _t_mb_start = _time.perf_counter() while True: item = _extraction_queue.get() if item is _SENTINEL: break - aln_batch, chunk_list, n_rids = item - _t_mb_start = _time.perf_counter() - - logger.info( - f"Mega-batch: {len(aln_batch)} alignments, {n_rids} for Rust extraction" - ) + chunk_list, aln_batch, n_rids, is_last = item if chunk_list: _consume_rust_chunks(iter(chunk_list)) + if not is_last: + continue _t_consume = _time.perf_counter() + assert aln_batch is not None + logger.info( + f"Mega-batch: {len(aln_batch)} alignments, {n_rids} for Rust extraction" + ) _finalize_mega_batch(aln_batch) _t_finalize = _time.perf_counter() @@ -1220,6 +1345,7 @@ def _extraction_producer(): f" Timing: consume+gpu={_t_consume - _t_mb_start:.2f}s " f"finalize={_t_finalize - _t_consume:.2f}s" ) + _t_mb_start = _t_finalize progress.update( task, advance=0, @@ -1315,18 +1441,18 @@ def _extraction_producer(): f"{total_reads / _t_total:.0f} reads/s)" ) - # wait=True, not False. Every one of these pools is already drained here - # (`_drain_gpu` and `_wait_for_bam_write` above), so waiting costs - # nothing -- but `wait=False` leaves worker threads alive past the - # return, and the parallel path below forks an `mp.Pool`. A fork - # inherits the memory of a process with running threads, including any - # lock those threads hold, but not the threads themselves, so nothing - # ever releases it: calling `run_inference` with `num_workers=0` and - # then with `num_workers>0` in one process hung forever, with no error. + # wait=True, not False, and the writer thread is joined rather than + # left daemonized. Everything here is already drained (`_drain_gpu` + # above), so waiting costs nothing -- but `wait=False` leaves worker + # threads alive past the return, and the parallel path above forks an + # `mp.Pool`. A fork inherits the memory of a process with running + # threads, including any lock those threads hold, but not the threads + # themselves, so nothing ever releases it: calling `run_inference` with + # `num_workers=0` and then with `num_workers>0` in one process hung + # forever, with no error. _extract_pool.shutdown(wait=True) _gpu_executor.shutdown(wait=True) - _wait_for_bam_write() # Ensure final BAM write completes before close - _bam_write_executor.shutdown(wait=True) + _wait_for_bam_write() # Drain queued writes before closing the file if tsv_writer is not None: tsv_writer.close() diff --git a/tests/test_inference.py b/tests/test_inference.py index ada995b4..c6fea854 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -1585,8 +1585,27 @@ def test_the_sequential_teardown_waits(self): source = Path(inference_single.__file__).read_text() teardown = source[source.index("_extract_pool.shutdown") :] - teardown = teardown[: teardown.index("_bam_write_executor.shutdown") + 60] + teardown = teardown[: teardown.index("_wait_for_bam_write()") + 60] assert "wait=False" not in teardown, ( "an executor in the sequential teardown shuts down with wait=False; " "the mp.Pool fork in the parallel path will inherit its held locks" ) + + def test_the_bam_writer_thread_is_joined(self): + """Same hazard, other shape: the writer is a bare thread, not an executor. + + It is created `daemon=True` so a crash elsewhere cannot wedge the + process, which also means nothing joins it implicitly -- and a live + thread at fork time is exactly what this class exists to prevent. + """ + from pathlib import Path + + from leech.inference import single as inference_single + + source = Path(inference_single.__file__).read_text() + waiter = source[source.index("def _wait_for_bam_write") :] + waiter = waiter[: waiter.index("def _finalize_mega_batch")] + assert "_writer_thread.join()" in waiter, ( + "the BAM writer thread is not joined in _wait_for_bam_write; a " + "daemon thread left running will be inherited by the mp.Pool fork" + )