Preserve partner-recovery progress around slow queries

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

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.

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

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

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)

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"

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: