diff --git a/scripts/test_bank_kwargs.py b/scripts/test_bank_kwargs.py new file mode 100644 index 0000000..8c443fe --- /dev/null +++ b/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() diff --git a/src/memory_bench/memory/hindsight.py b/src/memory_bench/memory/hindsight.py index 1e010c3..0503fa1 100644 --- a/src/memory_bench/memory/hindsight.py +++ b/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) ─────────────────────────────────────────────────