fix: unify predict's config resolution and chunk extraction (#269, #268)

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.

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.

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

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

src/leech/chunking/extractor.py

@@ -51,6 +51,119 @@ def resolve_feature_window(
     return start, end, end - start + 1
 
 
+def _assert_constant_column(column: np.ndarray, name: str) -> None:
+    """Refuse a corpus whose stored feature window is not the same on every row.
+
+    Reading column 0 (or any single row) and trusting it for the whole corpus
+    is issue #230's failure mode generalized to this field: a corpus written
+    by two prepare runs with different ``--feature-start``/``--feature-end``
+    would otherwise train silently on whichever value the first chunk happened
+    to carry.
+    """
+    if column.size and not np.all(column == column[0]):
+        raise ValueError(
+            f"Corpus has inconsistent {name!r} across chunks: {np.unique(column).tolist()}. "
+            f"A training run needs one feature window throughout; this corpus mixes prepare "
+            f"runs with different --feature-start/--feature-end and must be re-prepared "
+            f"consistently before it can be trained on."
+        )
+
+
+def _feature_window_from_table(table, kmer_context: int) -> tuple[int | None, int | None]:
+    """Column form of :func:`feature_window_from_metadata`.
+
+    ``feature_left``/``feature_right`` are never written as columns (only as
+    legacy per-chunk-dict fields on the row path), so the fallback chain here
+    is shorter than the mapping form's.
+    """
+    start_col = table.values("feature_start")
+    if start_col is not None:
+        _assert_constant_column(start_col, "feature_start")
+        start = int(start_col[0])
+    else:
+        margin_col = table.values("dwell_margin_left")
+        if margin_col is not None:
+            _assert_constant_column(margin_col, "dwell_margin_left")
+            start = -(kmer_context + int(margin_col[0]))
+        else:
+            start = None
+
+    end_col = table.values("feature_end")
+    if end_col is not None:
+        _assert_constant_column(end_col, "feature_end")
+        end = int(end_col[0])
+    else:
+        end = None
+
+    return start, end
+
+
+def _feature_window_from_mapping(source, kmer_context: int) -> tuple[int | None, int | None]:
+    """Mapping form of :func:`feature_window_from_metadata`.
+
+    ``source`` is a model config dict or a single chunk (a plain dict or a
+    :class:`~leech.chunking.table.ChunkRow`). Mirrors the fallback chain that
+    used to be pasted into ``single.py``, ``bundle.py``, ``training.py`` and
+    ``dataset.py`` separately: ``feature_start``/``feature_end`` (current),
+    then ``feature_left``/``feature_right`` (legacy), then
+    ``dwell_margin_left``/``dwell_margin_right`` (older legacy, ``right``
+    resolved from the stored feature array's width when one is present).
+    """
+    if source.get("feature_start") is not None:
+        start = int(source["feature_start"])
+    elif source.get("feature_left") is not None:
+        start = -int(source["feature_left"])
+    elif source.get("dwell_margin_left") is not None:
+        start = -(kmer_context + int(source["dwell_margin_left"]))
+    else:
+        start = None
+
+    if source.get("feature_end") is not None:
+        end = int(source["feature_end"])
+    elif source.get("feature_right") is not None:
+        end = int(source["feature_right"])
+    elif source.get("dwell_margin_right") is not None:
+        raw_features = source.get("features")
+        if raw_features is not None and getattr(raw_features, "ndim", 1) > 1 and start is not None:
+            # A chunk dict: derive from the stored feature array's actual width.
+            end = raw_features.shape[1] - 1 + start
+        else:
+            # A config dict never carries "features" -- this is the only branch
+            # that fires for it, and single.py/bundle.py's original inline code
+            # computed exactly this (kmer_context + dwell_margin_right) with no
+            # array to check. Do not collapse this to None: that silently
+            # narrows the window to +-kmer_context instead of using the margin
+            # the config actually recorded.
+            end = kmer_context + int(source["dwell_margin_right"])
+    else:
+        end = None
+
+    return start, end
+
+
+def feature_window_from_metadata(source, kmer_context: int) -> tuple[int | None, int | None]:
+    """Resolve the stored ``(feature_start, feature_end)`` from a config or corpus.
+
+    ``None`` for either element means nothing was stored for it -- callers
+    fall back to the k-mer window via :func:`resolve_feature_window`, exactly
+    as a fresh corpus with no feature-window fields at all always has.
+
+    ``source`` is either a :class:`~leech.chunking.table.ChunkTable` (the
+    window is read as a column and asserted constant across every chunk -- see
+    :func:`_assert_constant_column`) or a mapping -- a model's ``config.json``
+    dict or a single chunk (a plain dict or a
+    :class:`~leech.chunking.table.ChunkRow`). This is the one place the
+    fallback chain is written; ``training.py``, ``dataset.py`` and
+    ``InferenceSpec`` all resolve through it rather than carrying their own
+    copies (issue #269).
+    """
+    from leech.chunking.table import ChunkTable
+
+    if isinstance(source, ChunkTable):
+        return _feature_window_from_table(source, kmer_context)
+    return _feature_window_from_mapping(source, kmer_context)
+
+
 def merge_feature_channels(
     dwell_features: dict[str, np.ndarray],
     signal_features: dict[str, np.ndarray],

src/leech/cli.py

@@ -1828,13 +1828,27 @@ 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",
             FutureWarning,
             stacklevel=2,
         )
         anchor = "reference"
+        # The deprecated flag is itself an explicit request for anchor, same
+        # as --anchor reference would be -- without this, a model config
+        # recording anchor="basecall" silently wins over it instead of
+        # raising the conflict InferenceSpec.from_config exists to catch.
+        parameter_sources["anchor"] = True
 
     # Parse --copy-tags into a list
     parsed_copy_tags = [t.strip() for t in copy_tags.split(",") if t.strip()] if copy_tags else None
@@ -1846,8 +1860,6 @@ def predict(
     elif output_name.endswith(".bam"):
         output_format = "bam"
     else:
-        import rich_click as click
-
         raise click.UsageError(
             f"Cannot determine output format from extension: {output}. "
             "Use .bam for BAM output or .tsv / .tsv.gz for TSV output."
@@ -1880,6 +1892,7 @@ def predict(
         no_compile=no_compile,
         output_format=output_format,
         copy_tags=parsed_copy_tags,
+        parameter_sources=parameter_sources,
     )
 
 

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."""

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]")

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

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

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"