From bd8fc035964971cacaa82f334dbf391194ca2ee1 Mon Sep 17 00:00:00 2001 From: Alex Rubinsteyn Date: Fri, 2 Oct 2026 01:05:02 -0400 Subject: [PATCH] Bound partner recovery batches and preserve progress after timeouts --- docs/changelog.md | 7 +++ docs/partner-recovery-spec.md | 18 +++++++ docs/reads.md | 9 ++++ osteosarc/__init__.py | 2 +- osteosarc/recovery.py | 93 ++++++++++++++++++++++++----------- pyproject.toml | 2 +- tests/test_recovery.py | 67 +++++++++++++++++++++++++ 7 files changed, 167 insertions(+), 31 deletions(-) create mode 100644 docs/partner-recovery-spec.md diff --git a/docs/changelog.md b/docs/changelog.md index 2ab03dc..d4a4ff6 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -5,6 +5,13 @@ your version with `osteosarc --version`. Pin both the package version and your snapshot (by name or download date) for reproducible analyses. Full release notes are on [GitHub](https://github.com/iskandr/osteosarc/releases). +## 0.15.1 + +- Optional bounded batches for mate/SA recovery preserve available partners + after another query times out. Separate partner-query count and timeout + limits retain explicit missing/truncated outcomes (#104). Default acquisition + requests and cache identities are unchanged. + ## 0.15.0 (2026-09-29) - `subset_bundle` and `test-data make --from` carve out members offline while diff --git a/docs/partner-recovery-spec.md b/docs/partner-recovery-spec.md new file mode 100644 index 0000000..ec51741 --- /dev/null +++ b/docs/partner-recovery-spec.md @@ -0,0 +1,18 @@ +# Bounded partner-query batches (#104) + +Large regional RNA inputs can nominate thousands of mate/SA locations. A single +slow combined query currently times out before any of its available partners are +retained, even though the seed BAM is verified. + +Add optional limits for regions per partner query, total partner queries and +per-partner-query seconds. Preserve default requests/cache identities and all +source/RG/segment/placement matching. When `on_timeout="incomplete"`, retain +completed batches, label only the failed batch's leads as timed out, and continue +with other batches. Every attempted query is charged to the query budget. Failed +queries never become visited regions or evidence of absence. Subsequent calls +reuse verified completed extractions and retry missing work. + +Test late partners after an early timeout, exact record multiplicity, source +drift, exhausted budgets, input validation and offline cache reuse. Keep seed +acquisition semantics unchanged. This change supports Isovar's catalogue-wide +SV RNA audit without synthesizing observations from SAM pointers. diff --git a/docs/reads.md b/docs/reads.md index 972943e..33996d1 100644 --- a/docs/reads.md +++ b/docs/reads.md @@ -102,6 +102,15 @@ result truncated, and too many records is an error. With `RecoveryPolicy(on_timeout="incomplete")`, a query that times out gives what was found so far, marked incomplete, and the next call retries it. +For large seed regions, `partner_batch_size=256` divides partner locations into +separate indexed queries. With `on_timeout="incomplete"`, a slow batch is marked +missing while later batches continue; successfully retrieved partners survive. +`max_partner_queries=64` caps the number of attempted partner queries, including +timeouts, and `partner_timeout=30` sets their timeout in seconds independently of +seed extraction. These optional limits leave the default request strategy and +cache identities unchanged. The receipt records failed batches, visited regions +and exhausted limits; a failed query never establishes absence. + ## Genome builds ```python diff --git a/osteosarc/__init__.py b/osteosarc/__init__.py index e7f8dab..bc25745 100644 --- a/osteosarc/__init__.py +++ b/osteosarc/__init__.py @@ -57,7 +57,7 @@ from .sv_candidates import load_sv_candidates from .timeline import Event, Timeline -__version__ = "0.15.0" +__version__ = "0.15.1" # The public API. A few other helpers are importable from here because the OpenVax # libraries use them (digest, parse_variants, inspect_alignment, ...); everything diff --git a/osteosarc/recovery.py b/osteosarc/recovery.py index 4172031..b52e46e 100644 --- a/osteosarc/recovery.py +++ b/osteosarc/recovery.py @@ -27,6 +27,12 @@ class RecoveryPolicy: matched when a later partner query times out, instead of failing. The result's receipt has status "incomplete" and names the failed query; it is stored apart from complete results, so the next call retries the query. + + partner_batch_size splits each recovery round into deterministic regional + queries. With incomplete-on-timeout, later batches still run. Optional + max_partner_queries counts every partner query (including timed-out ones), + and partner_timeout overrides only partner-query seconds, not the seed. + None preserves the existing query strategy and request identity. """ mates: bool = True @@ -36,6 +42,9 @@ class RecoveryPolicy: max_bases: int = 1_000_000 max_records: int = 100_000 on_timeout: str = "fail" + partner_batch_size: int | None = None + max_partner_queries: int | None = None + partner_timeout: int | None = None def __post_init__(self): if self.on_timeout not in ("fail", "incomplete"): @@ -46,6 +55,10 @@ def __post_init__(self): for name in ("max_rounds", "max_intervals", "max_bases", "max_records"): if type(getattr(self, name)) is not int or getattr(self, name) < 1: raise ValueError(f"{name} must be a positive integer") + for name in ("partner_batch_size", "max_partner_queries", "partner_timeout"): + value = getattr(self, name) + if value is not None and (type(value) is not int or value < 1): + raise ValueError(f"{name} must be a positive integer or None") def _leads(record, policy): @@ -118,9 +131,11 @@ def recover_reads(source, regions, *, policy=None, cache=None, **kwargs): if len(resolved) > policy.max_intervals or sum(r.end - r.start for r in resolved) > policy.max_bases: raise IntegrityError("Seed acquisition exceeds recovery interval/base limit") seed = extract_reads(source, regions, cache=cache, max_records=policy.max_records, **kwargs) + # Unset new controls retain existing default request/cache identities. + recovery_parameters = {k: v for k, v in asdict(policy).items() if k != "on_timeout" and v is not None} request = dict(schema_version=2, operation="recover_reads", seed=seed.receipt["request"], seed_files=seed.receipt["files"], - recovery={k: v for k, v in asdict(policy).items() if k != "on_timeout"}) + recovery=recovery_parameters) directory = cache.workspace / "derived" / stable_id(request) with file_lock(cache.workspace / "locks" / (directory.name + ".lock")): if cached := _cached_subset(directory, request): @@ -136,6 +151,8 @@ def recover_reads(source, regions, *, policy=None, cache=None, **kwargs): kwargs.get("filters") or ReadFilter(), query_names=tuple(sorted({r.read.query_name for r in seed_records if r.read.query_name not in (None, "*")})))) + if policy.partner_timeout is not None: + partner_kwargs["timeout"] = policy.partner_timeout counts = Counter(r.digest for r in seed_records) records = {r.digest: r for r in seed_records} # An earlier query can contain a partner whose pointer is discovered later. @@ -172,7 +189,7 @@ def source_identity(receipt): receipt["request"]["reference_sha256"], receipt["request"]["reference_index_sha256"], (receipt.get("index_receipt") or {}).get("sha256")] - round_number = 1 + round_number, partner_queries = 1, 0 while True: leads = {} for key in sorted(records.keys() - handled_records): @@ -229,32 +246,48 @@ def source_identity(receipt): handled_leads.update(pending) if not intervals: continue - partner_regions = [Region(c, start, end, assembly, lengths[c]) for c, start, end in sorted(intervals)] - try: - subset = extract_reads(source, partner_regions, cache=cache, max_records=policy.max_records, - **partner_kwargs) - except (subprocess.TimeoutExpired, OsteosarcError) as error: - if policy.on_timeout != "incomplete" or not _timed_out(error): - raise - # Nothing from the failed query is kept; its leads stay unresolved. - waiting = sorted(i for i in pending if i not in unresolved) - for identity in waiting: - unresolved[identity] = "partner query timed out" - failed_queries.append(dict(round=round_number, error="timeout", - timeout_seconds=kwargs.get("timeout", 600), - regions=[asdict(r) for r in partner_regions], leads=waiting)) - break - if source_identity(subset.receipt) != source_identity(seed.receipt): - raise IntegrityError("Alignment/header/index changed across recovery acquisitions") - receipts.append(subset.receipt) - queried.update(intervals) - bases += len(intervals) - visited.extend(dict(asdict(r), round=round_number) for r in partner_regions) - fetched = list(read_records(subset.path)) - candidate_counts |= Counter(r.digest for r in fetched) - if sum(candidate_counts.values()) > policy.max_records: - raise RecordLimitError("Acquired candidates exceed recovery record limit") - add_candidates(fetched) + by_interval = defaultdict(list) + for identity, lead in pending.items(): + by_interval[lead["contig"], lead["start"], lead["start"] + 1].append(identity) + intervals.sort() + batch_size = policy.partner_batch_size or len(intervals) + timed_out = set() + for offset in range(0, len(intervals), batch_size): + batch = intervals[offset:offset + batch_size] + if policy.max_partner_queries is not None and partner_queries >= policy.max_partner_queries: + limits.add("max_partner_queries") + for interval in intervals[offset:]: + for identity in by_interval[interval]: + unresolved[identity] = "max_partner_queries" + break + partner_regions = [Region(c, start, end, assembly, lengths[c]) for c, start, end in batch] + partner_queries += 1 + try: + subset = extract_reads(source, partner_regions, cache=cache, max_records=policy.max_records, + **partner_kwargs) + except (subprocess.TimeoutExpired, OsteosarcError) as error: + if policy.on_timeout != "incomplete" or not _timed_out(error): + raise + # A timeout says nothing about another batch or absence. + waiting = sorted(i for interval in batch for i in by_interval[interval]) + timed_out.update(waiting) + for identity in waiting: + unresolved[identity] = "partner query timed out" + failed_queries.append(dict(round=round_number, error="timeout", + timeout_seconds=partner_kwargs.get("timeout", 600), + regions=[asdict(r) for r in partner_regions], leads=waiting)) + continue + if source_identity(subset.receipt) != source_identity(seed.receipt): + raise IntegrityError("Alignment/header/index changed across recovery acquisitions") + receipts.append(subset.receipt) + queried.update(batch) + bases += len(batch) + visited.extend(dict(asdict(r), round=round_number) for r in partner_regions) + fetched = list(read_records(subset.path)) + candidate_counts |= Counter(r.digest for r in fetched) + if sum(candidate_counts.values()) > policy.max_records: + raise RecordLimitError("Acquired candidates exceed recovery record limit") + add_candidates(fetched) found = Counter() found_records = {} for identity, lead in pending.items(): @@ -268,7 +301,7 @@ def source_identity(receipt): unresolved.pop(identity, None) if len(multiplicities) > 1: problems[identity] = dict(lead=identity, reason="multiple matching placements") - else: + elif identity not in timed_out: unresolved.setdefault(identity, "missing partner or conflicting alignment fields") if sum((counts | found).values()) > policy.max_records: raise RecordLimitError("Retained records exceed recovery record limit") @@ -305,6 +338,8 @@ def source_identity(receipt): observations=list(problems.values()), repeated_or_cyclic_leads=sorted(repeated), limits=sorted(limits), reasons={k: sorted(v) for k, v in sorted(reasons.items())}, pysam_version=pysam.__version__) + if policy.partner_batch_size is not None or policy.max_partner_queries is not None: + receipt["partner_queries"] = partner_queries if failed_queries: receipt["failed_queries"] = failed_queries write_json(work / "receipt.json", receipt) diff --git a/pyproject.toml b/pyproject.toml index 06af2a4..204fba0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "osteosarc" -version = "0.15.0" +version = "0.15.1" description = "Reproducible access to the public osteosarc.com dataset" readme = "README.md" requires-python = ">=3.9" diff --git a/tests/test_recovery.py b/tests/test_recovery.py index 6b7dd57..c5081ce 100644 --- a/tests/test_recovery.py +++ b/tests/test_recovery.py @@ -96,6 +96,73 @@ def test_seed_limits_checked_before_extraction(split_bam, tmp_path, monkeypatch, extract_reads(split_bam, iter(regions), cache=tmp_path / "cache", recovery=policy) +def test_partner_batches_keep_later_available_records_after_timeout(split_bam, tmp_path, monkeypatch): + import subprocess + + import osteosarc.recovery as recovery + + original = recovery.extract_reads + queried = [] + + def interrupted(source, regions, **kwargs): + regions = list(regions) + if regions[0].contig == "chr2": + assert len(regions) == 1 + assert kwargs["timeout"] == 7 + queried.append(regions[0].start) + if regions[0].start == 700: + raise subprocess.TimeoutExpired("samtools", 7) + else: + assert kwargs["timeout"] == 30 # The seed timeout is unchanged. + return original(source, regions, **kwargs) + + monkeypatch.setattr(recovery, "extract_reads", interrupted) + policy = RecoveryPolicy(partner_batch_size=1, partner_timeout=7, on_timeout="incomplete") + result = extract_reads(split_bam, [Region("chr1", 100, 150, "GRCh38")], cache=tmp_path / "cache", + recovery=policy, timeout=30) + assert queried == [500, 700, 800, 900] + assert result.receipt["status"] == "incomplete" + assert result.receipt["records"] == 7 # Both split pieces and the late mate survive. + assert not (record_multiset(result.path) - record_multiset(split_bam)) + assert result.receipt["partner_queries"] == 4 + failure, = result.receipt["failed_queries"] + assert len(failure["leads"]) == 1 and failure["timeout_seconds"] == 7 + assert not any(row["contig"] == "chr2" and row["start"] == 700 for row in result.receipt["visited_intervals"]) + assert Counter(result.receipt["unresolved"].values()) == { + "partner query timed out": 1, "missing partner or conflicting alignment fields": 1} + monkeypatch.setattr(recovery, "extract_reads", original) + resumed = extract_reads(split_bam, [Region("chr1", 100, 150, "GRCh38")], cache=tmp_path / "cache", + recovery=policy, timeout=30) + assert resumed.receipt["status"] == "bounded" + assert "failed_queries" not in resumed.receipt + assert record_multiset(resumed.path) == record_multiset(result.path) + + +def test_unset_batch_options_keep_existing_request_identity(split_bam, tmp_path): + result = extract_reads(split_bam, [Region("chr1", 100, 150, "GRCh38")], cache=tmp_path / "cache", + recovery=RecoveryPolicy()) + assert result.receipt["request"]["recovery"] == dict( + mates=True, supplementary=True, max_rounds=4, max_intervals=128, + max_bases=1_000_000, max_records=100_000) + + +def test_partner_query_budget_keeps_progress_and_reports_skipped_leads(split_bam, tmp_path): + result = extract_reads(split_bam, [Region("chr1", 100, 150, "GRCh38")], cache=tmp_path / "cache", + recovery=RecoveryPolicy(partner_batch_size=1, max_partner_queries=1)) + assert result.receipt["status"] == "truncated" + assert result.receipt["partner_queries"] == 1 + assert result.receipt["records"] == 6 + assert result.receipt["limits"] == ["max_partner_queries"] + assert set(result.receipt["unresolved"].values()) == {"max_partner_queries"} + + +@pytest.mark.parametrize("field", ["partner_batch_size", "max_partner_queries", "partner_timeout"]) +@pytest.mark.parametrize("value", [0, -1, True, 1.5, "1"]) +def test_partner_query_options_reject_invalid_limits(field, value): + with pytest.raises(ValueError, match=field): + RecoveryPolicy(**{field: value}) + + def test_later_lead_reuses_candidate_in_visited_window(split_bam, tmp_path): path = tmp_path / "recursive.bam" with pysam.AlignmentFile(split_bam) as original: