Skip to content

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

Merged
jayhesselberth merged 4 commits into
mainfrom
wave2-inference
Sep 13, 2026
Merged

fix: unify predict's config resolution and chunk extraction (#269, #268)#308
jayhesselberth merged 4 commits into
mainfrom
wave2-inference

Conversation

@jayhesselberth

@jayhesselberth jayhesselberth commented Sep 13, 2026

Copy link
Copy Markdown
Member

Summary

  • inference: one InferenceSpec resolver for model config, CLI overrides and the feature window #269: InferenceSpec.from_config is now the one place run_inference and run_bundle_inference resolve a model/bundle config into geometry, motif/anchor/justify, refinement settings, the feature window, and is_multiclass — replacing two ~90-line blocks that had drifted (is_multiclass was >1 in predict but >2 in eval/calibrate; seq_encoding's missing-key fallback disagreed between predict and calibration/export). chunking.feature_window_from_metadata is the shared read-side feature-window resolver, used by the spec, training.py, and dataset.py; its table form asserts a corpus's feature window is constant across every chunk rather than trusting chunk 0 (training's own chunks[0] read was exactly issue dataset: signal_kmer silently degrades to base_onehot — a warning, decided from one chunk, and the saved config still says signal_kmer #230's failure mode).
  • inference: run bundle predict through the single-model streaming pipeline #268: chunks_for_read is now the one function both predict paths' Python serial extraction uses for the per-chunk array-building triplet (signal padding/cropping, sequence encoding, feature-window narrowing) — bundle's copy had drifted (only it applied dwell templates, and it silently skipped signal_len padding/cropping on the raw signal). GpuBatchRunner extracts PR fix(predict): make torch.compile take effect, and unblock the GPU pipeline #253's double-buffered single-worker async GPU submission into a reusable primitive, now used by all three of run_bundle_inference's extraction paths via a new BundleScorer.
  • A follow-up commit addresses code-review findings: two confirmed regressions in the above (a config-dict feature-window fallback and a wide-feature-model margin fallback that read a config key no real checkpoint ever sets), plus three smaller fixes (a --reference-anchored parameter-source gap, a missing log line, a redundant alias) — see that commit message for full detail, including a self-caught corruption bug from a blind find-replace during cleanup (fixed before landing, not shipped).

Both issues' non-goals held: aggregation semantics untouched, no array value changed (verified by the full test_backend_parity.py/test_rust_refinement_parity.py/test_rust_python_parity.py suites).

Found and left out of scope for a future issue: chaining a real torch CPU-inference run_bundle_inference(num_workers=0) call with run_bundle_inference(num_workers=2) 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. Documented in tests/test_inference.py::TestBundleExtractionPathParity.

Test plan

  • Full tests/test_inference.py + tests/test_inference_config_keys.py + tests/test_chunk_table.py suites green, including new parity tests (single vs bundle chunk extraction agree on real fixture data; bundle's Rust/parallel/serial extraction paths agree with each other)
  • tests/test_evaluation.py, tests/test_calibration.py, tests/test_onnx_export.py green (is_multiclass/seq_encoding consumers)
  • tests/test_training.py, tests/test_training_advanced.py, tests/test_dataset.py, tests/test_backend_parity.py, tests/test_util.py green (feature_window_from_metadata consumers)
  • Full repo ruff check / ruff format --check clean
  • A real regression caught mid-implementation by the new TestExtractionPathParity::test_worker_path_matches_sequential_path test (a functools.partial-captured pending dict silently dropped predictions after the first mega-batch) — fixed and verified before landing
  • Two further regressions caught by /code-review high after implementation — fixed and verified before landing (see second commit)

Closes #269
Closes #268

🤖 Generated with Claude Code

https://claude.ai/code/session_01JJssLRQzDyJFFnruK7SoDu

jayhesselberth added a commit that referenced this pull request Sep 13, 2026
@jayhesselberth
fix(test): isolate run_bundle_inference's num_workers>0 case in a sub…
594c1b9 
…process

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.
@jayhesselberth
fix(inference): one InferenceSpec resolver, unify is_multiclass and s…
a10140d 
…eq_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
@jayhesselberth
fix(inference): share chunk extraction and GPU submission between sin…
f66e118 
…gle 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
@jayhesselberth
fix: address code review findings on InferenceSpec/chunks_for_read un…
9bdf2a5 
…ification

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.
@jayhesselberth
fix(test): isolate run_bundle_inference's num_workers>0 case in a sub…
0515053 
…process

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.
@jayhesselberth
jayhesselberth merged commit 8f24f42 into main Sep 13, 2026
3 checks passed
@jayhesselberth
jayhesselberth deleted the wave2-inference branch September 13, 2026 23:44
@jayhesselberth jayhesselberth mentioned this pull request Sep 14, 2026
Merged
jayhesselberth added a commit that referenced this pull request Sep 14, 2026
@jayhesselberth
ci: add job-level timeout backstop and pytest-timeout
0fadffc 
PR #308 hung ~55 minutes with nothing to stop it; add a 20-minute
job-level timeout plus a 120s per-test pytest-timeout (thread method,
since SIGALRM can't interrupt a C-level lock).
jayhesselberth added a commit that referenced this pull request Sep 14, 2026
@jayhesselberth
ci: add job-level timeout backstop and pytest-timeout (#309)
2ab253b 
* ci: add job-level timeout backstop and pytest-timeout

PR #308 hung ~55 minutes with nothing to stop it; add a 20-minute
job-level timeout plus a 120s per-test pytest-timeout (thread method,
since SIGALRM can't interrupt a C-level lock).

* chore: add changelog fragment for #309

* fix: install pytest-timeout where CI actually reads it, fix hang-test races

pytest-timeout was declared under [dependency-groups].dev, which CI's
install step never installs (it runs `uv pip install -e ".[test,rust,onnx]"`,
not `uv sync`) -- the whole per-test timeout was a silent no-op in CI.
Move it to the `test` extra instead, verified by simulating CI's exact
install command in a fresh venv.

Also mark the three existing hang-regression tests that carry their own
subprocess.run(timeout=120) (test_inference.py's
test_parallel_path_alone_writes_tags and
test_a_fork_after_a_waited_shutdown_completes, test_parallel_prep.py's
test_python_backend_real_pool_matches_rust_chunk_set) with
@pytest.mark.timeout(150), so their own internal timeout fires and reports
cleanly instead of racing the new global 120s timeout.

And distinguish "pytest finished with nothing to report" from "pytest was
killed mid-run" in the job-summary action, so a future timeout kill doesn't
get a falsely reassuring "no skipped/xfailed tests" summary.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

1 participant

@jayhesselberth