def run_inference(
model_and_config: tuple[torch.nn.Module | ModelInferenceWrapper | RemoraModelWrapper, dict]
| None = None,
model_path: Path | None = None,
pod5_path: Path | None = None,
bam_path: Path | None = None,
output_path: Path | None = None,
device: str = "cuda",
min_mapq: int = 0,
motif: str | None = None,
motif_offset: int = 0,
batch_size: int = 256,
base_justify: str = "center",
reverse_signal: bool = True,
num_workers: int = 0,
chunk_size: int = 100,
anchor: str = "reference",
reference_fasta: Path | None = None,
raw: bool = False,
min_confidence: int = 0,
min_margin: int = 0,
abstain_on_junction_indel: bool = False,
read_batch_size: int = 10_000,
backend: str = "auto",
no_compile: bool = False,
output_format: str = "bam",
copy_tags: list[str] | None = None,
parameter_sources: dict[str, bool] | None = None,
) -> None:
"""
Run inference on POD5 and BAM files.
Supports both leech native models and Remora TorchScript models (auto-detected).
Supports parallel chunk extraction via num_workers > 0.
Args:
model_and_config: Pre-loaded (wrapper_or_model, config) tuple.
model_path: Path to model checkpoint directory or Remora .pt file.
pod5_path: Path to POD5 file with raw signal
bam_path: Path to input BAM file with alignments
output_path: Path to output BAM file with predictions
raw: Write full float probabilities (default: compact uint8)
min_confidence: Confidence threshold in 0-255 uint8 space
min_margin: Margin threshold in 0-255 uint8 space
abstain_on_junction_indel: Multiclass only (BAM ``aa`` tag / TSV
``predicted_aa``). When True, ``min_margin`` is enforced only on
chunks whose motif junction is disrupted (unmapped, or a nonzero
CIGAR indel) rather than on every chunk -- issue #282's compound
abstention rule. A ``junction_indel``/``junction_mapped`` value is
always attached to a multiclass prediction when a motif
measurement was made, regardless of this flag; it only changes
whether that value gates the call.
device: Device for inference
min_mapq: Minimum mapping quality
motif: Optional motif to filter predictions (auto-read from config if None)
motif_offset: Offset within motif for prediction (auto-read from config if 0)
batch_size: Chunks per forward pass
base_justify: Signal justification within focus base
reverse_signal: Whether to reverse signal for RNA
num_workers: Parallel chunk extraction workers (0=sequential).
Only beneficial with GPU inference, where CPU chunk extraction
overlaps with GPU forward passes. For CPU-only inference, the
sequential path (0) is faster due to batched POD5 access and
no multiprocessing overhead.
backend: Extraction backend. "auto" uses Rust if available, "rust"
forces Rust (error if unavailable), "python" forces Python.
chunk_size: Reads per worker batch
anchor: "basecall" or "reference" for reference-anchored mode
reference_fasta: Path to reference FASTA (for reference-anchored mode)
read_batch_size: Reads per mega-batch for memory-bounded streaming (default 50K).
Each mega-batch loads BAM alignments + POD5 signals, runs inference,
writes predictions, then frees memory. Set to 0 to disable (load all).
output_format: "bam" for BAM output with tags, "tsv" for gzipped TSV.
TSV mode requires a multiclass model.
parameter_sources: Optional ``{param_name: was_explicit}`` from
click's ``get_parameter_source``, so an explicit CLI value that
equals its own default is still checked against the model config
rather than assumed to be "not passed" (see
``InferenceSpec.from_config``).
"""
logger.info(f"Extraction backend: {backend}")
# Load model
if model_and_config is not None:
wrapper_or_model, config = model_and_config
elif model_path is not None:
logger.info(f"Loading model from {model_path}")
wrapper_or_model, config = load_model_auto(model_path, device=device)
else:
raise ValueError("Either model_and_config or model_path must be provided")
# Determine if this is a Remora model or a leech model
is_remora = config.get("is_remora", False)
if is_remora:
model_wrapper = wrapper_or_model
elif isinstance(wrapper_or_model, ModelInferenceWrapper):
model_wrapper = wrapper_or_model
else:
model_wrapper = ModelInferenceWrapper(wrapper_or_model, config["model_name"])
# The checkpoint's own directory, for resolving a non-absolute
# dwell_template_table (issue #369) -- a checkpoint is a directory
# (config.json + model_best.pt); a Remora/leech-exported model is a
# single .pt file, so its containing directory is the closest analogue.
# None when the caller passed a pre-loaded (model, config) with no
# model_path of its own -- the recorded value (if any) then stays
# whatever it was, and only resolves if it happens to be absolute.
model_dir: Path | None = None
if model_path is not None:
_model_path = Path(model_path)
model_dir = _model_path if _model_path.is_dir() else _model_path.parent
spec = InferenceSpec.from_config(
config,
model_type=getattr(model_wrapper, "model_type", ""),
motif=motif,
motif_offset=motif_offset,
base_justify=base_justify,
anchor=anchor,
reference_fasta=reference_fasta,
default_refine_scale_iters=-1 if is_remora else 2,
parameter_sources=parameter_sources,
model_dwell_margin=lambda: (
getattr(model_wrapper.model, "dwell_margin", 0)
if hasattr(model_wrapper, "model")
else 0
),
base_dir=model_dir,
)
signal_len = spec.signal_len
kmer_len = spec.kmer_len
seq_encoding = spec.seq_encoding
signal_kmer_context = spec.signal_kmer_context
dwell_offset = spec.dwell_offset
signal_context = spec.signal_context
kmer_context = spec.kmer_context
requires_features = spec.requires_features
wide_features = spec.wide_features
_feature_start = spec.feature_start
_feature_end = spec.feature_end
num_out = spec.num_out
is_multiclass = spec.is_multiclass
motif = spec.motif
motif_offset = spec.motif_offset
base_justify = spec.base_justify
anchor = spec.anchor
reference_fasta = spec.reference_fasta
# Load per-AA dwell templates if this checkpoint was trained with them
# (issue #370). `run_bundle_inference` has always done this;
# `run_inference` had no equivalent, so a template checkpoint
# (num_features = base + n_aa) failed shape validation on the first
# chunk. `load_inference_dwell_template` is the one definition of this
# load, shared with bundle.py.
dwell_template = load_inference_dwell_template(spec, model_dir)
dwell_templates_arr: np.ndarray | None = (
dwell_template.templates if dwell_template is not None else None
)
dwell_template_min_pos: int = dwell_template.min_pos if dwell_template is not None else 0
dwell_row: int | None = dwell_template.dwell_row if dwell_template is not None else None
int_to_label = {v: k for k, v in spec.label_map.items()} if spec.label_map else None
# Signal map refinement setup
refine_signal_map = spec.refine_signal_map
signal_refiner = None
if refine_signal_map:
from leech.data import get_kmer_table
from leech.inference.helpers import _warn_if_kmer_table_drifted
from leech.signal_refine import SigMapRefiner
kmer_table_path = get_kmer_table()
_warn_if_kmer_table_drifted(config.get("kmer_table_sha256"), kmer_table_path)
signal_refiner = SigMapRefiner.from_table(
kmer_table_path,
half_bandwidth=spec.refine_half_bandwidth,
scale_iters=spec.refine_scale_iters,
center_idx=spec.refine_kmer_center_idx,
)
logger.info(
f"Signal map refinement enabled: half_bw={spec.refine_half_bandwidth}, "
f"scale_iters={spec.refine_scale_iters}, center_idx={spec.refine_kmer_center_idx}"
)
logger.info(f"Signal length: {signal_len}, K-mer length: {kmer_len}")
if is_multiclass:
logger.info(f"Multi-class model: num_out={num_out}")
if spec.signal_context_bases is not None:
logger.info(
f"Base-defined signal window: {spec.signal_context_bases} "
f"(padded/centre-cropped to signal_len={signal_len})"
)
else:
logger.info(f"Signal context: {signal_context}")
logger.info(f"Sequence encoding: {seq_encoding}, base_justify: {base_justify}")
if _feature_start is not None or _feature_end is not None:
_fs = _feature_start if _feature_start is not None else -kmer_context
_fe = _feature_end if _feature_end is not None else kmer_context
logger.info(f"Feature window: [{_fs}, {_fe}] relative to focus (width={_fe - _fs + 1})")
logger.info(f"Motif: {motif} (offset={motif_offset})")
logger.info(f"anchor: {anchor}")
if reference_fasta is not None:
logger.info(f"reference_fasta: {reference_fasta}")
# Open BAM for header and normalization detection
bam_in = pysam.AlignmentFile(str(bam_path), "rb")
# Detect normalization method: read from config, with sm/sd tag override for Remora
norm_method = config.get("signal_norm", "median_mad")
pa_mean = config.get("pa_mean")
pa_stdev = config.get("pa_stdev")
if is_remora:
# Peek at first alignment for pa_scaling tags
for first_aln in bam_in.fetch(until_eof=True):
if first_aln.query_name is not None:
if first_aln.has_tag("sm") and first_aln.has_tag("sd"):
pa_mean = float(first_aln.get_tag("sm"))
pa_stdev = float(first_aln.get_tag("sd"))
norm_method = "pa_scaling"
logger.info(f"Using pa_scaling normalization (sm={pa_mean}, sd={pa_stdev})")
break
if abstain_on_junction_indel and not is_multiclass:
raise RuntimeError(
"--abstain-on-junction-indel is only supported for multiclass models: it "
"gates the multiclass 'aa'/predicted_aa call, which a binary/pairwise model's "
"MP/ML BAM tags have no equivalent of. Drop the flag for a binary model."
)
# Create output writer (BAM or TSV)
output_path.parent.mkdir(parents=True, exist_ok=True)
tsv_writer = None
bam_out = None
if output_format == "tsv":
if not is_multiclass:
raise RuntimeError(
"TSV output is only supported for multiclass models. "
"Use .bam extension for binary models."
)
from leech.io.tsv_writer import TsvPredictionWriter
_has_cl = (
config.get("cl_regression", False)
and config.get("cl_regression_head_state_dict") is not None
)
if int_to_label:
_tsv_class_names = [int_to_label[i] for i in range(num_out)]
else:
_tsv_class_names = [str(i) for i in range(num_out)]
tsv_writer = TsvPredictionWriter(
output_path,
_tsv_class_names,
_has_cl,
copy_tags=copy_tags,
min_margin=min_margin,
abstain_on_junction_indel=abstain_on_junction_indel,
)
logger.info(f"TSV output: {output_path} ({len(_tsv_class_names)} classes)")
else:
bam_out = pysam.AlignmentFile(str(output_path), "wb", template=bam_in)
bam_in.close()
total_reads = 0
total_predictions = 0
if hasattr(model_wrapper, "model"):
model_wrapper.model.eval()
if hasattr(model_wrapper, "eval"):
model_wrapper.eval()
# Set up CL regression head if present in config (multiclass bundles)
_cl_head: torch.nn.Module | None = None
if config.get("cl_regression") and isinstance(model_wrapper, ModelInferenceWrapper):
from leech.losses import RegressionHead
cl_state = config.get("cl_regression_head_state_dict")
if cl_state is not None:
repr_dim = model_wrapper.enable_repr_capture()
_cl_head = RegressionHead(input_dim=repr_dim)
_cl_head.load_state_dict(cl_state)
_cl_head.to(device)
_cl_head.eval()
logger.info(f"CL regression head loaded (repr_dim={repr_dim})")
# Enable TF32 matmul for better performance on Ampere+ GPUs.
if device.startswith("cuda"):
torch.set_float32_matmul_precision("high")
# torch.compile decision deferred until after BAM read count is known (see below)
# Skip feature computation when model doesn't need them (big speedup)
# But always compute when signal_in_channels > 1 (needed for kmer residual)
signal_in_channels = spec.signal_in_channels
compute_features = requires_features or signal_in_channels > 1
# Load reference sequences for reference-anchored mode and/or reference-based motif search
reference_sequences = None
if anchor == "reference" or motif is not None:
from leech.io import get_reference_sequences
reference_sequences = get_reference_sequences(bam_path, reference_fasta)
logger.info(f"Loaded {len(reference_sequences)} reference sequences")
# Create motif searcher (reference-based when reference_sequences available)
motif_searcher = get_motif_searcher(
mode="fasta" if reference_sequences else "bam",
reference_sequences=reference_sequences,
skip_indels=spec.skip_motif_indels,
anchor=anchor,
# Recorded by `data prepare` and carried through `model train`. Without
# it, a corpus prepared with --no-require-query-mapping was scored at
# predict time with the gate back on, i.e. on a different read
# population than the model was trained on.
require_query_mapping=spec.require_query_mapping,
)
# Prepare class_names_str for multiclass (shared across mega-batches)
class_names_str = None
if is_multiclass:
if int_to_label:
class_names = [int_to_label[i] for i in range(num_out)]
class_names_str = ",".join(class_names)
else:
class_names_str = ",".join(str(i) for i in range(num_out))
logger.info(f"Streaming inference with read_batch_size={read_batch_size}")
n_total_reads = count_bam_reads(bam_path)
n_total_mega_batches = math.ceil(n_total_reads / read_batch_size) if n_total_reads > 0 else 0
logger.info(
f"BAM contains ~{n_total_reads} mapped reads -> ~{n_total_mega_batches} mega-batches "
f"of {read_batch_size}"
)
mega_batch_idx = 0
# torch.compile the model for faster inference (CUDA graph + kernel fusion).
# Auto-skip for small runs (< TORCH_COMPILE_MIN_SAMPLES reads) where
# compilation overhead (~15-30s) outweighs the speedup. Also skip when
# repr capture hooks are active or when --no-compile is set. `eval test`
# (evaluation.py) reuses this same threshold and flag (#264).
_has_repr_hook = (
isinstance(model_wrapper, ModelInferenceWrapper) and model_wrapper._repr_hook is not None
)
if no_compile:
logger.info("torch.compile disabled (--no-compile)")
elif n_total_reads < TORCH_COMPILE_MIN_SAMPLES:
logger.info(
f"torch.compile auto-skipped ({n_total_reads} reads < "
f"{TORCH_COMPILE_MIN_SAMPLES} threshold)"
)
elif (
isinstance(model_wrapper, ModelInferenceWrapper)
and device.startswith("cuda")
and hasattr(torch, "compile")
and not _has_repr_hook
):
try:
# forward_batch() dispatches through forward_module, not model --
# ModelInferenceWrapper keeps them separate so DDP can wrap one
# without the other (training.py's _wrap_ddp sets forward_module
# the same way). Compiling model alone left the compiled graph
# unreferenced: every real forward pass stayed eager despite the
# "torch.compile enabled" log line below (#264).
compiled_model = torch.compile(model_wrapper.model, mode="reduce-overhead")
model_wrapper.model = compiled_model # ty: ignore[invalid-assignment]
model_wrapper.forward_module = compiled_model # ty: ignore[invalid-assignment]
logger.info("torch.compile enabled (mode=reduce-overhead)")
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) ----
from leech.io.bam_reader import ReadInfo
logger.info(f"Parallel inference with {num_workers} workers")
inf_config = InferenceConfig(
pod5_path=pod5_path,
signal=SignalConfig(
reverse_signal=reverse_signal,
anchor=anchor,
norm_method=norm_method,
pa_mean=pa_mean,
pa_stdev=pa_stdev,
refine_signal_map=refine_signal_map,
signal_refiner=signal_refiner,
feature_channels=spec.feature_channels,
),
motif=MotifConfig(
motif=motif,
motif_offset=motif_offset,
reference_sequences=reference_sequences,
skip_motif_indels=spec.skip_motif_indels,
require_query_mapping=spec.require_query_mapping,
),
chunk=ChunkConfig(
base_justify=base_justify,
feature_start=_feature_start,
feature_end=_feature_end,
signal_context=signal_context,
kmer_context=kmer_context,
recover_softclip_signal=spec.recover_softclip_signal,
signal_context_bases=spec.signal_context_bases,
signal_len=signal_len,
mask_seq_side=spec.mask_seq_side,
),
seq_encoding=seq_encoding,
signal_kmer_context=signal_kmer_context,
signal_len=signal_len,
kmer_len=kmer_len,
dwell_offset=dwell_offset,
wide_features=wide_features,
requires_features=requires_features,
signal_in_channels=signal_in_channels,
)
calibration = spec.calibration
_batch_fn_p = (
functools.partial(
_run_batch_multiclass,
calibration=calibration,
cl_regression_head=_cl_head,
)
if is_multiclass
else _run_batch
)
def _run_worker_batch(sigs, seqs, feats, meta) -> None:
"""Flush callback: score one batch into the current mega-batch's ``pending``."""
_batch_fn_p(
sigs,
seqs,
feats,
meta,
model_wrapper,
requires_features,
device,
pending,
)
with Progress() as progress:
task = progress.add_task("[cyan]Running inference...", total=None)
with mp.Pool(processes=num_workers) as pool:
for aln_batch in iter_bam_batches(
bam_path, batch_size=read_batch_size, min_mapq=min_mapq
):
logger.info(f"Mega-batch: {len(aln_batch)} alignments read from BAM")
# Build ReadInfo objects from this mega-batch
read_infos = []
for aln in aln_batch:
try:
read_infos.append(ReadInfo(aln))
except Exception as e:
logger.warning(f"Skipping read {aln.query_name}: {e}")
logger.info(
f"Built {len(read_infos)} ReadInfo objects, "
f"dispatching to {num_workers} workers"
)
if not read_infos:
if bam_out is not None:
for aln in aln_batch:
bam_out.write(aln)
total_reads += len(aln_batch)
continue
# Split into worker sub-batches and dispatch
worker_batches = [
read_infos[i : i + chunk_size]
for i in range(0, len(read_infos), chunk_size)
]
worker_args = [(wb, inf_config, dwell_template) for wb in worker_batches]
pending: dict[str, list] = {}
accumulator = BatchAccumulator(batch_size, _run_worker_batch)
for worker_results in pool.imap_unordered(_inference_worker, worker_args):
for read_id, base_idx, sig, enc_seq, feat, ji, jm in worker_results:
# The template, if any, is already applied inside
# the worker -- before it narrowed the feature
# window to the model's k-mer width, the same
# order training used (issue #376).
accumulator.add(sig, enc_seq, feat, (read_id, base_idx, ji, jm))
# One worker's results never share a batch with the
# next worker's -- imap_unordered hands them back
# whole, and this is the batching the path has always
# had.
accumulator.flush()
# Write this mega-batch's predictions
if tsv_writer is not None:
batch_preds = tsv_writer.write_predictions(aln_batch, pending, int_to_label)
else:
batch_preds = _write_mega_batch_predictions(
aln_batch,
pending,
bam_out,
is_multiclass,
int_to_label,
class_names_str,
raw,
min_confidence,
min_margin,
abstain_on_junction_indel,
)
total_reads += len(aln_batch)
total_predictions += batch_preds
mega_batch_idx += 1
logger.info(
f"Mega-batch {mega_batch_idx}/{n_total_mega_batches} complete: "
f"wrote {batch_preds} predictions for {len(aln_batch)} reads"
)
progress.update(
task,
advance=0,
description=(
f"[cyan]Processed {total_reads} reads "
f"({total_predictions} predictions)..."
),
)
else:
# ---- Sequential path (mega-batched, double-buffered GPU) ----
from concurrent.futures import Future, ThreadPoolExecutor
pending: dict[str, list] = {}
_shape_validated = False
calibration = spec.calibration
_batch_fn = (
functools.partial(
_run_batch_multiclass,
calibration=calibration,
cl_regression_head=_cl_head,
)
if is_multiclass
else _run_batch
)
def _score_batch(sigs, seqs, feats, meta) -> None:
"""Score one batch into the *current* ``pending`` dict.
A plain closure, not ``functools.partial(..., pending=pending)``:
``_finalize_mega_batch`` below rebinds the name ``pending`` to a
fresh dict every mega-batch (so the async BAM-write thread can
keep draining the old one without racing the next mega-batch's
writes). A partial captures the dict *object* bound at
construction and keeps writing into it forever; a closure looks
``pending`` up again on every call, which is what a rebind needs.
"""
_batch_fn(sigs, seqs, feats, meta, model_wrapper, requires_features, device, pending)
# Double-buffered, single-worker GPU submission (PR #253): the
# accumulator has already detached each batch's buffers when it
# flushes, so the GPU thread owns them and extraction keeps filling
# the next batch while this one scores.
gpu_runner = GpuBatchRunner(_score_batch)
accumulator = BatchAccumulator(batch_size, gpu_runner.submit)
_bam_write_executor = ThreadPoolExecutor(max_workers=1)
_bam_write_future: Future | None = None
seq_signal_config = SignalConfig(
reverse_signal=reverse_signal,
anchor=anchor,
norm_method=norm_method,
pa_mean=pa_mean,
pa_stdev=pa_stdev,
refine_signal_map=refine_signal_map,
signal_refiner=signal_refiner,
compute_features=compute_features,
feature_channels=spec.feature_channels,
)
seq_chunk_config = ChunkConfig(
base_justify=base_justify,
feature_start=_feature_start,
feature_end=_feature_end,
signal_context=signal_context,
kmer_context=kmer_context,
recover_softclip_signal=spec.recover_softclip_signal,
signal_context_bases=spec.signal_context_bases,
signal_len=signal_len,
mask_seq_side=spec.mask_seq_side,
)
# Extraction thread count + rust setup (all three shared with
# run_bundle_inference via helpers.py).
_MAX_THREADS = 8 # sensible cap to avoid oversubscription on shared nodes
_avail_cpus = cap_rayon_threads_for_slurm(max_cap=_MAX_THREADS)
n_extract = max(1, _avail_cpus - 6) # reserve headroom for main + GPU I/O
logger.info(f"Sequential path with {n_extract} extraction threads, double-buffered GPU")
(
_use_rust_extraction,
_rs_extract_inference_chunks,
_rs_preload_pod5_signals,
_rs_extract_chunks_from_preloaded,
) = check_rust_extraction_available(
backend,
norm_method,
seq_chunk_config.recover_softclip_signal,
signal_context_bases=spec.signal_context_bases,
mask_seq_side=seq_chunk_config.mask_seq_side,
)
if _use_rust_extraction:
logger.info("Using Rust monolithic extraction (escapepod-rs + leech_core)")
else:
logger.info(
"Using Python extraction path"
+ (" (forced via --backend python)" if backend == "python" else "")
)
# Check for prefetch support (split POD5 preload + chunk extraction)
_has_prefetch = (
_use_rust_extraction
and _rs_preload_pod5_signals is not None
and _rs_extract_chunks_from_preloaded is not None
)
if _has_prefetch:
logger.info("POD5 prefetch pipeline enabled (overlapped I/O)")
def _extract_one_read(
aln: pysam.AlignedSegment,
) -> list[tuple[np.ndarray, np.ndarray, np.ndarray | None, tuple[str, int, int, bool]]]:
"""Extract ready-to-batch chunks from one alignment. Thread-safe."""
read_id = aln.query_name
read_seq = aln.query_sequence
if read_id is None or read_seq is None:
return []
try:
move_table = extract_move_table(aln)
raw_signal, pod5_metadata = pod5_reader.get_signal(read_id)
ref_seq = None
cigar_tuples = None
if seq_signal_config.anchor == "reference":
if reference_sequences and aln.reference_name in reference_sequences:
full_ref = reference_sequences[aln.reference_name]
ref_seq = full_ref[aln.reference_start : aln.reference_end]
else:
try:
ref_seq = aln.get_reference_sequence()
except Exception:
ref_seq = None
cigar_tuples = aln.cigartuples
leech_read = build_leech_read(
read_id=read_id,
sequence=read_seq,
raw_signal=raw_signal,
move_table=move_table,
signal_config=seq_signal_config,
metadata={},
reference_sequence=ref_seq,
cigar_tuples=cigar_tuples,
cal_offset=pod5_metadata.get("calibration_offset"),
cal_scale=pod5_metadata.get("calibration_scale"),
)
except Exception as e:
logger.warning(f"Skipping read {read_id}: {e}")
return []
# Find positions to predict
assert motif is not None
matches = motif_searcher.find_motif_positions(
leech_read.read_id, leech_read.sequence, aln, motif
)
positions = [m.position + motif_offset for m in matches]
# Junction disruption at each position's motif span (issue #282),
# looked up by base_idx after chunks_for_read (which may drop
# some positions) rather than threaded through it.
junction_by_base = {
m.position + motif_offset: (m.junction_indel, m.junction_mapped) for m in matches
}
return [
(
sig,
seq_arr,
feat,
(read_id, base_idx, *junction_by_base.get(base_idx, (0, False))),
)
for sig, seq_arr, feat, base_idx in chunks_for_read(
leech_read,
positions,
chunk_config=seq_chunk_config,
signal_len=signal_len,
seq_encoding=seq_encoding,
signal_kmer_context=signal_kmer_context,
requires_features=requires_features,
kmer_len=kmer_len,
dwell_offset=dwell_offset,
wide_features=wide_features,
dwell_templates=dwell_templates_arr,
template_min_pos=dwell_template_min_pos,
dwell_row=dwell_row,
)
]
_extract_pool = ThreadPoolExecutor(max_workers=n_extract)
# Shared Rust kwargs + metadata collection (shared with
# run_bundle_inference via helpers.py).
assert motif is not None # type narrowing for helpers below
_rs_kwargs = build_rust_extraction_kwargs(
signal_context=signal_context,
kmer_context=kmer_context,
signal_len=signal_len,
compute_features=compute_features,
reverse_signal=reverse_signal,
feature_start=_feature_start,
feature_end=_feature_end,
anchor=anchor,
seq_encoding=seq_encoding,
signal_kmer_context=signal_kmer_context,
refine_signal_map=refine_signal_map,
signal_refiner=signal_refiner,
refine_half_bandwidth=spec.refine_half_bandwidth,
refine_scale_iters=spec.refine_scale_iters,
signal_in_channels=signal_in_channels,
base_justify=base_justify,
signal_context_bases=spec.signal_context_bases,
feature_channels=spec.feature_channels,
)
def _collect_bam_metadata(aln_batch: list) -> tuple:
return collect_bam_metadata_for_rust(
aln_batch,
motif=motif,
motif_offset=motif_offset,
motif_searcher=motif_searcher,
anchor=anchor,
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.
"""
(
rs_rids,
rs_seqs,
rs_strides,
rs_mvs,
rs_ns,
rs_trims,
rs_motifs,
rs_cigars,
rs_refs,
_rs_junctions,
) = 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(
preloaded,
read_ids=rs_rids[start:end],
sequences=rs_seqs[start:end],
mv_strides=rs_strides[start:end],
mv_arrays=rs_mvs[start:end],
num_samples_list=rs_ns[start:end],
trim_offsets=rs_trims[start:end],
motif_positions=rs_motifs[start:end],
cigar_tuples=rs_cigars[start:end] if anchor == "reference" else None,
reference_sequences=rs_refs[start:end] if anchor == "reference" else None,
**_rs_kwargs,
)
yield from sub_chunks
def _consume_rust_chunks(chunks, junction_lookup: dict[tuple[str, int], tuple[int, bool]]):
"""Iterate Rust chunks into batch buffers, flushing to GPU as needed.
``junction_lookup`` (from :func:`junction_lookup_from_rs_meta`) is
built by the caller from the same metadata that produced
``chunks``' ``motif_positions`` -- Rust's own chunk tuples carry
no junction information (issue #282).
"""
nonlocal _shape_validated
for sig, seq_arr, feat, read_id, base_idx in chunks:
junction_indel, junction_mapped = junction_lookup.get(
(read_id, base_idx), (0, False)
)
if signal_in_channels > 1 and sig.ndim == 1:
sig = sig.reshape(signal_in_channels, -1)
# Rust returns features at the full requested window width, the
# same as a Python chunk's `features`. Both need the same
# narrowing to the model's k-mer window -- this loop used to
# skip it, so a wide-window corpus fed the model the wrong
# width and `dwell_offset` was inert on the Rust path.
feat = prepare_inference_features(
feat,
kmer_len=kmer_len,
feature_start=_feature_start,
dwell_offset=dwell_offset,
wide_features=wide_features,
dwell_templates=dwell_templates_arr,
template_min_pos=dwell_template_min_pos,
dwell_row=dwell_row,
)
if not _shape_validated:
validate_inference_shapes(sig, feat, config)
_shape_validated = True
accumulator.add(
sig, seq_arr, feat, (read_id, base_idx, junction_indel, junction_mapped)
)
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
def _finalize_mega_batch(aln_batch_to_write):
"""Flush GPU, submit async BAM write, update counters.
BAM writes are overlapped with the next mega-batch's extraction.
We serialize writes (wait for previous) since pysam is not thread-safe.
"""
nonlocal total_reads, total_predictions, mega_batch_idx
nonlocal pending, _bam_write_future
accumulator.flush()
gpu_runner.drain()
# 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,
abstain_on_junction_indel,
)
total_reads += len(aln_batch_to_write)
total_predictions += batch_preds
mega_batch_idx += 1
logger.info(
f"Mega-batch {mega_batch_idx}/{n_total_mega_batches} complete: "
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:
task = progress.add_task("[cyan]Running inference...", total=None)
# Skip opening Python DatasetReader when Rust handles all POD5 I/O.
# Opening a 40+ GB POD5 file in Python just to index it is expensive
# and completely unused on the Rust extraction paths.
_pod5_ctx = nullcontext() if _use_rust_extraction else POD5Reader(pod5_path)
with _pod5_ctx as pod5_reader:
if _has_prefetch:
# ---- Queue-based extraction pipeline ----
# A producer thread handles BAM reading, metadata collection,
# POD5 prefetch, and Rust extraction. It pushes (aln_batch, chunks)
# to a bounded queue. The main thread consumes from the queue,
# runs GPU inference, and writes results.
#
# Benefits over the previous prefetch pipeline:
# - Consumer (GPU + finalize) runs concurrently with producer
# - Metadata for batch N+1 overlaps with extraction of batch N
# - No synchronization gap between mega-batches
assert _rs_preload_pod5_signals is not None
assert _rs_extract_chunks_from_preloaded is not None
_SENTINEL = object()
_extraction_queue: _queue.Queue = _queue.Queue(maxsize=2)
_producer_error: BaseException | None = None
def _extraction_producer():
"""Background thread: reads BAM -> metadata -> prefetch -> extract -> queue."""
nonlocal _producer_error
try:
_meta_exec = ThreadPoolExecutor(max_workers=1)
_prefetch_exec = ThreadPoolExecutor(max_workers=1)
# Pipeline state for overlapping metadata and prefetch
_prev = None # (prefetch_future, rs_meta, aln_batch) or None
_meta_future = None
for aln_batch in iter_bam_batches(
bam_path, batch_size=read_batch_size, min_mapq=min_mapq
):
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)
_meta_future = _meta_exec.submit(
_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).
# The junction lookup travels with the chunks: the
# consumer thread never sees p_meta itself.
_extraction_queue.put(
(
p_aln,
chunk_list,
len(p_rids),
junction_lookup_from_rs_meta(p_meta),
)
)
# Get metadata result (should be done by now)
rs_meta = _meta_future.result()
else:
# First batch -- no previous, collect metadata sync
rs_meta = _collect_bam_metadata(aln_batch)
# Submit prefetch for current batch (overlaps with
# consumer processing + next BAM read)
cur_future = _prefetch_exec.submit(
_rs_preload_pod5_signals,
str(pod5_path),
rs_meta[0],
)
_prev = (cur_future, rs_meta, aln_batch)
# 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),
junction_lookup_from_rs_meta(p_meta),
)
)
_meta_exec.shutdown(wait=True)
_prefetch_exec.shutdown(wait=True)
except BaseException as exc:
_producer_error = exc
finally:
_extraction_queue.put(_SENTINEL)
_producer_thread = _threading.Thread(target=_extraction_producer, daemon=True)
_producer_thread.start()
logger.info("Queue-based extraction pipeline started (producer thread)")
# Consumer loop: pull from queue -> GPU -> finalize
while True:
item = _extraction_queue.get()
if item is _SENTINEL:
break
aln_batch, chunk_list, n_rids, chunk_junction_lookup = item
_t_mb_start = _time.perf_counter()
logger.info(
f"Mega-batch: {len(aln_batch)} alignments, {n_rids} for Rust extraction"
)
if chunk_list:
_consume_rust_chunks(iter(chunk_list), chunk_junction_lookup)
_t_consume = _time.perf_counter()
_finalize_mega_batch(aln_batch)
_t_finalize = _time.perf_counter()
logger.debug(
f" Timing: consume+gpu={_t_consume - _t_mb_start:.2f}s "
f"finalize={_t_finalize - _t_consume:.2f}s"
)
progress.update(
task,
advance=0,
description=(
f"[cyan]Processed {total_reads} reads "
f"({total_predictions} predictions)..."
),
)
_producer_thread.join()
if _producer_error is not None:
raise RuntimeError("Extraction producer thread failed") from _producer_error
else:
# ---- Non-prefetch fallback paths ----
for aln_batch in iter_bam_batches(
bam_path, batch_size=read_batch_size, min_mapq=min_mapq
):
if _use_rust_extraction:
# ---- Rust monolithic hot path (no prefetch) ----
_t_mb_start = _time.perf_counter()
rs_meta = _collect_bam_metadata(aln_batch)
_t_meta = _time.perf_counter()
rs_read_ids = rs_meta[0]
logger.info(
f"Mega-batch: {len(aln_batch)} alignments, "
f"{len(rs_read_ids)} for Rust extraction"
)
if rs_read_ids:
assert _rs_extract_inference_chunks is not None
chunks = _rs_extract_inference_chunks(
str(pod5_path),
read_ids=rs_read_ids,
sequences=rs_meta[1],
mv_strides=rs_meta[2],
mv_arrays=rs_meta[3],
num_samples_list=rs_meta[4],
trim_offsets=rs_meta[5],
motif_positions=rs_meta[6],
cigar_tuples=(rs_meta[7] if anchor == "reference" else None),
reference_sequences=(
rs_meta[8] if anchor == "reference" else None
),
**_rs_kwargs,
)
_consume_rust_chunks(chunks, junction_lookup_from_rs_meta(rs_meta))
_t_extract = _time.perf_counter()
logger.debug(
f" Timing: metadata={_t_meta - _t_mb_start:.2f}s "
f"extract+gpu={_t_extract - _t_meta:.2f}s"
)
else:
# ---- Python extraction path ----
batch_read_ids = [
aln.query_name
for aln in aln_batch
if aln.query_name is not None and aln.query_sequence is not None
]
if batch_read_ids:
pod5_reader.preload(batch_read_ids)
logger.info(
f"Mega-batch: {len(aln_batch)} alignments, "
f"{len(batch_read_ids)} preloaded from POD5"
)
# Parallel extraction -> GPU batching
for chunks in _extract_pool.map(_extract_one_read, aln_batch):
for sig, seq_arr, feat, meta in chunks:
if not _shape_validated:
validate_inference_shapes(sig, feat, config)
_shape_validated = True
accumulator.add(sig, seq_arr, feat, meta)
_finalize_mega_batch(aln_batch)
progress.update(
task,
advance=0,
description=(
f"[cyan]Processed {total_reads} reads "
f"({total_predictions} predictions)..."
),
)
_t_total = _time.perf_counter() - _t_total_start
logger.info(
f"Inference wall time: {_t_total:.1f}s "
f"({total_reads} reads, {total_predictions} predictions, "
f"{total_reads / _t_total:.0f} reads/s)"
)
# wait=True, not False. Every one of these pools is already drained here
# (`gpu_runner.drain()` 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.
_extract_pool.shutdown(wait=True)
gpu_runner.shutdown()
_wait_for_bam_write() # Ensure final BAM write completes before close
_bam_write_executor.shutdown(wait=True)
if tsv_writer is not None:
tsv_writer.close()
if bam_out is not None:
bam_out.close()
logger.info("Inference complete!")
logger.info(f"Reads processed: {total_reads}")
logger.info(f"Total predictions: {total_predictions}")
logger.info(f"Output written to: {output_path}")