From a10140dbdd67315df09c40a16c2d014a6ca5ee61 Mon Sep 17 00:00:00 2001 From: Jay Hesselberth Date: Sun, 13 Sep 2026 14:05:57 -0600 Subject: [PATCH 1/4] fix(inference): one InferenceSpec resolver, unify is_multiclass and seq_encoding defaults run_inference and run_bundle_inference each carried their own ~90-line config-resolution block (motif/anchor/base_justify, refinement setup, the feature-window fallback chain) that had drifted apart: is_multiclass was num_out > 1 in predict but > 2 in eval/calibrate, and seq_encoding's missing-key fallback disagreed between predict/eval ("signal_kmer") and calibration/export ("base_onehot") -- base_onehot is correct, since it is the model classes' own [params] default, not the CLI's. InferenceSpec.from_config is now the one place both predict paths resolve a model config into geometry, motif/anchor/justify, refinement settings, the feature window and is_multiclass. chunking.feature_window_from_metadata is the shared read-side resolver, used by the spec, training.py and dataset.py; the table form asserts a corpus's feature window is constant across every chunk rather than trusting chunk 0 (training's own `train_dataset.chunks[0]` read was exactly that pattern, generalized from issue #230). _check_config_consistency now accepts click's real parameter source so an explicit --motif-offset 0 that disagrees with the checkpoint raises instead of losing to a "was this the default" heuristic. Closes #269 --- changelog.d/269.fixed.md | 18 ++ src/leech/calibration.py | 5 +- src/leech/chunking/__init__.py | 2 + src/leech/chunking/extractor.py | 106 +++++++++++ src/leech/cli.py | 12 +- src/leech/commands/calibrate.py | 11 +- src/leech/commands/predict.py | 4 + src/leech/constants.py | 11 ++ src/leech/dataset.py | 15 +- src/leech/evaluation.py | 10 +- src/leech/inference/bundle.py | 151 +++++---------- src/leech/inference/helpers.py | 278 +++++++++++++++++++++++++++- src/leech/inference/single.py | 264 ++++++++------------------ src/leech/model_export.py | 4 +- src/leech/onnx_export.py | 4 +- src/leech/training.py | 44 ++--- tests/test_chunk_table.py | 70 +++++++ tests/test_inference_config_keys.py | 125 +++++++++++++ 18 files changed, 800 insertions(+), 334 deletions(-) create mode 100644 changelog.d/269.fixed.md 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..2214a548 100644 --- a/src/leech/chunking/extractor.py +++ b/src/leech/chunking/extractor.py @@ -51,6 +51,112 @@ 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: + end = raw_features.shape[1] - 1 + start + else: + end = None + 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..81794ce7 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", @@ -1846,8 +1855,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 +1887,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..ce6803c3 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,7 +21,7 @@ ) from leech.inference.helpers import ( BatchAccumulator, - _check_config_consistency, + InferenceSpec, _encode_sequence_for_inference, _write_prediction_tags, build_rust_extraction_kwargs, @@ -37,7 +36,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 @@ -108,6 +106,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,97 +145,44 @@ 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 + 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, ) - 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." - ) + 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 + _kmer_context = kmer_context + _feature_start = spec.feature_start + _feature_end = spec.feature_end - base_justify = _check_config_consistency( - "base-justify", base_justify, config.get("base_justify"), "center" - ) 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 @@ -327,13 +273,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 +294,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 +309,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 +326,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 @@ -499,7 +444,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 +496,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 +522,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 +531,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, diff --git a/src/leech/inference/helpers.py b/src/leech/inference/helpers.py index 14a461a4..0ac4dd3a 100644 --- a/src/leech/inference/helpers.py +++ b/src/leech/inference/helpers.py @@ -4,7 +4,8 @@ import json import logging import threading -from collections.abc import Callable, Hashable +from collections.abc import Callable, Hashable, Mapping +from dataclasses import dataclass from pathlib import Path import numpy as np @@ -17,11 +18,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.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, 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 +35,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 +151,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 +160,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 +181,241 @@ 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, + ) -> "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. + """ + 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 = config.get("dwell_margin", 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, diff --git a/src/leech/inference/single.py b/src/leech/inference/single.py index 4a482d21..5f5d9494 100644 --- a/src/leech/inference/single.py +++ b/src/leech/inference/single.py @@ -13,11 +13,11 @@ 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, + InferenceSpec, _encode_sequence_for_inference, _run_batch, _run_batch_multiclass, @@ -35,7 +35,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 +234,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 +273,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 +303,67 @@ 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) - 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 = ( - 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()} + elif isinstance(wrapper_or_model, ModelInferenceWrapper): + model_wrapper = wrapper_or_model 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" + 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, ) - logger.info(f"base_justify: {base_justify}") - - anchor = _check_config_consistency("anchor", anchor, config.get("anchor"), "reference") - logger.info(f"anchor: {anchor}") + 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 + _kmer_context = kmer_context + 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 - 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: @@ -480,8 +374,10 @@ def run_inference( _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 +455,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 +470,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 +556,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 +565,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 +577,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 +687,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, @@ -855,7 +751,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 @@ -1023,10 +919,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, ) 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_config_keys.py b/tests/test_inference_config_keys.py index b4790208..8b10ef72 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,126 @@ 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.""" + spec = InferenceSpec.from_config(self._config(dwell_margin_left=3, dwell_margin_right=3)) + # kmer_context=5, so feature_start = -(5+3) = -8 + assert spec.feature_start == -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): + spec = InferenceSpec.from_config( + self._config(dwell_margin=4), model_type="TCNDwellResidualMotor" + ) + assert spec.feature_start == -(5 + 4) + assert spec.feature_end == 5 + 4 + + 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 From f66e118dbf345c10d82a79c2afbce33d4a09febe Mon Sep 17 00:00:00 2001 From: Jay Hesselberth Date: Sun, 13 Sep 2026 15:14:12 -0600 Subject: [PATCH 2/4] fix(inference): share chunk extraction and GPU submission between single and bundle predict run_bundle_inference's serial Python extraction path carried its own copy of the per-chunk array-building loop (signal channels, sequence encoding, feature-window narrowing) that run_inference's sequential path also had, and the two had drifted: 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. chunks_for_read in inference/helpers.py is now the one place that loop runs, used by both. GpuBatchRunner extracts the double-buffered single-worker async GPU submission pattern PR #253 introduced for run_inference's sequential path into a reusable primitive, now used by all three of run_bundle_inference's extraction paths (Rust, num_workers>0, serial) via a new BundleScorer class wrapping the vmap/sequential-wrappers/Platt-scaling logic that used to be a closure. Fixed a real bug surfaced while wiring GpuBatchRunner into run_inference's own sequential path: its score closure bound `pending` by value via functools.partial at construction, so when _finalize_mega_batch rebound the name to a fresh dict for the next mega-batch (letting the async BAM-write thread keep draining the old one safely), every mega-batch after the first silently lost its predictions instead of writing them. Fixed by using a plain closure (which looks pending up again on every call) instead of a partial; caught by the new TestExtractionPathParity::test_worker_path_matches_sequential_path regression, not by inspection. Also found and left out of scope: chaining a real torch CPU-inference run_bundle_inference(num_workers=0) call with a real run_bundle_inference(num_workers=2) call in the same process hangs, forking mp.Pool after libtorch's native CPU thread pool has started -- a pre-existing fork-safety hazard unrelated to this change (GpuBatchRunner's own thread is cleanly joined by the time the second call starts, and a second real-inference call that does not fork does not hang). Noted in tests/test_inference.py::TestBundleExtractionPathParity for a future issue. Non-goals held: aggregation semantics (leech.inference.aggregation) untouched; InferenceSpec (issue #269) untouched by this commit. Closes #268 --- changelog.d/268.fixed.md | 22 +++ src/leech/inference/bundle.py | 256 +++++++++++++++++++-------------- src/leech/inference/helpers.py | 185 +++++++++++++++++++++++- src/leech/inference/single.py | 141 ++++++------------ tests/test_inference.py | 184 ++++++++++++++++++++++++ 5 files changed, 578 insertions(+), 210 deletions(-) create mode 100644 changelog.d/268.fixed.md 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/src/leech/inference/bundle.py b/src/leech/inference/bundle.py index ce6803c3..74f64166 100644 --- a/src/leech/inference/bundle.py +++ b/src/leech/inference/bundle.py @@ -21,12 +21,13 @@ ) from leech.inference.helpers import ( BatchAccumulator, + GpuBatchRunner, InferenceSpec, - _encode_sequence_for_inference, _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, @@ -85,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, @@ -376,54 +466,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 @@ -637,6 +702,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( @@ -665,7 +731,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)..." ), ) @@ -673,6 +739,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}") @@ -785,6 +852,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( @@ -813,7 +881,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)..." ), ) @@ -821,6 +889,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}") @@ -900,68 +969,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( @@ -989,13 +1028,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 0ac4dd3a..0eb4d686 100644 --- a/src/leech/inference/helpers.py +++ b/src/leech/inference/helpers.py @@ -5,6 +5,7 @@ import logging import threading from collections.abc import Callable, Hashable, Mapping +from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass from pathlib import Path @@ -24,7 +25,7 @@ DEFAULT_REFINE_HALF_BANDWIDTH, DEFAULT_SEQ_ENCODING_FALLBACK, ) -from leech.features import encode_signal_kmer, sequence_to_int +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 @@ -751,6 +752,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], @@ -806,6 +934,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 5f5d9494..b75822ee 100644 --- a/src/leech/inference/single.py +++ b/src/leech/inference/single.py @@ -17,14 +17,15 @@ from leech.features import encode_signal_kmer_batch, extract_move_table, sequence_to_int from leech.inference.helpers import ( BatchAccumulator, + GpuBatchRunner, InferenceSpec, - _encode_sequence_for_inference, _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, @@ -698,42 +699,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, @@ -838,68 +824,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) @@ -1019,7 +958,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 @@ -1278,15 +1217,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/tests/test_inference.py b/tests/test_inference.py index ada995b4..f8336427 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -1136,6 +1136,77 @@ 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). + + Not chained after another real CPU-inference call in the same + process: doing so hangs forking `mp.Pool` after libtorch's native CPU + thread pool has started, a pre-existing fork-safety hazard unrelated + to this issue's GpuBatchRunner/BundleScorer work (confirmed by: + GpuBatchRunner's own thread is cleanly joined by the time the second + call starts -- `threading.enumerate()` shows only MainThread -- and a + second real-inference call that does *not* fork, e.g. `backend="rust"` + after a serial Python call, does not hang). Out of scope for #268; + worth its own issue if `leech.run_bundle_inference` is ever called + with `num_workers>0` more than once per process. + """ + bundle_path = _create_pairwise_bundle(tmp_path, self.PAIR_NAMES) + parallel = self._predict(tmp_path, bundle_path, "parallel", backend="python", num_workers=2) + 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 +1610,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. From 9bdf2a55aa5b25e62aee436f0da3793a6747be4d Mon Sep 17 00:00:00 2001 From: Jay Hesselberth Date: Sun, 13 Sep 2026 15:45:31 -0600 Subject: [PATCH 3/4] fix: address code review findings on InferenceSpec/chunks_for_read unification Two confirmed regressions from unifying config resolution (#269) and chunk extraction (#268), both verified independently by multiple review angles: - feature_window_from_metadata's dwell_margin_right fallback required a "features" array to compute feature_end, which a model/bundle config dict never has -- so every real checkpoint recording only dwell_margin_left/dwell_margin_right (pre-#189 format) got feature_end silently narrowed to None instead of kmer_context + dwell_margin_right, the direct arithmetic single.py/bundle.py's original inline code used. Now tries the array-shape derivation first (training.py's per-chunk case) and falls back to the direct arithmetic (the config-dict case). - InferenceSpec.from_config's wide-feature-model margin fallback read config.get("dwell_margin", 0), but dwell_margin is a model constructor default never written into config.json (confirmed by grep) -- so this fallback was silently dead for every real WIDE_FEATURE_MODELS checkpoint missing feature_start/feature_end. Now takes an optional lazy model_dwell_margin callable, with single.py reading the instantiated model's attribute and bundle.py instantiating one only when the fallback is actually reached. Also fixed: an explicit --reference-anchored (the deprecated predecessor to --anchor reference) didn't mark anchor as an explicit parameter source, so it could still silently lose to a disagreeing model config; bundle.py's rewrite dropped a "Motif from bundle config" log line single.py kept; a redundant _kmer_context alias in both files. Caught while applying these fixes: a blind find-replace while cleaning up the redundant alias corrupted signal_kmer_context and sequence_with_kmer_context (both contain "_kmer_context" as a substring) to signalkmer_context / sequence_withkmer_context in single.py and bundle.py. Found immediately by the test suite, not shipped -- full inference/config-key/onnx/calibration/evaluation suites verified green after the fix. --- src/leech/chunking/extractor.py | 9 +++++- src/leech/cli.py | 5 ++++ src/leech/inference/bundle.py | 9 +++--- src/leech/inference/helpers.py | 11 +++++++- src/leech/inference/single.py | 10 +++++-- tests/test_inference_config_keys.py | 44 +++++++++++++++++++++++++++-- 6 files changed, 76 insertions(+), 12 deletions(-) diff --git a/src/leech/chunking/extractor.py b/src/leech/chunking/extractor.py index 2214a548..772092e4 100644 --- a/src/leech/chunking/extractor.py +++ b/src/leech/chunking/extractor.py @@ -125,9 +125,16 @@ def _feature_window_from_mapping(source, kmer_context: int) -> tuple[int | None, 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: - end = None + # 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 diff --git a/src/leech/cli.py b/src/leech/cli.py index 81794ce7..0ffb34ed 100644 --- a/src/leech/cli.py +++ b/src/leech/cli.py @@ -1844,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 diff --git a/src/leech/inference/bundle.py b/src/leech/inference/bundle.py index 74f64166..5c3270a8 100644 --- a/src/leech/inference/bundle.py +++ b/src/leech/inference/bundle.py @@ -247,6 +247,7 @@ def run_bundle_inference( 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 @@ -261,10 +262,10 @@ def run_bundle_inference( signal_context = spec.signal_context kmer_context = spec.kmer_context wide_features = spec.wide_features - _kmer_context = kmer_context _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}") logger.info(f"anchor: {anchor}") if reference_fasta is not None: @@ -275,8 +276,8 @@ def run_bundle_inference( ) 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})") @@ -831,7 +832,7 @@ def _single_forward(params, buffers, signal, sequence, features): 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, ) diff --git a/src/leech/inference/helpers.py b/src/leech/inference/helpers.py index 0eb4d686..ecd0f94a 100644 --- a/src/leech/inference/helpers.py +++ b/src/leech/inference/helpers.py @@ -248,6 +248,7 @@ def from_config( 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. @@ -288,6 +289,14 @@ def from_config( 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 {} @@ -359,7 +368,7 @@ def from_config( 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 = config.get("dwell_margin", 0) + 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 diff --git a/src/leech/inference/single.py b/src/leech/inference/single.py index b75822ee..a2598a56 100644 --- a/src/leech/inference/single.py +++ b/src/leech/inference/single.py @@ -321,6 +321,11 @@ def run_inference( 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 + ), ) signal_len = spec.signal_len kmer_len = spec.kmer_len @@ -333,7 +338,6 @@ def run_inference( wide_features = spec.wide_features _feature_start = spec.feature_start _feature_end = spec.feature_end - _kmer_context = kmer_context num_out = spec.num_out is_multiclass = spec.is_multiclass motif = spec.motif @@ -372,8 +376,8 @@ 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})") logger.info(f"Motif: {motif} (offset={motif_offset})") logger.info(f"anchor: {anchor}") diff --git a/tests/test_inference_config_keys.py b/tests/test_inference_config_keys.py index 8b10ef72..511d3b00 100644 --- a/tests/test_inference_config_keys.py +++ b/tests/test_inference_config_keys.py @@ -333,10 +333,20 @@ def test_explicit_seq_encoding_is_honored(self): 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.""" + 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 + # 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 @@ -381,12 +391,40 @@ def test_non_explicit_cli_value_equal_to_default_defers_to_config(self): 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(dwell_margin=4), model_type="TCNDwellResidualMotor" + 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` From 0515053b0fdda5c18ab671168f76c986be9db593 Mon Sep 17 00:00:00 2001 From: Jay Hesselberth Date: Sun, 13 Sep 2026 16:45:46 -0600 Subject: [PATCH 4/4] fix(test): isolate run_bundle_inference's num_workers>0 case in a subprocess test_parallel_path_alone_writes_tags called run_bundle_inference(num_workers=2) (an mp.Pool fork) in-process. It passed running alone or early in a session, but hung PR #308's CI for ~55 minutes: by the point pytest's default full-suite collection order reaches this test, many earlier tests have already run real CPU torch inference, and forking after libtorch's native thread pool has started is the same class of hazard test_parallel_prep.py's test_python_backend_real_pool_matches_rust_chunk_set already documents and works around for escapepod's rayon pool (#275, #307) -- same fix, same reason: a brand new interpreter has never touched torch, so it forks safely. --- tests/test_inference.py | 57 +++++++++++++++++++++++++++++++++-------- 1 file changed, 46 insertions(+), 11 deletions(-) diff --git a/tests/test_inference.py b/tests/test_inference.py index f8336427..fea11343 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -1180,19 +1180,54 @@ 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). - Not chained after another real CPU-inference call in the same - process: doing so hangs forking `mp.Pool` after libtorch's native CPU - thread pool has started, a pre-existing fork-safety hazard unrelated - to this issue's GpuBatchRunner/BundleScorer work (confirmed by: - GpuBatchRunner's own thread is cleanly joined by the time the second - call starts -- `threading.enumerate()` shows only MainThread -- and a - second real-inference call that does *not* fork, e.g. `backend="rust"` - after a serial Python call, does not hang). Out of scope for #268; - worth its own issue if `leech.run_bundle_inference` is ever called - with `num_workers>0` more than once per process. + 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) - parallel = self._predict(tmp_path, bundle_path, "parallel", backend="python", num_workers=2) + 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):