beam: scale the async-operation timeout to the corpus instead of a fixed 300s

scripts/test_bank_kwargs.py

@@ -0,0 +1,61 @@
+"""Check that the retain-side ingest controls reach bank creation.
+
+Ingest cost on a large split is dominated by fact extraction, and extraction cost is
+``corpus_chars / retain_chunk_size`` LLM calls — for BEAM-10M (~468M characters) that is ~156,000
+calls per run at the server-side default chunk size of 3000. These two env vars are the levers, so
+they are worth a check that they actually arrive rather than being silently dropped.
+
+    uv run python scripts/test_bank_kwargs.py
+"""
+import os
+
+from memory_bench.memory.hindsight import _HindsightBase
+
+
+def _kwargs(dataset: str | None, **env) -> dict:
+    prev = {k: os.environ.get(k) for k in env}
+    os.environ.update({k: v for k, v in env.items() if v is not None})
+    for k, v in env.items():
+        if v is None:
+            os.environ.pop(k, None)
+    try:
+        p = _HindsightBase.__new__(_HindsightBase)
+        p._dataset = dataset
+        return p._bank_kwargs()
+    finally:
+        for k, v in prev.items():
+            if v is None:
+                os.environ.pop(k, None)
+            else:
+                os.environ[k] = v
+
+
+def main() -> None:
+    # Default: neither control is sent, so an existing run is byte-for-byte the run it always was.
+    k = _kwargs("beam", AMB_HINDSIGHT_EXTRACTION_MODE=None, AMB_HINDSIGHT_CHUNK_SIZE=None)
+    assert "retain_extraction_mode" not in k, k
+    assert "retain_chunk_size" not in k, k
+    assert "retain_mission" in k, "BEAM must still get its extraction mission by default"
+    print("default            -> unchanged, mission present  ok")
+
+    # chunks mode skips the LLM entirely. It is a DIFFERENT measurement — the mission above is an
+    # extraction prompt and chunks mode ignores it — so this is opt-in, never a default.
+    k = _kwargs("beam", AMB_HINDSIGHT_EXTRACTION_MODE="chunks", AMB_HINDSIGHT_CHUNK_SIZE=None)
+    assert k["retain_extraction_mode"] == "chunks", k
+    print("mode=chunks        -> retain_extraction_mode=chunks  ok")
+
+    # Chunk size is the gentler lever: it keeps extraction and halves the call count per doubling.
+    k = _kwargs("beam", AMB_HINDSIGHT_EXTRACTION_MODE=None, AMB_HINDSIGHT_CHUNK_SIZE="12000")
+    assert k["retain_chunk_size"] == 12000 and isinstance(k["retain_chunk_size"], int), k
+    print("chunk_size=12000   -> retain_chunk_size=12000 (int)  ok")
+
+    # Non-BEAM datasets get the controls too, just no BEAM mission.
+    k = _kwargs("locomo", AMB_HINDSIGHT_EXTRACTION_MODE="chunks", AMB_HINDSIGHT_CHUNK_SIZE=None)
+    assert k["retain_extraction_mode"] == "chunks" and "retain_mission" not in k, k
+    print("non-beam           -> controls apply, no beam mission  ok")
+
+    print("\nall ok")
+
+
+if __name__ == "__main__":
+    main()

src/memory_bench/memory/hindsight.py

@@ -125,6 +125,28 @@ def _bank_kwargs(self, bank_id: str | None = None) -> dict:
         kwargs: dict = dict(enable_observations=False)
         if self._dataset == "beam":
             kwargs["retain_mission"] = self._BEAM_RETAIN_MISSION
+
+        # Retain-side ingest controls, both unset by default so every existing result stays
+        # comparable. They exist because ingest cost is dominated by fact extraction, and extraction
+        # cost is `corpus_chars / retain_chunk_size` LLM calls:
+        #
+        #   BEAM-10M is 10 conversations totalling ~468M characters. At the server-side default
+        #   chunk size of 3000 chars that is ~156,000 extraction calls for a single run.
+        #
+        # `AMB_HINDSIGHT_EXTRACTION_MODE=chunks` skips the LLM entirely and stores each chunk as its
+        # own unit. That makes a run of that size tractable, but it is NOT the same measurement:
+        # `_BEAM_RETAIN_MISSION` above is an extraction prompt, so chunks mode ignores it and stores
+        # raw text with no fact extraction and no entities. Use it deliberately, and do not compare
+        # a chunks-mode score against an extracted one.
+        #
+        # `AMB_HINDSIGHT_CHUNK_SIZE` trades the same axis more gently: doubling it roughly halves the
+        # call count while keeping extraction, at coarser granularity.
+        mode = os.environ.get("AMB_HINDSIGHT_EXTRACTION_MODE")
+        if mode:
+            kwargs["retain_extraction_mode"] = mode
+        chunk_size = os.environ.get("AMB_HINDSIGHT_CHUNK_SIZE")
+        if chunk_size:
+            kwargs["retain_chunk_size"] = int(chunk_size)
         return kwargs
 
     def _create_bank(self, bank_id: str, force_reset: bool = True) -> None:
@@ -136,8 +158,26 @@ def _create_bank(self, bank_id: str, force_reset: bool = True) -> None:
                 pass
         self._client.create_bank(bank_id=bank_id, name=f"Benchmark Bank ({bank_id})", **kwargs)
 
-    async def _await_operation(self, client, bank_id: str, operation_id: str, max_wait_s: int = 300) -> None:
-        """Poll until an async retain operation completes (5-minute timeout)."""
+    async def _await_operation(self, client, bank_id: str, operation_id: str, max_wait_s: int | None = None) -> None:
+        """Poll until an async retain operation completes.
+
+        The timeout has to scale with the corpus, not sit at a constant. A retain is submitted
+        asynchronously and returns an operation id immediately -- measured in-cluster, submitting a
+        20,000,000-char document takes 3.9s -- but the work behind it runs at roughly 31,000 chars/s,
+        so a BEAM-10M conversation (median 47,280,119 chars) needs ~25 MINUTES to finish.
+
+        At the old fixed 300s this method gave up five minutes in, logged a warning, and let the
+        harness query a corpus that was still loading. That is what produced the published 10m
+        result: `ingested_docs: 1`, `accuracy: 0.0`. The run did not fail, it scored an empty bank.
+
+        `AMB_OPERATION_TIMEOUT_S` overrides it; the default is two hours, which covers the largest
+        single document in any current split with room to spare.
+        """
+        if max_wait_s is None:
+            try:
+                max_wait_s = int(os.environ.get("AMB_OPERATION_TIMEOUT_S", "7200"))
+            except ValueError:
+                max_wait_s = 7200
         from hindsight_client_api.api.operations_api import OperationsApi
         ops_api = OperationsApi(client._api_client)
         waited = 0
@@ -157,9 +197,10 @@ async def _await_operation(self, client, bank_id: str, operation_id: str, max_wa
             waited += 1
         if waited >= max_wait_s:
             import logging
-            logging.getLogger(__name__).warning(
-                f"_await_operation timed out after {max_wait_s}s for bank={bank_id} op={operation_id} "
-                f"last_status={last_status!r}; continuing anyway."
+            logging.getLogger(__name__).error(
+                f"_await_operation GAVE UP after {max_wait_s}s for bank={bank_id} op={operation_id} "
+                f"last_status={last_status!r}. Ingestion is INCOMPLETE and any score from this run "
+                f"measures a partially loaded corpus, not the system. Raise AMB_OPERATION_TIMEOUT_S."
             )
 
     # ── Bank creation (async) ─────────────────────────────────────────────────