diff --git a/changelog.d/268.fixed.md b/changelog.d/268.fixed.md new file mode 100644 index 00000000..5fef86d6 --- /dev/null +++ b/changelog.d/268.fixed.md @@ -0,0 +1,22 @@ +**Bundle predict shares the single-model pipeline's chunk extraction and GPU +submission instead of carrying its own copies.** `chunks_for_read` is now the +one place the per-chunk array-building triplet (signal padding/cropping, +sequence encoding, feature-window narrowing) runs for the Python serial +extraction path, used by both `run_inference` and `run_bundle_inference` — +before this, only the bundle copy applied dwell templates, and only the +bundle copy skipped `prepare_signal_channels`'s `signal_len` padding/cropping +on the raw signal itself (a real, if narrow, correctness gap found while +unifying, not by design). `GpuBatchRunner` is the shared double-buffered, +single-worker async GPU submission PR #253 introduced for predict, now also +used by every one of `run_bundle_inference`'s three extraction paths (Rust, +`num_workers>0`, serial) via a new `BundleScorer`. Bundle predict on the +serial Python path picks up the same overlap-extraction-with-scoring +speedup single-model predict already had. + +Fixed along the way: `run_inference`'s sequential-path GPU submission bound +its scoring closure's `pending` dict by value at construction, so once +`_finalize_mega_batch` rebound the name to a fresh dict for the next +mega-batch (the mechanism that lets the async BAM-write thread keep draining +the old one safely), every later mega-batch's predictions were computed and +silently dropped instead of written — caught by a new parity test, not by +inspection. diff --git a/changelog.d/269.fixed.md b/changelog.d/269.fixed.md new file mode 100644 index 00000000..0a6fa35e --- /dev/null +++ b/changelog.d/269.fixed.md @@ -0,0 +1,18 @@ +**One `InferenceSpec` resolver replaces four copies of predict's config-resolution +logic.** `run_inference` and `run_bundle_inference` each carried their own +~90-line block resolving motif/anchor/base_justify, signal-map refinement, +and the feature-window fallback chain; the two had drifted into real +disagreements. `is_multiclass` was `num_out > 1` in predict but `> 2` in +`eval test` and `model calibrate`, so a 2-output cross-entropy model got +multiclass BAM tags from predict but binary Platt calibration from +calibrate. `seq_encoding`'s missing-key fallback was `"signal_kmer"` in +predict/eval and `"base_onehot"` in calibration/export — `"base_onehot"` is +correct (it's the model classes' own default, not the CLI's), so every site +now agrees on it. The Rust-path `refine_scale_iters` default was hardcoded +`2` even for Remora models, which use `-1`. `_check_config_consistency` now +takes click's real parameter source, so an explicit `--motif-offset 0` that +disagrees with the checkpoint raises instead of silently losing to it. The +feature-window resolver (`chunking.feature_window_from_metadata`) reads a +whole corpus's column and asserts it is constant, rather than trusting chunk +0 — training's own `_raw_chunk = train_dataset.chunks[0]` read was exactly +that failure mode, generalized from issue #230. diff --git a/src/leech/calibration.py b/src/leech/calibration.py index 753a020d..d1faae97 100644 --- a/src/leech/calibration.py +++ b/src/leech/calibration.py @@ -23,6 +23,7 @@ from torch.optim import LBFGS from torch.utils.data import DataLoader +from leech.constants import DEFAULT_SEQ_ENCODING_FALLBACK from leech.dataset import LeechDataset, collate_fn, resolve_val_dataloader_workers from leech.models import get_model from leech.models.inference_wrapper import ModelInferenceWrapper @@ -141,7 +142,7 @@ def calibrate_model( model_name = config["model_name"] signal_len = config["signal_len"] kmer_len = config["kmer_len"] - seq_encoding = config.get("seq_encoding", "base_onehot") + seq_encoding = config.get("seq_encoding", DEFAULT_SEQ_ENCODING_FALLBACK) signal_kmer_context = tuple(config.get("signal_kmer_context", (4, 4))) # Build model init kwargs @@ -498,7 +499,7 @@ def calibrate_model_multiclass( signal_len = config["signal_len"] kmer_len = config["kmer_len"] num_out = config.get("num_out", 1) - seq_encoding = config.get("seq_encoding", "base_onehot") + seq_encoding = config.get("seq_encoding", DEFAULT_SEQ_ENCODING_FALLBACK) signal_kmer_context = tuple(config.get("signal_kmer_context", (4, 4))) # Build model init kwargs diff --git a/src/leech/chunking/__init__.py b/src/leech/chunking/__init__.py index 0610042a..284f3710 100644 --- a/src/leech/chunking/__init__.py +++ b/src/leech/chunking/__init__.py @@ -9,6 +9,7 @@ LeechRead, extract_training_chunks, extraction_sequence, + feature_window_from_metadata, find_focus_bases, resolve_feature_window, ) @@ -33,6 +34,7 @@ "LeechRead", "extract_training_chunks", "extraction_sequence", + "feature_window_from_metadata", "find_focus_bases", "resolve_feature_window", # Serialization diff --git a/src/leech/chunking/extractor.py b/src/leech/chunking/extractor.py index e2ba113b..772092e4 100644 --- a/src/leech/chunking/extractor.py +++ b/src/leech/chunking/extractor.py @@ -51,6 +51,119 @@ def resolve_feature_window( return start, end, end - start + 1 +def _assert_constant_column(column: np.ndarray, name: str) -> None: + """Refuse a corpus whose stored feature window is not the same on every row. + + Reading column 0 (or any single row) and trusting it for the whole corpus + is issue #230's failure mode generalized to this field: a corpus written + by two prepare runs with different ``--feature-start``/``--feature-end`` + would otherwise train silently on whichever value the first chunk happened + to carry. + """ + if column.size and not np.all(column == column[0]): + raise ValueError( + f"Corpus has inconsistent {name!r} across chunks: {np.unique(column).tolist()}. " + f"A training run needs one feature window throughout; this corpus mixes prepare " + f"runs with different --feature-start/--feature-end and must be re-prepared " + f"consistently before it can be trained on." + ) + + +def _feature_window_from_table(table, kmer_context: int) -> tuple[int | None, int | None]: + """Column form of :func:`feature_window_from_metadata`. + + ``feature_left``/``feature_right`` are never written as columns (only as + legacy per-chunk-dict fields on the row path), so the fallback chain here + is shorter than the mapping form's. + """ + start_col = table.values("feature_start") + if start_col is not None: + _assert_constant_column(start_col, "feature_start") + start = int(start_col[0]) + else: + margin_col = table.values("dwell_margin_left") + if margin_col is not None: + _assert_constant_column(margin_col, "dwell_margin_left") + start = -(kmer_context + int(margin_col[0])) + else: + start = None + + end_col = table.values("feature_end") + if end_col is not None: + _assert_constant_column(end_col, "feature_end") + end = int(end_col[0]) + else: + end = None + + return start, end + + +def _feature_window_from_mapping(source, kmer_context: int) -> tuple[int | None, int | None]: + """Mapping form of :func:`feature_window_from_metadata`. + + ``source`` is a model config dict or a single chunk (a plain dict or a + :class:`~leech.chunking.table.ChunkRow`). Mirrors the fallback chain that + used to be pasted into ``single.py``, ``bundle.py``, ``training.py`` and + ``dataset.py`` separately: ``feature_start``/``feature_end`` (current), + then ``feature_left``/``feature_right`` (legacy), then + ``dwell_margin_left``/``dwell_margin_right`` (older legacy, ``right`` + resolved from the stored feature array's width when one is present). + """ + if source.get("feature_start") is not None: + start = int(source["feature_start"]) + elif source.get("feature_left") is not None: + start = -int(source["feature_left"]) + elif source.get("dwell_margin_left") is not None: + start = -(kmer_context + int(source["dwell_margin_left"])) + else: + start = None + + if source.get("feature_end") is not None: + end = int(source["feature_end"]) + elif source.get("feature_right") is not None: + end = int(source["feature_right"]) + elif source.get("dwell_margin_right") is not None: + raw_features = source.get("features") + if raw_features is not None and getattr(raw_features, "ndim", 1) > 1 and start is not None: + # A chunk dict: derive from the stored feature array's actual width. + end = raw_features.shape[1] - 1 + start + else: + # A config dict never carries "features" -- this is the only branch + # that fires for it, and single.py/bundle.py's original inline code + # computed exactly this (kmer_context + dwell_margin_right) with no + # array to check. Do not collapse this to None: that silently + # narrows the window to +-kmer_context instead of using the margin + # the config actually recorded. + end = kmer_context + int(source["dwell_margin_right"]) + else: + end = None + + return start, end + + +def feature_window_from_metadata(source, kmer_context: int) -> tuple[int | None, int | None]: + """Resolve the stored ``(feature_start, feature_end)`` from a config or corpus. + + ``None`` for either element means nothing was stored for it -- callers + fall back to the k-mer window via :func:`resolve_feature_window`, exactly + as a fresh corpus with no feature-window fields at all always has. + + ``source`` is either a :class:`~leech.chunking.table.ChunkTable` (the + window is read as a column and asserted constant across every chunk -- see + :func:`_assert_constant_column`) or a mapping -- a model's ``config.json`` + dict or a single chunk (a plain dict or a + :class:`~leech.chunking.table.ChunkRow`). This is the one place the + fallback chain is written; ``training.py``, ``dataset.py`` and + ``InferenceSpec`` all resolve through it rather than carrying their own + copies (issue #269). + """ + from leech.chunking.table import ChunkTable + + if isinstance(source, ChunkTable): + return _feature_window_from_table(source, kmer_context) + return _feature_window_from_mapping(source, kmer_context) + + def merge_feature_channels( dwell_features: dict[str, np.ndarray], signal_features: dict[str, np.ndarray], diff --git a/src/leech/cli.py b/src/leech/cli.py index fda3c397..0ffb34ed 100644 --- a/src/leech/cli.py +++ b/src/leech/cli.py @@ -1828,6 +1828,15 @@ def predict( from leech.commands.predict import handle_predict + # An explicit CLI value that happens to equal its own default (e.g. + # `--motif-offset 0`) must still be checked against the model config + # rather than silently losing to it -- see InferenceSpec.from_config. + _ctx = click.get_current_context() + parameter_sources = { + name: _ctx.get_parameter_source(name) is ParameterSource.COMMANDLINE + for name in ("motif", "motif_offset", "base_justify", "anchor") + } + if reference_anchored: warnings.warn( "--reference-anchored is deprecated, use --anchor reference instead", @@ -1835,6 +1844,11 @@ def predict( stacklevel=2, ) anchor = "reference" + # The deprecated flag is itself an explicit request for anchor, same + # as --anchor reference would be -- without this, a model config + # recording anchor="basecall" silently wins over it instead of + # raising the conflict InferenceSpec.from_config exists to catch. + parameter_sources["anchor"] = True # Parse --copy-tags into a list parsed_copy_tags = [t.strip() for t in copy_tags.split(",") if t.strip()] if copy_tags else None @@ -1846,8 +1860,6 @@ def predict( elif output_name.endswith(".bam"): output_format = "bam" else: - import rich_click as click - raise click.UsageError( f"Cannot determine output format from extension: {output}. " "Use .bam for BAM output or .tsv / .tsv.gz for TSV output." @@ -1880,6 +1892,7 @@ def predict( no_compile=no_compile, output_format=output_format, copy_tags=parsed_copy_tags, + parameter_sources=parameter_sources, ) diff --git a/src/leech/commands/calibrate.py b/src/leech/commands/calibrate.py index 08007bdb..069eb72d 100644 --- a/src/leech/commands/calibrate.py +++ b/src/leech/commands/calibrate.py @@ -48,12 +48,19 @@ def handle_calibrate( import json from leech.calibration import calibrate_model, calibrate_model_multiclass + from leech.inference.helpers import is_multiclass as _is_multiclass_config def _is_multiclass(config_path: Path) -> bool: - """Check if model config indicates multiclass (num_out > 2).""" + """Check if model config indicates multiclass output. + + One definition, shared with predict and `eval test`: num_out > 1 + (issue #269). This used to be `> 2` here, so a 2-output cross-entropy + model got binary Platt calibration even though predict tags and + expects it to be calibrated as multiclass. + """ with open(config_path) as f: cfg = json.load(f) - return cfg.get("num_out", 1) > 2 + return _is_multiclass_config(cfg) def _calibrate_single(mdir: Path, output_path: Path | None = None) -> str: """Calibrate a single model dir, auto-detecting binary vs multiclass.""" diff --git a/src/leech/commands/predict.py b/src/leech/commands/predict.py index b1abaab8..9fca6a12 100644 --- a/src/leech/commands/predict.py +++ b/src/leech/commands/predict.py @@ -61,6 +61,7 @@ def handle_predict( no_compile: bool = False, output_format: str = "bam", copy_tags: list[str] | None = None, + parameter_sources: dict[str, bool] | None = None, ) -> None: """ Handle the predict command logic. @@ -128,6 +129,7 @@ def handle_predict( no_compile=no_compile, output_format=output_format, copy_tags=copy_tags, + parameter_sources=parameter_sources, ) else: if output_format == "tsv": @@ -158,6 +160,7 @@ def handle_predict( num_workers=workers, read_batch_size=read_batch_size, backend=backend, + parameter_sources=parameter_sources, ) else: if output_format == "tsv": @@ -203,6 +206,7 @@ def handle_predict( read_batch_size=read_batch_size, backend=backend, no_compile=no_compile, + parameter_sources=parameter_sources, ) console.print("[bold green]Inference complete![/bold green]") diff --git a/src/leech/constants.py b/src/leech/constants.py index 8b7d4b4f..8e1ba673 100644 --- a/src/leech/constants.py +++ b/src/leech/constants.py @@ -96,6 +96,17 @@ def generate_random_seed() -> int: # Sequence encoding defaults DEFAULT_SIGNAL_KMER_CONTEXT = (4, 4) # Kmer context for signal-level kmer encoding +# What `seq_encoding` means for a checkpoint whose config lacks the key. +# Every consumer of an old/bare config must agree on this fallback (issue +# #269) -- `base_onehot`, because that is the model classes' OWN default +# (see e.g. `models/configs/conv_lstm.toml`'s `seq_encoding = "base_onehot"` +# [params] entry): a model built via `get_model(name, **kwargs)` with no +# seq_encoding kwarg is a base_onehot model, full stop. `--seq-encoding +# signal_kmer` is only the *CLI*'s default for `model train`, which always +# records the effective encoding it used into config.json -- so a config +# that lacks the key was never touched by that flag at all. +DEFAULT_SEQ_ENCODING_FALLBACK = "base_onehot" + # Signal map refinement defaults. # # The single source of truth for the refine-half-bandwidth default (5, not diff --git a/src/leech/dataset.py b/src/leech/dataset.py index 2d3521b2..46458945 100644 --- a/src/leech/dataset.py +++ b/src/leech/dataset.py @@ -48,6 +48,7 @@ from leech.chunking import ( ChunkTable, + feature_window_from_metadata, iter_npz_row_blocks, load_chunks, load_seq_to_sig_csr, @@ -1510,15 +1511,11 @@ def _prepare_features( if self._wide_features: pass # full-width features elif dwell_width > self.kmer_len: - # Determine feature_start (signed offset from focus). - # New chunks have it directly; old chunks need conversion. - if "feature_start" in chunk: - feat_start = int(chunk["feature_start"]) - elif "feature_left" in chunk: - feat_start = -int(chunk["feature_left"]) - elif "dwell_margin_left" in chunk: - feat_start = -(kmer_context + int(chunk["dwell_margin_left"])) - else: + # Determine feature_start (signed offset from focus). New chunks + # have it directly; old chunks need conversion -- the shared + # fallback chain (issue #269) covers both. + feat_start, _ = feature_window_from_metadata(chunk, kmer_context) + if feat_start is None: feat_start = -(dwell_width - 1) // 2 # symmetric fallback # kmer-aligned start within the feature array # Feature array starts at focus + feat_start, kmer starts at focus - kmer_context diff --git a/src/leech/evaluation.py b/src/leech/evaluation.py index 93cacb7a..1b9d43c7 100644 --- a/src/leech/evaluation.py +++ b/src/leech/evaluation.py @@ -20,8 +20,9 @@ ) from torch.utils.data import DataLoader -from leech.constants import TORCH_COMPILE_MIN_SAMPLES +from leech.constants import DEFAULT_SEQ_ENCODING_FALLBACK, TORCH_COMPILE_MIN_SAMPLES from leech.dataset import LeechDataset, collate_fn, resolve_dataloader_workers +from leech.inference.helpers import is_multiclass as _is_multiclass_config from leech.metrics import compute_metrics, print_metrics, save_metrics from leech.model_loading import load_model_from_checkpoint from leech.models.inference_wrapper import ModelInferenceWrapper @@ -125,7 +126,7 @@ def evaluate_model( num_out = config.get("num_out", 1) left_context = config.get("left_context") right_context = config.get("right_context") - seq_encoding = config.get("seq_encoding", "signal_kmer") + seq_encoding = config.get("seq_encoding", DEFAULT_SEQ_ENCODING_FALLBACK) dwell_offset = config.get("dwell_offset", 0) signal_kmer_context = tuple(config.get("signal_kmer_context", (4, 4))) signal_mode = config.get("signal_mode", "both") @@ -208,7 +209,10 @@ def evaluate_model( # Run evaluation logger.info("\nRunning evaluation...") - is_multiclass = num_out > 2 + # One definition, shared with predict: num_out > 1 (issue #269). This used + # to be `> 2` here, so a 2-output cross-entropy model was scored as binary + # by `eval test` while `predict` tagged and calibrated it as multiclass. + is_multiclass = _is_multiclass_config(config) # Same decision predict makes (never autocasts by default): opt in with # --mixed-precision rather than autocasting on CUDA unconditionally (#264). use_autocast = mixed_precision and device != "cpu" diff --git a/src/leech/inference/bundle.py b/src/leech/inference/bundle.py index 572cd869..5c3270a8 100644 --- a/src/leech/inference/bundle.py +++ b/src/leech/inference/bundle.py @@ -12,7 +12,6 @@ from rich.progress import Progress from leech.configs import ChunkConfig, InferenceConfig, MotifConfig, SignalConfig -from leech.constants import DEFAULT_REFINE_HALF_BANDWIDTH from leech.features import extract_move_table from leech.inference.aggregation import ( aggregate_one_vs_all, @@ -22,12 +21,13 @@ ) from leech.inference.helpers import ( BatchAccumulator, - _check_config_consistency, - _encode_sequence_for_inference, + GpuBatchRunner, + InferenceSpec, _write_prediction_tags, build_rust_extraction_kwargs, cap_rayon_threads_for_slurm, check_rust_extraction_available, + chunks_for_read, collect_bam_metadata_for_rust, prepare_inference_features, validate_inference_shapes, @@ -37,7 +37,6 @@ from leech.io.pod5_reader import POD5Reader from leech.model_export import deserialize_exported_model, deserialize_traced_model from leech.model_loading import _instantiate_model -from leech.models import wide_features as _wide_features from leech.models.inference_wrapper import ModelInferenceWrapper, TracedModelWrapper from leech.preparation.reader import build_leech_read @@ -87,6 +86,95 @@ def _write_bundle_mega_batch( return n_written +class BundleScorer: + """Runs every model in a bundle over one batch of chunks. + + Replaces the ``_run_bundle_batch`` closure that all three of + ``run_bundle_inference``'s extraction paths (Rust, parallel, serial) used + to share: same vmap-vs-sequential-wrappers branch, same Platt scaling. + Wrapping it as a scorer object rather than a closure is what lets it be + handed to :class:`~leech.inference.helpers.GpuBatchRunner` the same way + the single-model path's ``_run_batch`` is (issue #268). + """ + + def __init__( + self, + *, + pairs: list[str], + pair_to_idx: dict[str, int], + n_pairs: int, + needs_features: bool, + device: str, + is_vmap: bool, + wrappers: dict | None = None, + platt_params: dict[str, tuple[float, float]] | None = None, + vmapped_forward=None, + vmap_stacked_params: dict | None = None, + vmap_stacked_buffers: dict | None = None, + vmap_platt_a=None, + vmap_platt_b=None, + ): + self.pairs = pairs + self.pair_to_idx = pair_to_idx + self.n_pairs = n_pairs + self.needs_features = needs_features + self.device = device + self.is_vmap = is_vmap + self.wrappers = wrappers or {} + self.platt_params = platt_params or {} + self.vmapped_forward = vmapped_forward + self.vmap_stacked_params = vmap_stacked_params + self.vmap_stacked_buffers = vmap_stacked_buffers + self.vmap_platt_a = vmap_platt_a + self.vmap_platt_b = vmap_platt_b + self.n_batches_done = 0 + + def score(self, sigs: list, seqs: list, feats: list, rids: list, *, read_probs: dict) -> None: + """Score one batch, writing each read's ``(n_pairs,)`` probability + vector into ``read_probs``. Matches :class:`GpuBatchRunner`'s + ``score_fn(signals, sequences, features, meta)`` shape via + ``functools.partial(scorer.score, read_probs=read_probs)``.""" + sig_t = torch.stack(sigs).to(self.device) + seq_t = torch.stack(seqs).to(self.device) + feat_t = None + if self.needs_features: + valid_feats = [f for f in feats if f is not None] + if valid_feats: + feat_t = torch.stack(valid_feats).to(self.device) + + with torch.inference_mode(): + if self.is_vmap and self.vmapped_forward is not None: + if self.needs_features: + all_logits = self.vmapped_forward( + self.vmap_stacked_params, self.vmap_stacked_buffers, sig_t, seq_t, feat_t + ) + else: + all_logits = self.vmapped_forward( + self.vmap_stacked_params, self.vmap_stacked_buffers, sig_t, seq_t, None + ) + all_logits = ( + self.vmap_platt_a[:, None, None] * all_logits + self.vmap_platt_b[:, None, None] + ) + all_p = torch.sigmoid(all_logits).squeeze(-1).cpu().numpy() + else: + all_p = np.empty((self.n_pairs, len(sigs)), dtype=np.float32) + for pair in self.pairs: + batch_dict: dict[str, torch.Tensor] = {"signal": sig_t, "sequence": seq_t} + if feat_t is not None: + batch_dict["features"] = feat_t + logits = self.wrappers[pair].forward_batch(batch_dict, self.device) + pp = self.platt_params.get(pair) + if pp is not None: + a, b = pp + logits = a * logits + b + all_p[self.pair_to_idx[pair]] = torch.sigmoid(logits).cpu().numpy().flatten() + + for i, rid in enumerate(rids): + read_probs[rid] = all_p[:, i] + + self.n_batches_done += 1 + + def run_bundle_inference( bundle_path: Path, pod5_path: Path, @@ -108,6 +196,7 @@ def run_bundle_inference( num_workers: int = 0, read_batch_size: int = 10_000, backend: str = "auto", + parameter_sources: dict[str, bool] | None = None, ) -> None: """ Run all models from a bundle on each read, aggregate to a single AA prediction. @@ -146,101 +235,49 @@ def run_bundle_inference( comparison_type = metadata["comparison_type"] is_torchscript = metadata.get("torchscript", False) - signal_len = config["signal_len"] - kmer_len = config["kmer_len"] model_type = config.get("model_name", metadata.get("architecture", "")) - dwell_offset = config.get("dwell_offset", 0) - seq_encoding = config.get("seq_encoding", "signal_kmer") - signal_kmer_context = tuple(config.get("signal_kmer_context", (4, 4))) - - # Resolve motif/offset/justify from config, erroring on CLI conflict - motif = _check_config_consistency("motif", motif, config.get("motif"), None) - motif_offset = _check_config_consistency( - "motif-offset", motif_offset, config.get("motif_offset"), 0 - ) - if motif is not None: - logger.info(f"Motif from bundle config: {motif} (offset={motif_offset})") - - if motif is None: - raise ValueError( - "motif is None after auto-read from bundle config. " - "Either pass --motif on the CLI or ensure the bundle config contains a non-null 'motif' field. " - "Without a motif, inference predicts at every position, producing noise." - ) - - base_justify = _check_config_consistency( - "base-justify", base_justify, config.get("base_justify"), "center" + spec = InferenceSpec.from_config( + config, + model_type=model_type, + motif=motif, + motif_offset=motif_offset, + base_justify=base_justify, + anchor=anchor, + reference_fasta=reference_fasta, + default_refine_scale_iters=2, + parameter_sources=parameter_sources, + strict_model_type=True, + model_dwell_margin=lambda: getattr(_instantiate_model(config), "dwell_margin", 0), ) + signal_len = spec.signal_len + kmer_len = spec.kmer_len + dwell_offset = spec.dwell_offset + seq_encoding = spec.seq_encoding + signal_kmer_context = spec.signal_kmer_context + motif = spec.motif + motif_offset = spec.motif_offset + base_justify = spec.base_justify + anchor = spec.anchor + reference_fasta = spec.reference_fasta + signal_context = spec.signal_context + kmer_context = spec.kmer_context + wide_features = spec.wide_features + _feature_start = spec.feature_start + _feature_end = spec.feature_end + + logger.info(f"Motif from bundle config: {motif} (offset={motif_offset})") logger.info(f"base_justify: {base_justify}") - - anchor = _check_config_consistency("anchor", anchor, config.get("anchor"), "reference") logger.info(f"anchor: {anchor}") - - if reference_fasta is None: - cfg_ref = config.get("reference_fasta") - if cfg_ref is not None: - cfg_path = Path(cfg_ref) - if cfg_path.exists(): - reference_fasta = cfg_path - logger.info(f"reference_fasta from config: {reference_fasta}") - else: - logger.warning( - f"reference_fasta from config ({cfg_ref}) not found; " - f"pass --reference-fasta explicitly" - ) - - # Use asymmetric context if available, otherwise fall back to symmetric - left_ctx = config.get("left_context") - right_ctx = config.get("right_context") - if left_ctx is not None and right_ctx is not None: - signal_context = (left_ctx, right_ctx) - else: - signal_context = (signal_len // 2, signal_len // 2) - kmer_context = kmer_len // 2 - + if reference_fasta is not None: + logger.info(f"reference_fasta: {reference_fasta}") logger.info( f"Bundle: {metadata['architecture']}, {len(pairs)} models, v{metadata['bundle_version']}" f"{' (TorchScript)' if is_torchscript else ''}" ) - - # Determine feature_start/feature_end from config (must match training data) - # model_type is read back from a saved bundle's own metadata, so it should - # always be a real registry name; an unrecognized one means the bundle is - # stale or corrupted and must fail loudly rather than silently guess its - # feature-window convention. - try: - wide_features = bool(model_type) and _wide_features(model_type) - except KeyError as e: - raise KeyError( - f"Bundle architecture '{model_type}' is not a recognized model " - f"(renamed, removed, or a corrupted bundle). Cannot determine " - f"its feature-window convention." - ) from e - _kmer_context = kmer_len // 2 - _feature_start = config.get("feature_start") - _feature_end = config.get("feature_end") - if _feature_start is None and "feature_left" in config: - _feature_start = -config["feature_left"] - if _feature_end is None and "feature_right" in config: - _feature_end = config["feature_right"] - if _feature_start is None and "dwell_margin_left" in config: - _feature_start = -(_kmer_context + config["dwell_margin_left"]) - if _feature_end is None and "dwell_margin_right" in config: - _feature_end = _kmer_context + config["dwell_margin_right"] - if wide_features and _feature_start is None and _feature_end is None: - _model_margin = getattr(_instantiate_model(config), "dwell_margin", 0) - if _model_margin: - _feature_start = -(_kmer_context + _model_margin) - _feature_end = _kmer_context + _model_margin - logger.warning( - f"Bundle config missing feature_start/end, " - f"falling back to model default margin: {_model_margin}" - ) - logger.info(f"Signal context: {signal_context}, kmer_len: {kmer_len}") logger.info(f"seq_encoding: {seq_encoding}, base_justify: {base_justify}") - _fs = _feature_start if _feature_start is not None else -_kmer_context - _fe = _feature_end if _feature_end is not None else _kmer_context + _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"dwell_offset: {dwell_offset}, feature window: [{_fs}, {_fe}]") if _feature_start is not None or _feature_end is not None: logger.info(f"Feature window: [{_fs}, {_fe}] relative to focus (width={_fe - _fs + 1})") @@ -327,13 +364,13 @@ def run_bundle_inference( motif_searcher = get_motif_searcher( mode="fasta" if reference_sequences else "bam", reference_sequences=reference_sequences, - skip_indels=config.get("skip_motif_indels", False), + 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=config.get("require_query_mapping", True), + require_query_mapping=spec.require_query_mapping, ) # Open BAM for header only @@ -348,13 +385,13 @@ def run_bundle_inference( else: first_wrapper = next(iter(wrappers.values())) needs_features = first_wrapper.requires_features - bundle_signal_in_channels = config.get("signal_in_channels", 1) + bundle_signal_in_channels = spec.signal_in_channels bundle_compute_features = needs_features or bundle_signal_in_channels > 1 # Signal map refinement for bundle models (needed for kmer residual signal channel) - bundle_refine = False + bundle_refine = spec.refine_signal_map bundle_refiner = None - if config.get("refine_signal_map", True) or bundle_signal_in_channels > 1: + if bundle_refine: from leech.data import get_kmer_table from leech.inference.helpers import _warn_if_kmer_table_drifted from leech.signal_refine import SigMapRefiner @@ -363,12 +400,11 @@ def run_bundle_inference( _warn_if_kmer_table_drifted(config.get("kmer_table_sha256"), kmer_table_path) bundle_refiner = SigMapRefiner.from_table( kmer_table_path, - half_bandwidth=config.get("refine_half_bandwidth", DEFAULT_REFINE_HALF_BANDWIDTH), - do_rough_rescale=config.get("refine_do_rough_rescale", True), - scale_iters=config.get("refine_scale_iters", 2), - center_idx=config.get("refine_kmer_center_idx", -1), + half_bandwidth=spec.refine_half_bandwidth, + do_rough_rescale=spec.refine_do_rough_rescale, + scale_iters=spec.refine_scale_iters, + center_idx=spec.refine_kmer_center_idx, ) - bundle_refine = True logger.info( f"Signal map refinement enabled for bundle " f"(signal_in_channels={bundle_signal_in_channels})" @@ -381,7 +417,7 @@ def run_bundle_inference( # verbatim at inference time. dwell_templates_arr: np.ndarray | None = None dwell_template_min_pos: int = 0 - bundle_dwell_template_table = config.get("dwell_template_table") or None + bundle_dwell_template_table = spec.dwell_template_table if bundle_dwell_template_table: from leech.dataset import load_dwell_template_table @@ -431,54 +467,29 @@ def _single_forward(params, buffers, signal, sequence, features): read_probs: dict[str, np.ndarray] = {} # read_id -> shape (n_pairs,) _shape_validated = False n_chunks = 0 - n_batches_done = 0 - - def _run_bundle_batch(sigs: list, seqs: list, feats: list, rids: list) -> None: - """Flush callback: run every model in the bundle over one batch.""" - nonlocal n_batches_done - - sig_t = torch.stack(sigs).to(device) - seq_t = torch.stack(seqs).to(device) - feat_t = None - if needs_features: - valid_feats = [f for f in feats if f is not None] - if valid_feats: - feat_t = torch.stack(valid_feats).to(device) - - with torch.inference_mode(): - if is_vmap and vmapped_forward is not None: - if needs_features: - all_logits = vmapped_forward( - vmap_stacked_params, vmap_stacked_buffers, sig_t, seq_t, feat_t - ) - else: - all_logits = vmapped_forward( - vmap_stacked_params, vmap_stacked_buffers, sig_t, seq_t, None - ) - all_logits = vmap_platt_a[:, None, None] * all_logits + vmap_platt_b[:, None, None] - all_p = torch.sigmoid(all_logits).squeeze(-1).cpu().numpy() - else: - all_p = np.empty((n_pairs, len(sigs)), dtype=np.float32) - for pair in pairs: - batch_dict: dict[str, torch.Tensor] = { - "signal": sig_t, - "sequence": seq_t, - } - if feat_t is not None: - batch_dict["features"] = feat_t - logits = wrappers[pair].forward_batch(batch_dict, device) - pp = platt_params.get(pair) - if pp is not None: - a, b = pp - logits = a * logits + b - all_p[pair_to_idx[pair]] = torch.sigmoid(logits).cpu().numpy().flatten() - for i, rid in enumerate(rids): - read_probs[rid] = all_p[:, i] - - n_batches_done += 1 - - accumulator = BatchAccumulator(batch_size, _run_bundle_batch) + scorer = BundleScorer( + pairs=pairs, + pair_to_idx=pair_to_idx, + n_pairs=n_pairs, + needs_features=needs_features, + device=device, + is_vmap=is_vmap, + wrappers=wrappers, + platt_params=platt_params, + vmapped_forward=vmapped_forward, + vmap_stacked_params=vmap_stacked_params, + vmap_stacked_buffers=vmap_stacked_buffers, + vmap_platt_a=vmap_platt_a, + vmap_platt_b=vmap_platt_b, + ) + # Double-buffered, single-worker GPU submission (PR #253's pattern, + # extended here to bundle predict -- issue #268). All three extraction + # paths below (Rust, parallel, serial) share this one accumulator and + # runner; each must `gpu_runner.drain()` before reading `read_probs`, + # since a flush only guarantees the batch was *submitted*, not scored. + gpu_runner = GpuBatchRunner(partial(scorer.score, read_probs=read_probs)) + accumulator = BatchAccumulator(batch_size, gpu_runner.submit) n_reads = 0 n_predicted = 0 @@ -499,7 +510,7 @@ def _run_bundle_batch(sigs: list, seqs: list, feats: list, rids: list) -> None: feature_end=_feature_end, signal_context=signal_context, kmer_context=kmer_context, - recover_softclip_signal=config.get("recover_softclip_signal", False), + recover_softclip_signal=spec.recover_softclip_signal, ) logger.info(f"Streaming bundle inference with read_batch_size={read_batch_size}") @@ -551,10 +562,8 @@ def _run_bundle_batch(sigs: list, seqs: list, feats: list, rids: list) -> None: signal_kmer_context=signal_kmer_context, refine_signal_map=bundle_refine, signal_refiner=bundle_refiner, - refine_half_bandwidth=config.get( - "refine_half_bandwidth", DEFAULT_REFINE_HALF_BANDWIDTH - ), - refine_scale_iters=config.get("refine_scale_iters", 2), + refine_half_bandwidth=spec.refine_half_bandwidth, + refine_scale_iters=spec.refine_scale_iters, signal_in_channels=bundle_signal_in_channels, base_justify=base_justify, ) @@ -579,8 +588,8 @@ def _run_bundle_batch(sigs: list, seqs: list, feats: list, rids: list) -> None: motif=motif, motif_offset=motif_offset, reference_sequences=reference_sequences, - skip_motif_indels=config.get("skip_motif_indels", False), - require_query_mapping=config.get("require_query_mapping", True), + skip_motif_indels=spec.skip_motif_indels, + require_query_mapping=spec.require_query_mapping, ), chunk=ChunkConfig( base_justify=base_justify, @@ -588,7 +597,7 @@ def _run_bundle_batch(sigs: list, seqs: list, feats: list, rids: list) -> None: feature_end=_feature_end, signal_context=signal_context, kmer_context=kmer_context, - recover_softclip_signal=config.get("recover_softclip_signal", False), + recover_softclip_signal=spec.recover_softclip_signal, ), seq_encoding=seq_encoding, signal_kmer_context=signal_kmer_context, @@ -694,6 +703,7 @@ def _run_bundle_batch(sigs: list, seqs: list, feats: list, rids: list) -> None: # Flush remaining chunks for this mega-batch accumulator.flush() + gpu_runner.drain() # -- Aggregate per-read and write BAM for this mega-batch -- batch_preds = _write_bundle_mega_batch( @@ -722,7 +732,7 @@ def _run_bundle_batch(sigs: list, seqs: list, feats: list, rids: list) -> None: advance=0, description=( f"[cyan]Processed {n_chunks} chunks from " - f"{n_reads} reads ({n_batches_done} batches, " + f"{n_reads} reads ({scorer.n_batches_done} batches, " f"{n_predicted} predicted)..." ), ) @@ -730,6 +740,7 @@ def _run_bundle_batch(sigs: list, seqs: list, feats: list, rids: list) -> None: logger.info(f"Extracted and inferred {n_chunks} chunks from {n_reads} reads") logger.info(f"Predicted {n_predicted} reads") bam_out.close() + gpu_runner.shutdown() logger.info(f"Bundle inference complete: {n_reads} reads, {len(pairs)} models") logger.info(f"Output written to: {output_path}") @@ -821,7 +832,7 @@ def _run_bundle_batch(sigs: list, seqs: list, feats: list, rids: list) -> None: feat, feat_start=_feature_start if _feature_start is not None - else -_kmer_context, + else -kmer_context, dwell_templates=dwell_templates_arr, template_min_pos=dwell_template_min_pos, ) @@ -842,6 +853,7 @@ def _run_bundle_batch(sigs: list, seqs: list, feats: list, rids: list) -> None: # Flush remaining chunks for this mega-batch accumulator.flush() + gpu_runner.drain() # -- Aggregate per-read and write BAM for this mega-batch -- batch_preds = _write_bundle_mega_batch( @@ -870,7 +882,7 @@ def _run_bundle_batch(sigs: list, seqs: list, feats: list, rids: list) -> None: advance=0, description=( f"[cyan]Processed {n_chunks} chunks from " - f"{n_reads} reads ({n_batches_done} batches, " + f"{n_reads} reads ({scorer.n_batches_done} batches, " f"{n_predicted} predicted)..." ), ) @@ -878,6 +890,7 @@ def _run_bundle_batch(sigs: list, seqs: list, feats: list, rids: list) -> None: logger.info(f"Extracted and inferred {n_chunks} chunks from {n_reads} reads") logger.info(f"Predicted {n_predicted} reads") bam_out.close() + gpu_runner.shutdown() logger.info(f"Bundle inference complete: {n_reads} reads, {len(pairs)} models") logger.info(f"Output written to: {output_path}") @@ -957,68 +970,38 @@ def _run_bundle_batch(sigs: list, seqs: list, feats: list, rids: list) -> None: if not positions: continue - base_idx = positions[0] - chunk = leech_read.get_chunk(base_idx, config=bundle_chunk_config) - if chunk is None: - continue - - # Prepare tensors (with optional kmer residual channel) - signal_array = chunk["signal"] - assert isinstance(signal_array, np.ndarray) - sig = signal_array.astype(np.float32) - sig_residual = chunk.get("signal_residual") - if sig_residual is not None: - sig_residual = sig_residual.astype(np.float32) - if len(sig_residual) < len(sig): - sig_residual = np.pad( - sig_residual, - (0, len(sig) - len(sig_residual)), - mode="constant", - ) - elif len(sig_residual) > len(sig): - sig_residual = sig_residual[: len(sig)] - sig = np.stack([sig, sig_residual], axis=0) - signal_t = torch.from_numpy(sig) - - seq_t = _encode_sequence_for_inference( - chunk, seq_encoding, signal_len, signal_kmer_context - ) - if seq_t is None: - continue - - n_chunks += 1 - - feat_t = None - if needs_features: - features_array = chunk["features"] - assert isinstance(features_array, np.ndarray) - # Templates first, then narrowing -- the order - # `dataset.py` used at training time. This site had it - # the other way round, so the template channels were - # keyed to the pre-narrowing column 0 while the array - # had already been shifted out from under them. - features_array = prepare_inference_features( - features_array.astype(np.float32), - kmer_len=kmer_len, - feature_start=chunk.get("feature_start"), - dwell_offset=dwell_offset, - wide_features=wide_features, - dwell_templates=dwell_templates_arr, - template_min_pos=dwell_template_min_pos, - ) - feat_t = torch.from_numpy(features_array) + # Bundle semantics: one chunk per read (the first motif + # position), unlike single-model predict's all-positions + # loop -- see chunks_for_read's own note on why read + # construction and position-finding stay with the caller. + for sig, seq_arr, feat, _base_idx in chunks_for_read( + leech_read, + positions[:1], + chunk_config=bundle_chunk_config, + signal_len=signal_len, + seq_encoding=seq_encoding, + signal_kmer_context=signal_kmer_context, + requires_features=needs_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, + ): + n_chunks += 1 + signal_t = torch.from_numpy(sig) + seq_t = torch.from_numpy(seq_arr) + feat_t = torch.from_numpy(feat) if feat is not None else None - if not _shape_validated: - _feat_for_check = ( - features_array.astype(np.float32) if needs_features else None - ) - validate_inference_shapes(sig, _feat_for_check, config) - _shape_validated = True + if not _shape_validated: + validate_inference_shapes(sig, feat, config) + _shape_validated = True - accumulator.add(signal_t, seq_t, feat_t, leech_read.read_id) + accumulator.add(signal_t, seq_t, feat_t, leech_read.read_id) # Flush remaining chunks for this mega-batch accumulator.flush() + gpu_runner.drain() # -- Aggregate per-read and write BAM for this mega-batch -- batch_preds = _write_bundle_mega_batch( @@ -1046,13 +1029,14 @@ def _run_bundle_batch(sigs: list, seqs: list, feats: list, rids: list) -> None: advance=0, description=( f"[cyan]Processed {n_chunks} chunks from {n_reads} reads " - f"({n_batches_done} batches, {n_predicted} predicted)..." + f"({scorer.n_batches_done} batches, {n_predicted} predicted)..." ), ) logger.info(f"Extracted and inferred {n_chunks} chunks from {n_reads} reads") logger.info(f"Predicted {n_predicted} reads") bam_out.close() + gpu_runner.shutdown() logger.info(f"Bundle inference complete: {n_reads} reads, {len(pairs)} models") logger.info(f"Output written to: {output_path}") diff --git a/src/leech/inference/helpers.py b/src/leech/inference/helpers.py index 14a461a4..ecd0f94a 100644 --- a/src/leech/inference/helpers.py +++ b/src/leech/inference/helpers.py @@ -4,7 +4,9 @@ import json import logging import threading -from collections.abc import Callable, Hashable +from collections.abc import Callable, Hashable, Mapping +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass from pathlib import Path import numpy as np @@ -17,11 +19,16 @@ rust_supports_norm_method, rust_supports_softclip_recovery, ) -from leech.chunking import extraction_sequence -from leech.constants import BELOW_THRESHOLD_LABEL, DEFAULT_REFINE_HALF_BANDWIDTH -from leech.features import encode_signal_kmer, sequence_to_int +from leech.chunking import extraction_sequence, feature_window_from_metadata +from leech.constants import ( + BELOW_THRESHOLD_LABEL, + DEFAULT_REFINE_HALF_BANDWIDTH, + DEFAULT_SEQ_ENCODING_FALLBACK, +) +from leech.features import encode_signal_kmer, encode_signal_kmer_batch, sequence_to_int from leech.model_loading import load_model_from_checkpoint from leech.models import requires_features as _requires_features +from leech.models import wide_features as _wide_features from leech.models.inference_wrapper import ModelInferenceWrapper, TracedModelWrapper from leech.models.remora_compat import RemoraModelWrapper from leech.preparation import encode_kmer @@ -29,6 +36,24 @@ logger = logging.getLogger("leech.inference") +def is_multiclass(config: Mapping) -> bool: + """Whether a model's config describes a categorical (softmax) output. + + One definition, for every place that decides how predictions are + represented, tagged and calibrated: ``num_out > 1``. Before this, + ``evaluation.py`` and ``commands/calibrate.py`` used ``> 2`` here while + predict used ``> 1``, so a 2-output cross-entropy model got multiclass BAM + tags (``aa``/``pn``/``pp``) from predict but binary Platt calibration that + predict never applied to it (issue #269). + + This is a different question from ``Trainer``'s internal ``_num_out > 2`` + branches in the validation loop, which pick a *metric* (macro-F1 plus + accuracy for 3+ classes vs. an AUROC-style treatment for a 1- or 2-output + model) given an already-known loss type, and are left alone here. + """ + return config.get("num_out", 1) > 1 + + def _write_prediction_tags( aln: pysam.AlignedSegment, predicted_aa: str, @@ -127,6 +152,8 @@ def _check_config_consistency[T]( cli_value: T, config_value: T | None, cli_default: T, + *, + explicit: bool | None = None, ) -> T: """Resolve inference param from config, erroring on CLI conflict. @@ -134,9 +161,18 @@ def _check_config_consistency[T]( - config has value + CLI is default -> use config (normal auto-read) - config has value + CLI differs -> raise InferenceConfigError - config is None/missing -> use CLI value (old models without field) + + ``explicit``, when given, comes from click's own + ``get_parameter_source(param_name) is COMMANDLINE`` rather than being + inferred from ``cli_value != cli_default`` -- the inferred form cannot + tell an explicit ``--motif-offset 0`` from "not passed" when the default + is also 0, so that call silently lost to the config instead of raising + on a real conflict. Callers with no click context (tests, programmatic + use) omit it and keep the inferred heuristic. """ if config_value is not None: - if cli_value != cli_default and cli_value != config_value: + is_explicit = explicit if explicit is not None else cli_value != cli_default + if is_explicit and cli_value != config_value: raise InferenceConfigError( f"CLI --{param_name}={cli_value!r} conflicts with training config " f"{param_name}={config_value!r}. Inference must use the same " @@ -146,6 +182,250 @@ def _check_config_consistency[T]( return cli_value +@dataclass(frozen=True) +class InferenceSpec: + """Everything predict needs to know about a model, resolved once. + + Replaces the config-resolution block that used to be pasted into + ``run_inference`` and ``run_bundle_inference`` separately (motif/anchor/ + base_justify via ``_check_config_consistency``, refinement setup, the + feature-window fallback chain, ``wide_features`` detection, the + ``is_multiclass``/``seq_encoding`` default disagreements) -- issue #269. + Both call :meth:`from_config` once and read every other field off the + result rather than re-deriving it. + """ + + # Model geometry + signal_len: int + kmer_len: int + signal_context: tuple[int, int] + kmer_context: int + dwell_offset: int + seq_encoding: str + signal_kmer_context: tuple[int, int] + wide_features: bool + requires_features: bool + signal_in_channels: int + feature_start: int | None + feature_end: int | None + + # Motif / anchoring + motif: str + motif_offset: int + anchor: str + base_justify: str + reference_fasta: Path | None + skip_motif_indels: bool + require_query_mapping: bool + recover_softclip_signal: bool + + # Signal map refinement + refine_signal_map: bool + refine_half_bandwidth: int + refine_do_rough_rescale: bool + refine_scale_iters: int + refine_kmer_center_idx: int + + # Output representation + is_multiclass: bool + num_out: int + label_map: dict[str, int] | None + calibration: dict | None + cl_regression: bool + dwell_template_table: str | None + + @classmethod + def from_config( + cls, + config: Mapping, + *, + model_type: str = "", + motif: str | None = None, + motif_offset: int = 0, + base_justify: str = "center", + anchor: str = "reference", + reference_fasta: Path | None = None, + default_refine_scale_iters: int = 2, + parameter_sources: Mapping[str, bool] | None = None, + strict_model_type: bool = False, + model_dwell_margin: Callable[[], int] | None = None, + ) -> "InferenceSpec": + """Resolve one :class:`InferenceSpec` from a model/bundle config dict. + + Args: + config: A leech checkpoint's ``config.json``, a bundle's + ``config``, or the dict :func:`load_model_auto` synthesizes + for a Remora model. All three share this shape. + model_type: The registry name (``ModelInferenceWrapper.model_type`` + or a bundle's ``model_name``), used only to resolve + ``wide_features``/``requires_features`` structurally via + :func:`leech.models.wide_features` / + :func:`leech.models.requires_features`. An unresolvable name + (e.g. a Remora model, or the empty default) falls back to + ``False`` for both rather than raising -- unlike bundling, + grid search and export, which read ``model_type`` from a real + config and treat an unresolvable name as a bug worth raising. + motif, motif_offset, base_justify, anchor, reference_fasta: + CLI-provided values, checked against the config for conflicts + via :func:`_check_config_consistency`. + default_refine_scale_iters: The literal fallback for + ``refine_scale_iters`` when the config lacks the key -- + Remora configs default to ``-1`` (no refinement DP pass), + leech configs to ``2``. Callers must pass the right one; this + function does not infer it from the config shape. + parameter_sources: Optional ``{param_name: was_explicit}`` map + from click's ``get_parameter_source``, passed through to + :func:`_check_config_consistency` so an explicit CLI value + that happens to equal the default is still checked against + the config rather than assumed to be "not passed". + strict_model_type: When ``True`` and ``model_type`` is truthy but + not a recognized registry name, raise ``KeyError`` instead of + falling back to ``False`` for ``wide_features``/ + ``requires_features``. Bundle callers pass this: a bundle's + ``model_type`` is read back from its own saved metadata, so + an unresolvable name means the bundle is stale, renamed out + from under it, or corrupted -- a real bug worth raising + loudly rather than silently guessing its feature-window + convention. Single-model callers leave this ``False``: a + Remora model has no registry name at all (``model_type`` is + the empty-string default), which is expected, not an error. + model_dwell_margin: Optional zero-arg callable returning a wide- + feature model's real ``dwell_margin`` attribute (e.g. + ``lambda: getattr(model_wrapper.model, "dwell_margin", 0)``), + for the rare fallback below. ``dwell_margin`` is a model + constructor default, never written into ``config.json``, so + reading it needs the model itself -- lazy because building + one (bundle.py's caller has to ``_instantiate_model(config)``) + is wasted work on every call that never reaches this branch. + """ + sources = parameter_sources or {} + + motif = _check_config_consistency( + "motif", motif, config.get("motif"), None, explicit=sources.get("motif") + ) + motif_offset = _check_config_consistency( + "motif-offset", + motif_offset, + config.get("motif_offset"), + 0, + explicit=sources.get("motif_offset"), + ) + if motif is None: + raise InferenceConfigError( + "motif is None after auto-read from config. Either pass --motif on " + "the CLI or ensure the config contains a non-null 'motif' field. " + "Without a motif, inference predicts at every position, producing noise." + ) + + base_justify = _check_config_consistency( + "base-justify", + base_justify, + config.get("base_justify"), + "center", + explicit=sources.get("base_justify"), + ) + anchor = _check_config_consistency( + "anchor", anchor, config.get("anchor"), "reference", explicit=sources.get("anchor") + ) + + resolved_reference_fasta = reference_fasta + if resolved_reference_fasta is None: + cfg_ref = config.get("reference_fasta") + if cfg_ref is not None: + cfg_path = Path(cfg_ref) + if cfg_path.exists(): + resolved_reference_fasta = cfg_path + else: + logger.warning( + f"reference_fasta from config ({cfg_ref}) not found; " + f"pass --reference-fasta explicitly" + ) + + signal_len = int(config.get("signal_len", 100)) + kmer_len = int(config.get("kmer_len", 9)) + kmer_context = kmer_len // 2 + left_ctx = config.get("left_context") + right_ctx = config.get("right_context") + signal_context = ( + (int(left_ctx), int(right_ctx)) + if left_ctx is not None and right_ctx is not None + else (signal_len // 2, signal_len // 2) + ) + + seq_encoding = config.get("seq_encoding", DEFAULT_SEQ_ENCODING_FALLBACK) + signal_kmer_context = tuple(config.get("signal_kmer_context", (4, 4))) + dwell_offset = int(config.get("dwell_offset", 0)) + + try: + wide_features = _wide_features(model_type) + except KeyError as e: + if strict_model_type and model_type: + raise KeyError( + f"Model architecture '{model_type}' is not a recognized " + f"model (renamed, removed, or a corrupted checkpoint/bundle). " + f"Cannot determine its feature-window convention." + ) from e + wide_features = False + feature_start, feature_end = feature_window_from_metadata(config, kmer_context) + if wide_features and feature_start is None and feature_end is None: + model_margin = model_dwell_margin() if model_dwell_margin is not None else 0 + if model_margin: + feature_start = -(kmer_context + model_margin) + feature_end = kmer_context + model_margin + logger.warning( + f"Config missing feature_start/end, " + f"falling back to model default margin: {model_margin}" + ) + + num_out = int(config.get("num_out", 1)) + label_map = config.get("label_map") + + signal_in_channels = int(config.get("signal_in_channels", 1)) + try: + requires_features = _requires_features(model_type) + except KeyError: + requires_features = False + + refine_signal_map = bool(config.get("refine_signal_map", True) or signal_in_channels > 1) + refine_scale_iters = int(config.get("refine_scale_iters", default_refine_scale_iters)) + + return cls( + signal_len=signal_len, + kmer_len=kmer_len, + signal_context=signal_context, + kmer_context=kmer_context, + dwell_offset=dwell_offset, + seq_encoding=seq_encoding, + signal_kmer_context=signal_kmer_context, + wide_features=wide_features, + requires_features=requires_features, + signal_in_channels=signal_in_channels, + feature_start=feature_start, + feature_end=feature_end, + motif=motif, + motif_offset=motif_offset, + anchor=anchor, + base_justify=base_justify, + reference_fasta=resolved_reference_fasta, + skip_motif_indels=bool(config.get("skip_motif_indels", False)), + require_query_mapping=bool(config.get("require_query_mapping", True)), + recover_softclip_signal=bool(config.get("recover_softclip_signal", False)), + refine_signal_map=refine_signal_map, + refine_half_bandwidth=int( + config.get("refine_half_bandwidth", DEFAULT_REFINE_HALF_BANDWIDTH) + ), + refine_do_rough_rescale=bool(config.get("refine_do_rough_rescale", True)), + refine_scale_iters=refine_scale_iters, + refine_kmer_center_idx=int(config.get("refine_kmer_center_idx", -1)), + is_multiclass=num_out > 1, + num_out=num_out, + label_map=label_map, + calibration=config.get("calibration") if num_out > 1 else None, + cl_regression=bool(config.get("cl_regression", False)), + dwell_template_table=config.get("dwell_template_table") or None, + ) + + def validate_inference_shapes( signal: np.ndarray, features: np.ndarray | None, @@ -481,6 +761,133 @@ def _encode_sequence_for_inference( return encode_kmer(chunk["sequence"]) +def chunks_for_read( + leech_read, + positions: list[int], + *, + chunk_config, + signal_len: int, + seq_encoding: str, + signal_kmer_context: tuple[int, int], + requires_features: bool, + kmer_len: int, + dwell_offset: int, + wide_features: bool, + dwell_templates: np.ndarray | None = None, + template_min_pos: int = 0, +) -> list[tuple[np.ndarray, np.ndarray, np.ndarray | None, int]]: + """Extract ready-to-batch ``(signal, sequence, features, base_idx)`` tuples + from an already-built :class:`~leech.chunking.LeechRead` at each of + ``positions``. + + The one place the per-chunk array-building triplet runs -- + :func:`prepare_signal_channels`, :func:`_encode_sequence_for_inference`, + :func:`prepare_inference_features` -- for the Python serial extraction + path. ``run_inference``'s sequential path and ``run_bundle_inference``'s + serial path each carried their own copy of this loop, and only the bundle + copy applied dwell templates; the bundle copy also built its signal array + by hand rather than through ``prepare_signal_channels``, so it never + padded/cropped ``signal`` to ``signal_len`` the way every other extraction + path does (issue #268 -- found while unifying, not by design). + + Building the ``LeechRead`` itself, and finding ``positions`` on it, stays + with the caller: a ``pysam.AlignedSegment`` and a ``ReadInfo`` differ too + much above this point (motif search takes a real alignment object one + caller has and the other reconstructs a mock of) for a shared protocol to + buy anything; this is the part that was actually duplicated array-copy + logic rather than read-shape plumbing. + + ``signal_kmer`` is batched once per read rather than once per chunk + (issue #260): every chunk's raw ``sequence_with_kmer_context``/ + ``seq_to_sig_map`` is collected first and run through + :func:`~leech.features.encode_signal_kmer_batch` in one call, instead of + one pyo3 call per chunk via :func:`_encode_sequence_for_inference`. Both + ``run_inference`` and ``run_bundle_inference`` get this from being routed + through this one function -- before #268 unified them, only the + ``run_inference`` copy of this loop had been given the batched path. + """ + results: list[tuple[np.ndarray, np.ndarray, np.ndarray | None, int]] = [] + + if seq_encoding != "signal_kmer": + for base_idx in positions: + chunk = leech_read.get_chunk(base_idx, config=chunk_config) + if chunk is None: + continue + + sig = prepare_signal_channels(chunk, signal_len) + + seq_enc = _encode_sequence_for_inference( + chunk, seq_encoding, signal_len, signal_kmer_context + ) + if seq_enc is None: + continue + seq_arr = seq_enc.numpy() if isinstance(seq_enc, torch.Tensor) else seq_enc + + feat = None + if requires_features: + features_array = chunk["features"] + assert isinstance(features_array, np.ndarray) + feat = prepare_inference_features( + features_array.astype(np.float32), + kmer_len=kmer_len, + feature_start=chunk.get("feature_start"), + dwell_offset=dwell_offset, + wide_features=wide_features, + dwell_templates=dwell_templates, + template_min_pos=template_min_pos, + ) + + results.append((sig, seq_arr, feat, base_idx)) + return results + + # signal_kmer: collect every chunk's raw inputs first, batch-encode once. + pending: list[tuple[np.ndarray, np.ndarray | None, int, np.ndarray, np.ndarray]] = [] + for base_idx in positions: + chunk = leech_read.get_chunk(base_idx, config=chunk_config) + if chunk is None: + continue + + seq_ctx = chunk.get("sequence_with_kmer_context") + seq_to_sig = chunk.get("seq_to_sig_map") + if seq_ctx is None or seq_to_sig is None: + continue + + sig = prepare_signal_channels(chunk, signal_len) + + feat = None + if requires_features: + features_array = chunk["features"] + assert isinstance(features_array, np.ndarray) + feat = prepare_inference_features( + features_array.astype(np.float32), + kmer_len=kmer_len, + feature_start=chunk.get("feature_start"), + dwell_offset=dwell_offset, + wide_features=wide_features, + dwell_templates=dwell_templates, + template_min_pos=template_min_pos, + ) + + pending.append((sig, feat, base_idx, sequence_to_int(seq_ctx), seq_to_sig)) + + if pending: + n = len(pending) + max_si = max(p[3].shape[0] for p in pending) + max_s2 = max(p[4].shape[0] for p in pending) + padded_seq_ints = np.full((n, max_si), -1, dtype=np.int8) + padded_s2s = np.full((n, max_s2), signal_len, dtype=np.int32) + for i, (_, _, _, seq_ints, s2s) in enumerate(pending): + padded_seq_ints[i, : seq_ints.shape[0]] = seq_ints + padded_s2s[i, : s2s.shape[0]] = s2s + batch_enc = encode_signal_kmer_batch( + padded_seq_ints, padded_s2s, signal_len, signal_kmer_context + ) + for i, (sig, feat, base_idx, _, _) in enumerate(pending): + results.append((sig, batch_enc[i], feat, base_idx)) + + return results + + def _write_mega_batch_predictions( aln_batch: list[pysam.AlignedSegment], pending: dict[str, list], @@ -536,6 +943,61 @@ def _write_mega_batch_predictions( return n_preds +class GpuBatchRunner: + """Double-buffered, single-worker async submission of one scoring function. + + The pattern ``run_inference``'s and ``run_bundle_inference``'s serial + paths each hand-rolled: a dedicated single-thread executor runs + ``score_fn`` on one batch while the caller's extraction loop keeps filling + the next one. Submitting a new batch waits for the previous one to finish + first, so at most one is ever in flight — the ``max_workers=1`` GPU + executor invariant PR #253 established, preserved here rather than + reintroduced independently on the bundle path (issue #268). + + ``score_fn(signals, sequences, features, meta)`` is the same shape + :class:`BatchAccumulator`'s flush callback expects, so a bound + :meth:`submit` is a drop-in ``flush_fn``:: + + runner = GpuBatchRunner(score_fn) + accumulator = BatchAccumulator(batch_size, runner.submit) + ... + accumulator.flush() + runner.drain() + runner.shutdown() + """ + + __slots__ = ("_executor", "_future", "score_fn") + + def __init__(self, score_fn: Callable[[list, list, list, list], None]): + self.score_fn = score_fn + self._executor = ThreadPoolExecutor(max_workers=1) + self._future: Future | None = None + + def submit(self, signals: list, sequences: list, features: list, meta: list) -> None: + """Wait for any in-flight batch, then submit this one (non-blocking).""" + if self._future is not None: + self._future.result() + self._future = self._executor.submit(self.score_fn, signals, sequences, features, meta) + + def drain(self) -> None: + """Block until the most recently submitted batch has finished.""" + if self._future is not None: + self._future.result() + self._future = None + + def shutdown(self) -> None: + """Wait for the executor's thread to exit. + + Always call after :meth:`drain` — every batch is done by then, so + this costs nothing, but a caller that goes on to fork (an ``mp.Pool`` + on the parallel path) must not leave the thread alive across it: a + fork inherits a running thread's memory but not the thread itself, so + any lock it held never releases (see ``run_inference``'s own note on + this at its ``_extract_pool.shutdown(wait=True)`` call). + """ + self._executor.shutdown(wait=True) + + class BatchAccumulator: """Four parallel chunk buffers, a size check, and a flush — in one place. diff --git a/src/leech/inference/single.py b/src/leech/inference/single.py index 4a482d21..a2598a56 100644 --- a/src/leech/inference/single.py +++ b/src/leech/inference/single.py @@ -13,18 +13,19 @@ from rich.progress import Progress from leech.configs import ChunkConfig, InferenceConfig, MotifConfig, SignalConfig -from leech.constants import DEFAULT_REFINE_HALF_BANDWIDTH, TORCH_COMPILE_MIN_SAMPLES +from leech.constants import TORCH_COMPILE_MIN_SAMPLES from leech.features import encode_signal_kmer_batch, extract_move_table, sequence_to_int from leech.inference.helpers import ( BatchAccumulator, - _check_config_consistency, - _encode_sequence_for_inference, + GpuBatchRunner, + InferenceSpec, _run_batch, _run_batch_multiclass, _write_mega_batch_predictions, build_rust_extraction_kwargs, cap_rayon_threads_for_slurm, check_rust_extraction_available, + chunks_for_read, collect_bam_metadata_for_rust, load_model_auto, prepare_inference_features, @@ -35,7 +36,6 @@ from leech.io.motif_search import get_motif_searcher from leech.io.pod5_reader import POD5Reader from leech.models.inference_wrapper import ModelInferenceWrapper -from leech.models.inference_wrapper import resolve_wide_features as _resolve_wide_features from leech.models.remora_compat import RemoraModelWrapper from leech.preparation.reader import build_leech_read @@ -235,6 +235,7 @@ def run_inference( 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. @@ -273,6 +274,11 @@ def run_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``). """ # Apply backend override to signal_refine module logger.info(f"Extraction backend: {backend}") @@ -298,178 +304,71 @@ def run_inference( # Determine if this is a Remora model or a leech model is_remora = config.get("is_remora", False) - # Signal map refinement setup - refine_signal_map = False - signal_refiner = None - if is_remora: model_wrapper = wrapper_or_model - signal_len = config.get("signal_len", 100) - kmer_len = config.get("kmer_len", 9) - seq_encoding = "signal_kmer" - signal_kmer_context = tuple(config.get("signal_kmer_context", (4, 4))) - dwell_offset = 0 - - # Resolve motif/offset from config, erroring on CLI conflict - motif = _check_config_consistency("motif", motif, config.get("motif"), None) - motif_offset = _check_config_consistency( - "motif-offset", motif_offset, config.get("motif_offset"), 0 - ) - if motif is not None: - logger.info(f"Motif from remora config: {motif} (offset={motif_offset})") - - if motif is None: - raise ValueError("--motif is required for Remora models (no config.json)") - - # Set up signal map refinement if model specifies it - if config.get("refine_signal_map", True): - 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) - half_bw = config.get("refine_half_bandwidth", DEFAULT_REFINE_HALF_BANDWIDTH) - do_rescale = config.get("refine_do_rough_rescale", True) - scale_iters = config.get("refine_scale_iters", -1) - center_idx = config.get("refine_kmer_center_idx", -1) - signal_refiner = SigMapRefiner.from_table( - kmer_table_path, - half_bandwidth=half_bw, - do_rough_rescale=do_rescale, - scale_iters=scale_iters, - center_idx=center_idx, - ) - refine_signal_map = True - logger.info( - f"Signal map refinement: half_bw={half_bw}, " - f"scale_iters={scale_iters}, center_idx={center_idx}" - ) - else: - # Leech model - if isinstance(wrapper_or_model, ModelInferenceWrapper): - model_wrapper = wrapper_or_model - else: - model_type = config["model_name"] - model_wrapper = ModelInferenceWrapper(wrapper_or_model, model_type) - - signal_len = config["signal_len"] - kmer_len = config["kmer_len"] - dwell_offset = config.get("dwell_offset", 0) - seq_encoding = config.get("seq_encoding", "signal_kmer") - signal_kmer_context = tuple(config.get("signal_kmer_context", (4, 4))) - - # Resolve motif/offset from config, erroring on CLI conflict - motif = _check_config_consistency("motif", motif, config.get("motif"), None) - motif_offset = _check_config_consistency( - "motif-offset", motif_offset, config.get("motif_offset"), 0 - ) - if motif is not None: - logger.info(f"Motif from config: {motif} (offset={motif_offset})") - - if motif is None: - raise ValueError( - "motif is None after auto-read from config. " - "Either pass --motif on the CLI or ensure config.json contains a non-null 'motif' field. " - "Without a motif, inference predicts at every position, producing noise." - ) - - # Signal map refinement for leech models (needed for kmer residual signal channel) - if config.get("refine_signal_map", True) or config.get("signal_in_channels", 1) > 1: - 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) - half_bw = config.get("refine_half_bandwidth", DEFAULT_REFINE_HALF_BANDWIDTH) - do_rescale = config.get("refine_do_rough_rescale", True) - scale_iters = config.get("refine_scale_iters", 2) - center_idx = config.get("refine_kmer_center_idx", -1) - signal_refiner = SigMapRefiner.from_table( - kmer_table_path, - half_bandwidth=half_bw, - do_rough_rescale=do_rescale, - scale_iters=scale_iters, - center_idx=center_idx, - ) - refine_signal_map = True - logger.info( - f"Signal map refinement enabled for leech model " - f"(signal_in_channels={config.get('signal_in_channels', 1)})" - ) - - # Use asymmetric context if available, otherwise fall back to symmetric - left_ctx = config.get("left_context") - right_ctx = config.get("right_context") - if left_ctx is not None and right_ctx is not None: - signal_context = (left_ctx, right_ctx) + elif isinstance(wrapper_or_model, ModelInferenceWrapper): + model_wrapper = wrapper_or_model else: - signal_context = (signal_len // 2, signal_len // 2) - kmer_context = kmer_len // 2 - requires_features = getattr(model_wrapper, "requires_features", False) - - # Determine feature_start/feature_end from config (must match training data) - _model_type = getattr(model_wrapper, "model_type", "") - wide_features = _resolve_wide_features(model_wrapper, _model_type) - _kmer_context = kmer_len // 2 - - # Read new params, falling back to old dwell_margin_* for backward compat - _feature_start = config.get("feature_start") - _feature_end = config.get("feature_end") - if _feature_start is None and "feature_left" in config: - _feature_start = -config["feature_left"] - if _feature_end is None and "feature_right" in config: - _feature_end = config["feature_right"] - if _feature_start is None and "dwell_margin_left" in config: - _feature_start = -(_kmer_context + config["dwell_margin_left"]) - if _feature_end is None and "dwell_margin_right" in config: - _feature_end = _kmer_context + config["dwell_margin_right"] - if wide_features and _feature_start is None and _feature_end is None: - _model_margin = ( + model_wrapper = ModelInferenceWrapper(wrapper_or_model, config["model_name"]) + + 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 - ) - if _model_margin: - _feature_start = -(_kmer_context + _model_margin) - _feature_end = _kmer_context + _model_margin - logger.warning( - f"Config missing feature_start/end, " - f"falling back to model default margin: {_model_margin}" - ) - - # Detect multi-class model - num_out = config.get("num_out", 1) - label_map = config.get("label_map") # {name: int} or None - if label_map: - # Invert to {int: name} - int_to_label = {v: k for k, v in label_map.items()} - else: - int_to_label = None - is_multiclass = num_out > 1 - - # Resolve base_justify from config, erroring on CLI conflict - base_justify = _check_config_consistency( - "base-justify", base_justify, config.get("base_justify"), "center" + ), ) - logger.info(f"base_justify: {base_justify}") + 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 + + int_to_label = {v: k for k, v in spec.label_map.items()} if spec.label_map else None - anchor = _check_config_consistency("anchor", anchor, config.get("anchor"), "reference") - logger.info(f"anchor: {anchor}") - - if reference_fasta is None: - cfg_ref = config.get("reference_fasta") - if cfg_ref is not None: - cfg_path = Path(cfg_ref) - if cfg_path.exists(): - reference_fasta = cfg_path - logger.info(f"reference_fasta from config: {reference_fasta}") - else: - logger.warning( - f"reference_fasta from config ({cfg_ref}) not found; " - f"pass --reference-fasta explicitly" - ) + # 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, + do_rough_rescale=spec.refine_do_rough_rescale, + 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: @@ -477,11 +376,13 @@ def run_inference( 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 + _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})") - if motif: - logger.info(f"Motif: {motif} (offset={motif_offset})") + 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") @@ -559,7 +460,7 @@ def run_inference( # 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 = config.get("signal_in_channels", 1) + 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 @@ -574,13 +475,13 @@ def run_inference( motif_searcher = get_motif_searcher( mode="fasta" if reference_sequences else "bam", reference_sequences=reference_sequences, - skip_indels=config.get("skip_motif_indels", False), + 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=config.get("require_query_mapping", True), + require_query_mapping=spec.require_query_mapping, ) # Prepare class_names_str for multiclass (shared across mega-batches) @@ -660,8 +561,8 @@ def run_inference( motif=motif, motif_offset=motif_offset, reference_sequences=reference_sequences, - skip_motif_indels=config.get("skip_motif_indels", False), - require_query_mapping=config.get("require_query_mapping", True), + skip_motif_indels=spec.skip_motif_indels, + require_query_mapping=spec.require_query_mapping, ), chunk=ChunkConfig( base_justify=base_justify, @@ -669,7 +570,7 @@ def run_inference( feature_end=_feature_end, signal_context=signal_context, kmer_context=kmer_context, - recover_softclip_signal=config.get("recover_softclip_signal", False), + recover_softclip_signal=spec.recover_softclip_signal, ), seq_encoding=seq_encoding, signal_kmer_context=signal_kmer_context, @@ -681,7 +582,7 @@ def run_inference( signal_in_channels=signal_in_channels, ) - calibration = config.get("calibration") if is_multiclass else None + calibration = spec.calibration _batch_fn_p = ( functools.partial( _run_batch_multiclass, @@ -791,7 +692,7 @@ def _run_worker_batch(sigs, seqs, feats, meta) -> None: pending: dict[str, list] = {} _shape_validated = False - calibration = config.get("calibration") if is_multiclass else None + calibration = spec.calibration _batch_fn = ( functools.partial( _run_batch_multiclass, @@ -802,42 +703,27 @@ def _run_worker_batch(sigs, seqs, feats, meta) -> None: else _run_batch ) - _gpu_executor = ThreadPoolExecutor(max_workers=1) - _gpu_future: Future | None = None - _bam_write_executor = ThreadPoolExecutor(max_workers=1) - _bam_write_future: Future | None = None - - def _submit_gpu_batch(sigs, seqs, feats, meta) -> None: - """Flush callback: hand one batch to the GPU thread (double-buffered). + def _score_batch(sigs, seqs, feats, meta) -> None: + """Score one batch into the *current* ``pending`` dict. - The accumulator has already detached these buffers, so the GPU - thread owns them and extraction can keep filling the next batch. + 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. """ - 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, - ) - - 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 + _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, @@ -855,7 +741,7 @@ def _drain_gpu() -> None: feature_end=_feature_end, signal_context=signal_context, kmer_context=kmer_context, - recover_softclip_signal=config.get("recover_softclip_signal", False), + recover_softclip_signal=spec.recover_softclip_signal, ) # Extraction thread count + rust setup (all three shared with @@ -942,68 +828,21 @@ def _extract_one_read( ) ] - results: list[tuple[np.ndarray, np.ndarray, np.ndarray | None, tuple[str, int]]] = [] - # signal_kmer encoding is batched once per read (issue #260) - # rather than once per chunk: `pending` holds every chunk's - # (signal, feat, key) plus its raw inputs, appended to `results` - # once the batch call below fills in the real sequence array. - pending: list[ - tuple[np.ndarray, np.ndarray | None, tuple[str, int], np.ndarray, np.ndarray] - ] = [] - for base_idx in positions: - chunk = leech_read.get_chunk(base_idx, config=seq_chunk_config) - if chunk is None: - continue - - # Signal (with optional kmer residual channel) - sig = prepare_signal_channels(chunk, signal_len) - - feat = None - if requires_features: - features_array = chunk["features"] - assert isinstance(features_array, np.ndarray) - feat = prepare_inference_features( - features_array.astype(np.float32), - kmer_len=kmer_len, - feature_start=chunk.get("feature_start"), - dwell_offset=dwell_offset, - wide_features=wide_features, - ) - - if seq_encoding == "signal_kmer": - seq_ctx = chunk.get("sequence_with_kmer_context") - seq_to_sig = chunk.get("seq_to_sig_map") - if seq_ctx is None or seq_to_sig is None: - continue - pending.append( - (sig, feat, (read_id, base_idx), sequence_to_int(seq_ctx), seq_to_sig) - ) - continue - - seq_enc = _encode_sequence_for_inference( - chunk, seq_encoding, signal_len, signal_kmer_context - ) - if seq_enc is None: - continue - seq_arr = seq_enc.numpy() if isinstance(seq_enc, torch.Tensor) else seq_enc - results.append((sig, seq_arr, feat, (read_id, base_idx))) - - if pending: - n = len(pending) - max_si = max(p[3].shape[0] for p in pending) - max_s2 = max(p[4].shape[0] for p in pending) - padded_seq_ints = np.full((n, max_si), -1, dtype=np.int8) - padded_s2s = np.full((n, max_s2), signal_len, dtype=np.int32) - for i, (_, _, _, seq_ints, s2s) in enumerate(pending): - padded_seq_ints[i, : seq_ints.shape[0]] = seq_ints - padded_s2s[i, : s2s.shape[0]] = s2s - batch_enc = encode_signal_kmer_batch( - padded_seq_ints, padded_s2s, signal_len, signal_kmer_context + return [ + (sig, seq_arr, feat, (read_id, base_idx)) + 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, ) - for i, (sig, feat, key, _, _) in enumerate(pending): - results.append((sig, batch_enc[i], feat, key)) - - return results + ] _extract_pool = ThreadPoolExecutor(max_workers=n_extract) @@ -1023,10 +862,8 @@ def _extract_one_read( signal_kmer_context=signal_kmer_context, refine_signal_map=refine_signal_map, signal_refiner=signal_refiner, - refine_half_bandwidth=config.get( - "refine_half_bandwidth", DEFAULT_REFINE_HALF_BANDWIDTH - ), - refine_scale_iters=config.get("refine_scale_iters", 2), + refine_half_bandwidth=spec.refine_half_bandwidth, + refine_scale_iters=spec.refine_scale_iters, signal_in_channels=signal_in_channels, base_justify=base_justify, ) @@ -1125,7 +962,7 @@ def _finalize_mega_batch(aln_batch_to_write): nonlocal total_reads, total_predictions, mega_batch_idx nonlocal pending, _bam_write_future accumulator.flush() - _drain_gpu() + 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 @@ -1384,15 +1221,15 @@ def _extraction_producer(): ) # 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 + # (`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_executor.shutdown(wait=True) + gpu_runner.shutdown() _wait_for_bam_write() # Ensure final BAM write completes before close _bam_write_executor.shutdown(wait=True) diff --git a/src/leech/model_export.py b/src/leech/model_export.py index abf0d4a6..de1eb994 100644 --- a/src/leech/model_export.py +++ b/src/leech/model_export.py @@ -14,6 +14,8 @@ import torch import torch.nn as nn +from leech.constants import DEFAULT_SEQ_ENCODING_FALLBACK + logger = logging.getLogger("leech.model_export") @@ -42,7 +44,7 @@ def _build_example_inputs( signal_len = config["signal_len"] kmer_len = config["kmer_len"] - seq_encoding = config.get("seq_encoding", "base_onehot") + seq_encoding = config.get("seq_encoding", DEFAULT_SEQ_ENCODING_FALLBACK) signal_kmer_context = config.get("signal_kmer_context", [4, 4]) if seq_encoding == "signal_kmer": diff --git a/src/leech/onnx_export.py b/src/leech/onnx_export.py index ae24da85..a2122e30 100644 --- a/src/leech/onnx_export.py +++ b/src/leech/onnx_export.py @@ -70,6 +70,8 @@ from pathlib import Path from typing import Any +from leech.constants import DEFAULT_SEQ_ENCODING_FALLBACK + __all__ = [ "OPSET", "InputSpec", @@ -106,7 +108,7 @@ def describe_inputs(config: dict, example_inputs: tuple) -> list[InputSpec]: the graph; the *roles* come from the config, which is the only place they exist. """ - seq_encoding = config.get("seq_encoding", "base_onehot") + seq_encoding = config.get("seq_encoding", DEFAULT_SEQ_ENCODING_FALLBACK) kmer_context = config.get("signal_kmer_context", [4, 4]) if seq_encoding == "signal_kmer": diff --git a/src/leech/training.py b/src/leech/training.py index 2314966d..83998f04 100644 --- a/src/leech/training.py +++ b/src/leech/training.py @@ -27,6 +27,7 @@ from torch.utils.data import DataLoader, DistributedSampler, Sampler, WeightedRandomSampler import leech +from leech.chunking import feature_window_from_metadata from leech.chunking.table import ChunkTable from leech.cli_config import make_console from leech.configs import AugmentConfig, AuxHeadConfig, OptimConfig, SchedulerConfig, TrainConfig @@ -2392,29 +2393,30 @@ def train_model( loss_type = "cross_entropy" logger.info(f"Model {model_name} has num_out={num_out}, switching to cross_entropy loss") - # Introspect feature_start/feature_end from raw training data - _raw_chunk = train_dataset.chunks[0] + # Introspect feature_start/feature_end from the training corpus. + # + # Resolved from the whole corpus (`train_dataset.chunks`, a ChunkTable + # when one is in use), not chunk 0 -- reading a single row and trusting it + # for the run is exactly the failure mode issue #230 fixed for + # seq_encoding, generalized here to the feature window (issue #269). + # `feature_window_from_metadata` raises if the corpus disagrees with + # itself; when `train_dataset.chunks` is a plain list (legacy in-memory + # path) it falls back to chunk 0, same as before. _kmer_context = kmer_len // 2 - if "feature_start" in _raw_chunk: - _feature_start = int(_raw_chunk["feature_start"]) - elif "feature_left" in _raw_chunk: - _feature_start = -int(_raw_chunk["feature_left"]) - elif "dwell_margin_left" in _raw_chunk: - _feature_start = -(_kmer_context + int(_raw_chunk["dwell_margin_left"])) - else: + _window_source = ( + train_dataset.chunks + if isinstance(train_dataset.chunks, ChunkTable) + else train_dataset.chunks[0] + ) + # For a mapping source (the legacy row-only path), the resolver already + # derives `feature_end` from the chunk's own feature array when only + # `dwell_margin_right` is stored; a ChunkTable has no single feature array + # to take a width from, so that one fallback is table-only unavailable and + # `feature_end` comes back None in that corner case. + _feature_start, _feature_end = feature_window_from_metadata(_window_source, _kmer_context) + if _feature_start is None: _feature_start = -_kmer_context - if "feature_end" in _raw_chunk: - _feature_end = int(_raw_chunk["feature_end"]) - elif "feature_right" in _raw_chunk: - _feature_end = int(_raw_chunk["feature_right"]) - elif "dwell_margin_right" in _raw_chunk: - _raw_features = _raw_chunk.get("features") - if _raw_features is not None and _raw_features.ndim > 1: - _feat_width = _raw_features.shape[1] - _feature_end = _feat_width - 1 + _feature_start - else: - _feature_end = _kmer_context - else: + if _feature_end is None: _feature_end = _kmer_context # Read preparation config sidecar if available (for provenance in config.json) diff --git a/tests/test_chunk_table.py b/tests/test_chunk_table.py index b08600a7..64e210c2 100644 --- a/tests/test_chunk_table.py +++ b/tests/test_chunk_table.py @@ -270,3 +270,73 @@ def test_round_trips(self, corpus): assert row["read_id"] == original["read_id"] assert row["label_int"] == original["label_int"] assert row["sequence"] == original["sequence"] + + +class TestFeatureWindowFromMetadataTableForm: + """The column form of ``feature_window_from_metadata`` -- read from the + whole corpus, not chunk 0 (issue #269, generalizing #230's failure mode). + """ + + def test_resolves_a_consistent_window(self, corpus): + from leech.chunking import feature_window_from_metadata + + path, _ = corpus # every chunk in `corpus` has feature_start=-6, feature_end=6 + table = ChunkTable.from_npz(path) + start, end = feature_window_from_metadata(table, kmer_context=5) + assert (start, end) == (-6, 6) + + def test_falls_back_to_dwell_margin_left_column(self, tmp_path): + from leech.chunking import feature_window_from_metadata + + chunks = make_chunks(4) + path = tmp_path / "legacy.npz" + np.savez( + path, + sequences=np.array([c["sequence"] for c in chunks], dtype=str), + labels=np.array([c["label"] for c in chunks], dtype=str), + labels_int=np.array([c["label_int"] for c in chunks], dtype=np.int64), + read_ids=np.array([c["read_id"] for c in chunks], dtype=str), + base_indices=np.array([c["base_idx"] for c in chunks], dtype=np.int64), + dwell_margin_lefts=np.full(len(chunks), 4, dtype=np.int64), + ) + table = ChunkTable.from_npz(path) + start, end = feature_window_from_metadata(table, kmer_context=5) + assert start == -(5 + 4) + assert end is None # no feature_end / dwell_margin_right column at all + + def test_no_window_fields_at_all_returns_none(self, tmp_path): + from leech.chunking import feature_window_from_metadata + + chunks = make_chunks(4) + path = tmp_path / "bare.npz" + np.savez( + path, + sequences=np.array([c["sequence"] for c in chunks], dtype=str), + labels=np.array([c["label"] for c in chunks], dtype=str), + labels_int=np.array([c["label_int"] for c in chunks], dtype=np.int64), + read_ids=np.array([c["read_id"] for c in chunks], dtype=str), + base_indices=np.array([c["base_idx"] for c in chunks], dtype=np.int64), + ) + table = ChunkTable.from_npz(path) + assert feature_window_from_metadata(table, kmer_context=5) == (None, None) + + def test_inconsistent_feature_start_raises(self, tmp_path): + """A corpus mixing two prepare runs with different --feature-start + values must not silently train on whichever value chunk 0 carried.""" + from leech.chunking import feature_window_from_metadata + + chunks = make_chunks(4) + path = tmp_path / "inconsistent.npz" + feature_starts = np.array([-6, -6, -5, -6], dtype=np.int64) + np.savez( + path, + sequences=np.array([c["sequence"] for c in chunks], dtype=str), + labels=np.array([c["label"] for c in chunks], dtype=str), + labels_int=np.array([c["label_int"] for c in chunks], dtype=np.int64), + read_ids=np.array([c["read_id"] for c in chunks], dtype=str), + base_indices=np.array([c["base_idx"] for c in chunks], dtype=np.int64), + feature_starts=feature_starts, + ) + table = ChunkTable.from_npz(path) + with pytest.raises(ValueError, match="inconsistent"): + feature_window_from_metadata(table, kmer_context=5) diff --git a/tests/test_inference.py b/tests/test_inference.py index ada995b4..fea11343 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -1136,6 +1136,112 @@ def test_pairwise_bundle_min_confidence_unc(self, tmp_path): assert all(r.get_tag("aa") == "unc" for r in tagged) +@pytest.mark.skipif(not TRNA_FIXTURES_AVAILABLE, reason="tRNA fixtures not available") +class TestBundleExtractionPathParity: + """``run_bundle_inference``'s three extraction paths (Rust, parallel, + serial) now share one accumulator, one ``GpuBatchRunner`` and one + ``BundleScorer`` (issue #268) -- so whichever path extracted the chunks, + the predictions must agree. Before this, only the serial path was + exercised by any test here.""" + + PAIR_NAMES = ["Ala_Gly", "Ala_Met", "Gly_Met"] + + def _predict(self, tmp_path, bundle_path, name, *, backend, num_workers=0): + from leech.inference import run_bundle_inference + + out = tmp_path / f"{name}.bam" + run_bundle_inference( + bundle_path=bundle_path, + pod5_path=TRNA_POD5, + bam_path=TRNA_BAM, + output_path=out, + device="cpu", + batch_size=8, + reverse_signal=True, + reference_fasta=TRNA_REF, + num_workers=num_workers, + backend=backend, + ) + with pysam.AlignmentFile(str(out), "rb") as bam: + return { + r.query_name: (r.get_tag("aa"), r.get_tag("ac")) for r in bam if r.has_tag("aa") + } + + def test_rust_path_writes_tags(self, tmp_path): + """``run_bundle_inference`` on the Rust monolithic extraction path, + the default for any run with leech_core installed and no explicit + --backend -- previously untested here.""" + pytest.importorskip("leech_core") + bundle_path = _create_pairwise_bundle(tmp_path, self.PAIR_NAMES) + rust_predicted = self._predict(tmp_path, bundle_path, "rust", backend="rust") + assert rust_predicted + + def test_parallel_path_alone_writes_tags(self, tmp_path): + """``num_workers > 0`` (mp.Pool, always Python extraction) -- + previously untested here (all prior bundle tests ran serial only). + + Runs in a **fresh subprocess**, deliberately, not in-process -- + same fork-after-native-threads hazard `test_parallel_prep.py`'s + ``TestPrepareTrainingDataParallel::test_python_backend_real_pool_matches_rust_chunk_set`` + documents for the Rust/rayon case (#275, #307), except the thread + pool here is libtorch's: by the time this test class runs in the + full suite, many earlier tests have already run real CPU inference, + so forking `mp.Pool` in-process hangs (confirmed: it does not hang + run alone or early in a session, only after prior torch CPU use -- + this actually hung PR #308's CI for ~55 minutes before being caught). + A brand new interpreter has never touched torch, so it forks safely. + """ + import subprocess + import sys + from pathlib import Path + + bundle_path = _create_pairwise_bundle(tmp_path, self.PAIR_NAMES) + out = tmp_path / "parallel.bam" + script = f""" +import pysam +from pathlib import Path +from leech.inference import run_bundle_inference + +run_bundle_inference( + bundle_path=Path({str(bundle_path)!r}), + pod5_path=Path({str(TRNA_POD5)!r}), + bam_path=Path({str(TRNA_BAM)!r}), + output_path=Path({str(out)!r}), + device="cpu", + batch_size=8, + reverse_signal=True, + reference_fasta=Path({str(TRNA_REF)!r}), + num_workers=2, + backend="python", +) +""" + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + timeout=120, + cwd=Path(__file__).parent.parent, + ) + assert result.returncode == 0, ( + f"subprocess failed (rc={result.returncode}):\n" + f"stdout={result.stdout}\nstderr={result.stderr}" + ) + with pysam.AlignmentFile(str(out), "rb") as bam: + parallel = {r.query_name: r.get_tag("aa") for r in bam if r.has_tag("aa")} + assert parallel + + def test_rust_path_matches_serial_python_path(self, tmp_path): + """The Rust and Python extraction paths must agree, same as + ``TestExtractionPathParity.test_rust_and_python_extraction_agree`` + does for single-model predict.""" + pytest.importorskip("leech_core") + bundle_path = _create_pairwise_bundle(tmp_path, self.PAIR_NAMES) + serial = self._predict(tmp_path, bundle_path, "py2", backend="python") + rust = self._predict(tmp_path, bundle_path, "rs2", backend="rust") + assert serial + assert rust == serial + + # --------------------------------------------------------------------------- # BatchAccumulator: the one batching state machine (was written three times — # multiprocessing workers and threaded Rust in single.py, and bundle.py) @@ -1539,6 +1645,119 @@ def test_worker_path_matches_sequential_path(self, tmp_path): assert parallel == sequential +@pytest.mark.skipif(not TRNA_FIXTURES_AVAILABLE, reason="tRNA fixtures not available") +class TestChunksForReadSharedBySingleAndBundle: + """``chunks_for_read`` is the one function both ``run_inference``'s and + ``run_bundle_inference``'s Python serial extraction paths call to turn a + ``LeechRead`` + a focus position into model-ready arrays (issue #268). + Before this they carried two copies, and only the bundle copy applied + dwell templates while only the bundle copy also silently skipped + ``prepare_signal_channels``'s ``signal_len`` padding/cropping on ``signal`` + itself (found while unifying, not by design). + + single.py predicts at every motif position a read has; bundle.py predicts + only at the first. The two call sites can therefore only be compared on + their first position, which is what this test does, on a real fixture + read (not a synthetic chunk) so the whole build_leech_read -> motif + search -> chunks_for_read chain is exercised end to end. + """ + + def test_single_and_bundle_style_calls_agree_on_a_real_read(self): + import numpy as np + + from leech.configs import ChunkConfig + from leech.features import extract_move_table + from leech.inference.helpers import InferenceSpec, chunks_for_read + from leech.io import get_reference_sequences + from leech.io.motif_search import get_motif_searcher + from leech.io.pod5_reader import POD5Reader + from leech.preparation.reader import build_leech_read + + config = { + "model_name": "ConvLSTMDwell", + **_ARCH_CONFIG, + "motif": "CCAGGC", + "motif_offset": 2, + } + spec = InferenceSpec.from_config(config, model_type="ConvLSTMDwell") + assert spec.requires_features, "test assumes a feature-branch model" + + reference_sequences = get_reference_sequences(TRNA_BAM, TRNA_REF) + motif_searcher = get_motif_searcher( + mode="fasta" if reference_sequences else "bam", + reference_sequences=reference_sequences, + skip_indels=spec.skip_motif_indels, + anchor=spec.anchor, + require_query_mapping=spec.require_query_mapping, + ) + signal_config = SignalConfig( + reverse_signal=True, + anchor=spec.anchor, + refine_signal_map=spec.refine_signal_map, + ) + chunk_config = ChunkConfig( + base_justify=spec.base_justify, + feature_start=spec.feature_start, + feature_end=spec.feature_end, + signal_context=spec.signal_context, + kmer_context=spec.kmer_context, + ) + + with ( + pysam.AlignmentFile(str(TRNA_BAM), "rb") as bam, + POD5Reader(TRNA_POD5) as pod5_reader, + ): + aln = next(a for a in bam if a.query_name is not None and a.query_sequence is not None) + raw_signal, pod5_metadata = pod5_reader.get_signal(aln.query_name) + ref_seq = reference_sequences.get(aln.reference_name) if reference_sequences else None + leech_read = build_leech_read( + read_id=aln.query_name, + sequence=aln.query_sequence, + raw_signal=raw_signal, + move_table=extract_move_table(aln), + signal_config=signal_config, + metadata={}, + reference_sequence=ref_seq, + cigar_tuples=aln.cigartuples, + cal_offset=pod5_metadata.get("calibration_offset"), + cal_scale=pod5_metadata.get("calibration_scale"), + ) + positions = [ + pos + spec.motif_offset + for pos in motif_searcher.find_motif_positions( + leech_read.read_id, leech_read.sequence, aln, spec.motif + ) + ] + + assert positions, "fixture read must have at least one motif match" + + common_kwargs = { + "chunk_config": chunk_config, + "signal_len": spec.signal_len, + "seq_encoding": spec.seq_encoding, + "signal_kmer_context": spec.signal_kmer_context, + "requires_features": spec.requires_features, + "kmer_len": spec.kmer_len, + "dwell_offset": spec.dwell_offset, + "wide_features": spec.wide_features, + } + + # single.py's call shape: every motif position. + single_style = chunks_for_read(leech_read, positions, **common_kwargs) + # bundle.py's call shape: only the first (one chunk per read). + bundle_style = chunks_for_read(leech_read, positions[:1], **common_kwargs) + + assert single_style[0][3] == positions[0] + assert bundle_style[0][3] == positions[0] + single_sig, single_seq, single_feat, _ = single_style[0] + bundle_sig, bundle_seq, bundle_feat, _ = bundle_style[0] + np.testing.assert_array_equal(single_sig, bundle_sig) + np.testing.assert_array_equal(single_seq, bundle_seq) + assert single_sig.shape == (spec.signal_len,) + np.testing.assert_array_equal(single_feat, bundle_feat) + assert single_feat.shape == (config["num_features"], spec.kmer_len) + + class TestSequentialThenParallelInProcess: """Two `run_inference` calls in one process must not deadlock. diff --git a/tests/test_inference_config_keys.py b/tests/test_inference_config_keys.py index b4790208..511d3b00 100644 --- a/tests/test_inference_config_keys.py +++ b/tests/test_inference_config_keys.py @@ -19,7 +19,9 @@ from leech.inference.helpers import ( InferenceConfigError, + InferenceSpec, build_rust_extraction_kwargs, + is_multiclass, prepare_inference_features, ) @@ -270,3 +272,164 @@ def test_prepare_and_inference_agree(self): reference_sequence=read_info.reference_sequence, cigar_tuples=read_info.cigar_tuples, ) + + +class TestIsMulticlass: + """One definition of "multiclass", not `> 1` in predict and `> 2` in + eval/calibrate -- issue #269.""" + + @pytest.mark.parametrize(("num_out", "expected"), [(1, False), (2, True), (3, True)]) + def test_matches_num_out_greater_than_one(self, num_out, expected): + assert is_multiclass({"num_out": num_out}) is expected + + def test_missing_num_out_defaults_to_binary(self): + assert is_multiclass({}) is False + + def test_two_output_model_is_multiclass(self): + """The exact regression #269 names: a 2-output cross-entropy model + used to get multiclass BAM tags from predict (`> 1`) but binary Platt + calibration from `eval`/`calibrate` (`> 2`). One call, one answer.""" + config = {"num_out": 2} + assert is_multiclass(config) is True + + +class TestInferenceSpecFromConfig: + """`InferenceSpec.from_config` replaces the config-resolution block that + used to be pasted into ``run_inference`` and ``run_bundle_inference`` + separately (issue #269).""" + + def _config(self, **overrides) -> dict: + base = { + "signal_len": 400, + "kmer_len": 11, + "motif": "CCAGGC", + "motif_offset": 0, + "num_out": 1, + } + base.update(overrides) + return base + + def test_resolves_geometry_from_config(self): + spec = InferenceSpec.from_config(self._config()) + assert spec.signal_len == 400 + assert spec.kmer_len == 11 + assert spec.kmer_context == 5 + assert spec.signal_context == (200, 200) + + def test_config_without_seq_encoding_defaults_to_base_onehot(self): + """A config with no `seq_encoding` key at all predates the field, or + was built directly via `get_model(name, **kwargs)` with no override -- + either way it is a base_onehot model, since that is the model + classes' own [params] default (e.g. `conv_lstm.toml`), not + `signal_kmer` (only the CLI's `model train` default, which always + records the encoding it actually used).""" + spec = InferenceSpec.from_config(self._config()) + assert spec.seq_encoding == "base_onehot" + + def test_explicit_seq_encoding_is_honored(self): + spec = InferenceSpec.from_config(self._config(seq_encoding="signal_kmer")) + assert spec.seq_encoding == "signal_kmer" + + def test_legacy_dwell_margin_corpus_resolves_feature_window(self): + """Old configs recorded `dwell_margin_left`/`dwell_margin_right` + instead of `feature_start`/`feature_end`; the resolver must still + recover a usable window rather than falling back to the k-mer window. + + A config dict never carries a "features" array, so `feature_end` + must come from the direct `kmer_context + dwell_margin_right` + arithmetic single.py/bundle.py's original inline code used, not from + the array-shape-based branch training.py's per-chunk resolution + needs -- collapsing to `None` here (issue #268 code review) silently + narrowed the window to +-kmer_context for every real checkpoint that + predates `feature_start`/`feature_end`. + """ + spec = InferenceSpec.from_config(self._config(dwell_margin_left=3, dwell_margin_right=3)) + # kmer_context=5, so feature_start = -(5+3) = -8, feature_end = 5+3 = 8 + assert spec.feature_start == -8 + assert spec.feature_end == 8 + + def test_two_output_model_is_flagged_multiclass_and_calibrated(self): + """The same regression as TestIsMulticlass, seen through the spec + every predict path now actually builds from.""" + spec = InferenceSpec.from_config(self._config(num_out=2, calibration={"pairs": {}})) + assert spec.is_multiclass is True + assert spec.calibration == {"pairs": {}} + + def test_binary_model_never_carries_calibration(self): + spec = InferenceSpec.from_config(self._config(num_out=1, calibration={"a": 1})) + assert spec.is_multiclass is False + assert spec.calibration is None + + def test_missing_motif_raises(self): + with pytest.raises(InferenceConfigError, match="motif is None"): + InferenceSpec.from_config(self._config(motif=None)) + + def test_cli_motif_offset_disagreeing_with_config_raises(self): + with pytest.raises(InferenceConfigError, match="conflicts"): + InferenceSpec.from_config(self._config(motif_offset=2), motif_offset=5, motif="CCAGGC") + + def test_explicit_cli_value_equal_to_default_is_still_checked(self): + """The bug `parameter_sources` exists to fix: an explicit + `--motif-offset 0` that disagrees with the config must raise even + though 0 is also click's own default, so the naive + `cli_value != cli_default` heuristic would wave it through.""" + with pytest.raises(InferenceConfigError, match="conflicts"): + InferenceSpec.from_config( + self._config(motif_offset=3), + motif="CCAGGC", + motif_offset=0, + parameter_sources={"motif_offset": True}, + ) + + def test_non_explicit_cli_value_equal_to_default_defers_to_config(self): + """Without a parameter source, the same call is indistinguishable + from "not passed" and the config wins -- the pre-existing, documented + heuristic for callers with no click context.""" + spec = InferenceSpec.from_config( + self._config(motif_offset=3), motif="CCAGGC", motif_offset=0 + ) + assert spec.motif_offset == 3 + + def test_wide_features_model_falls_back_to_model_dwell_margin(self): + """`dwell_margin` is a model constructor default never written into + config.json (issue #268 code review) -- the real fallback needs the + instantiated model, supplied lazily via `model_dwell_margin`.""" + spec = InferenceSpec.from_config( + self._config(), + model_type="TCNDwellResidualMotor", + model_dwell_margin=lambda: 4, + ) + assert spec.feature_start == -(5 + 4) + assert spec.feature_end == 5 + 4 + + def test_dwell_margin_in_config_alone_is_not_read(self): + """`dwell_margin` in the config dict itself must not be trusted -- + real checkpoints never write it there, only the model class does.""" + spec = InferenceSpec.from_config( + self._config(dwell_margin=4), model_type="TCNDwellResidualMotor" + ) + # No model_dwell_margin callable supplied -> falls back to 0 -> no + # override fires -> stays unresolved (callers fall back to + # +-kmer_context at the point of use, e.g. prepare_inference_features). + assert spec.feature_start is None + assert spec.feature_end is None + + def test_model_dwell_margin_is_not_called_when_not_needed(self): + """Building a model just to read one attribute is wasted work on + every call that never reaches this fallback -- must stay lazy.""" + calls = [] + InferenceSpec.from_config( + self._config(feature_start=-3, feature_end=3), + model_type="TCNDwellResidualMotor", + model_dwell_margin=lambda: calls.append(1) or 4, + ) + assert calls == [] + + def test_default_refine_scale_iters_differs_for_remora_vs_leech(self): + """single.py passed a hardcoded `2` to the Rust kwargs builder + regardless of is_remora even though its own refiner setup used `-1` + for Remora models; the spec is now the one place this is decided.""" + leech_spec = InferenceSpec.from_config(self._config(), default_refine_scale_iters=2) + remora_spec = InferenceSpec.from_config(self._config(), default_refine_scale_iters=-1) + assert leech_spec.refine_scale_iters == 2 + assert remora_spec.refine_scale_iters == -1