Compare commits
5
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6f004bda10 | ||
|
|
06078baf0a | ||
|
|
aff569cb87 | ||
|
|
774d6d15f5 | ||
|
|
886701a899 |
+61
@@ -0,0 +1,61 @@
|
||||
"""Add bank_stats_cache table for distributed get_bank_stats caching
|
||||
|
||||
Revision ID: b57a7c9e0d13
|
||||
Revises: c3f7a1b9d2e4
|
||||
Create Date: 2026-07-01
|
||||
|
||||
get_bank_stats aggregates over memory_links / unit_entities — a multi-second scan
|
||||
on banks with millions of rows. The result was cached per-process (in-memory), so
|
||||
every API worker recomputed it once per TTL and the first caller after expiry
|
||||
stalled. This table backs a shared, cross-process TTL cache: one worker's compute
|
||||
is written here and served to all the others.
|
||||
|
||||
PostgreSQL only. Oracle keeps the in-process cache (the runtime picks the backing
|
||||
store by dialect), so the Oracle upgrade slot is intentionally absent.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "b57a7c9e0d13"
|
||||
down_revision: str | Sequence[str] | None = "c3f7a1b9d2e4"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# One row per bank: payload is the full get_bank_stats result, computed_at
|
||||
# drives logical TTL expiry. Rows are overwritten in place (ON CONFLICT), so
|
||||
# the table never grows beyond the number of banks and needs no purge job.
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}bank_stats_cache (
|
||||
bank_id TEXT PRIMARY KEY,
|
||||
payload JSONB NOT NULL,
|
||||
computed_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}bank_stats_cache")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade) # oracle slot intentionally absent → no-op
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
@@ -4191,11 +4191,17 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
async def api_stats(
|
||||
bank_id: str,
|
||||
refresh: bool = Query(
|
||||
default=False,
|
||||
description="Force a fresh recompute, bypassing the cached value (and refreshing the cache).",
|
||||
),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Get statistics about memory nodes and links for a memory bank."""
|
||||
try:
|
||||
stats = await app.state.memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
stats = await app.state.memory.get_bank_stats(
|
||||
bank_id, request_context=request_context, force_refresh=refresh
|
||||
)
|
||||
nodes_by_type = stats["node_counts"]
|
||||
links_by_type = stats["link_counts"]
|
||||
links_by_fact_type = stats["link_counts_by_fact_type"]
|
||||
|
||||
@@ -13,9 +13,18 @@ in-flight task so that N concurrent callers produce one query rather than N.
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from typing import Any, Awaitable, Callable
|
||||
from typing import TYPE_CHECKING, Any, Awaitable, Callable
|
||||
|
||||
from .db_utils import acquire_with_retry
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .db.base import DatabaseBackend
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BankStatsCache:
|
||||
@@ -66,17 +75,28 @@ class BankStatsCache:
|
||||
schema: str,
|
||||
bank_id: str,
|
||||
loader: Callable[[], Awaitable[dict[str, Any]]],
|
||||
*,
|
||||
force_refresh: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""Return cached stats for `(schema, bank_id)` or call `loader()`.
|
||||
|
||||
Concurrent misses on the same key are coalesced onto a single
|
||||
in-flight loader.
|
||||
in-flight loader. When ``force_refresh`` is set the cached value is
|
||||
ignored: the loader runs and its result replaces the cached entry.
|
||||
"""
|
||||
if not self.enabled:
|
||||
return await loader()
|
||||
|
||||
key = (schema, bank_id)
|
||||
|
||||
if force_refresh:
|
||||
value = await loader()
|
||||
async with self._lock:
|
||||
self._store_unlocked(key, value)
|
||||
# Supersede any loader that was in flight for this key.
|
||||
self._in_flight.pop(key, None)
|
||||
return value
|
||||
|
||||
async with self._lock:
|
||||
cached = self._get_fresh_unlocked(key)
|
||||
if cached is not None:
|
||||
@@ -132,3 +152,103 @@ class BankStatsCache:
|
||||
async with self._lock:
|
||||
self._entries.clear()
|
||||
self._in_flight.clear()
|
||||
|
||||
|
||||
class DistributedBankStatsCache:
|
||||
"""Table-backed (cross-process) TTL cache for `get_bank_stats`.
|
||||
|
||||
Same ``get_or_load`` / ``invalidate`` / ``clear`` contract as
|
||||
:class:`BankStatsCache`, but the store is the per-schema ``bank_stats_cache``
|
||||
table instead of a per-process dict — so one worker's computation is shared
|
||||
with every other worker, and no caller recomputes while a fresh row exists.
|
||||
|
||||
On a hit, a call is a single primary-key ``SELECT`` (sub-millisecond); only a
|
||||
miss runs the (expensive) ``loader`` and writes the row back. Concurrent
|
||||
misses are *not* coalesced across processes (that would need a lock): they
|
||||
each compute and ``UPSERT``, last write wins — all results are correct, at the
|
||||
cost of a brief redundant compute at expiry.
|
||||
|
||||
Every DB touch is best-effort: if the cache table is unreachable or missing
|
||||
(e.g. a schema mid-migration), the call degrades to computing without caching
|
||||
rather than failing ``get_bank_stats``. PostgreSQL only — the engine keeps the
|
||||
in-process :class:`BankStatsCache` for Oracle.
|
||||
"""
|
||||
|
||||
def __init__(self, *, backend: "DatabaseBackend", ttl_seconds: float) -> None:
|
||||
self._backend = backend
|
||||
self._ttl = float(ttl_seconds)
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return self._ttl > 0
|
||||
|
||||
@staticmethod
|
||||
def _qualified(schema: str) -> str:
|
||||
return f'"{schema}".bank_stats_cache' if schema else "bank_stats_cache"
|
||||
|
||||
async def get_or_load(
|
||||
self,
|
||||
schema: str,
|
||||
bank_id: str,
|
||||
loader: Callable[[], Awaitable[dict[str, Any]]],
|
||||
*,
|
||||
force_refresh: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
if not self.enabled:
|
||||
return await loader()
|
||||
|
||||
table = self._qualified(schema)
|
||||
|
||||
# 1. Fresh row? Single PK lookup; ``payload::text`` sidesteps any
|
||||
# jsonb->object codec so we always decode the same way. Skipped when
|
||||
# the caller forces a refresh — then we recompute and overwrite below.
|
||||
if not force_refresh:
|
||||
try:
|
||||
async with acquire_with_retry(self._backend) as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"SELECT payload::text AS payload FROM {table} "
|
||||
f"WHERE bank_id = $1 AND computed_at > now() - make_interval(secs => $2::double precision)",
|
||||
bank_id,
|
||||
self._ttl,
|
||||
)
|
||||
if row is not None:
|
||||
return json.loads(row["payload"])
|
||||
except Exception as exc: # noqa: BLE001 — cache read must never break the endpoint
|
||||
logger.debug("bank_stats_cache read failed for %s.%s (%s); computing uncached", schema, bank_id, exc)
|
||||
return await loader()
|
||||
|
||||
# 2. Miss — compute, then write the row back (best-effort).
|
||||
value = await loader()
|
||||
try:
|
||||
async with acquire_with_retry(self._backend) as conn:
|
||||
await conn.execute(
|
||||
f"INSERT INTO {table} (bank_id, payload, computed_at) VALUES ($1, $2::jsonb, now()) "
|
||||
f"ON CONFLICT (bank_id) DO UPDATE SET payload = EXCLUDED.payload, computed_at = now()",
|
||||
bank_id,
|
||||
json.dumps(value),
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 — a failed write just means no caching this round
|
||||
logger.warning("bank_stats_cache write failed for %s.%s (%s)", schema, bank_id, exc)
|
||||
return value
|
||||
|
||||
async def invalidate(self, schema: str, bank_id: str) -> None:
|
||||
"""Drop the cached row so the next read recomputes."""
|
||||
if not self.enabled:
|
||||
return
|
||||
try:
|
||||
async with acquire_with_retry(self._backend) as conn:
|
||||
await conn.execute(f"DELETE FROM {self._qualified(schema)} WHERE bank_id = $1", bank_id)
|
||||
except Exception as exc: # noqa: BLE001 — invalidation must never break the write path
|
||||
logger.debug("bank_stats_cache invalidate failed for %s.%s (%s)", schema, bank_id, exc)
|
||||
|
||||
async def clear(self) -> None:
|
||||
"""Drop all cached rows in the current schema (best-effort)."""
|
||||
if not self.enabled:
|
||||
return
|
||||
from .memory_engine import get_current_schema
|
||||
|
||||
try:
|
||||
async with acquire_with_retry(self._backend) as conn:
|
||||
await conn.execute(f"DELETE FROM {self._qualified(get_current_schema())}")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("bank_stats_cache clear failed (%s)", exc)
|
||||
|
||||
@@ -449,6 +449,7 @@ class MemoryEngineInterface(ABC):
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
force_refresh: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Get statistics about memory nodes and links for a bank.
|
||||
@@ -456,6 +457,8 @@ class MemoryEngineInterface(ABC):
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
request_context: Request context for authentication.
|
||||
force_refresh: Bypass the cached value and recompute (also refreshes
|
||||
the cache for subsequent callers).
|
||||
|
||||
Returns:
|
||||
Dict with node_counts, link_counts, link_counts_by_fact_type
|
||||
|
||||
@@ -44,7 +44,7 @@ from ..utils import mask_network_location
|
||||
from ..worker.exceptions import DeferOperation, RetryTaskAt
|
||||
from ..worker.stage import set_stage
|
||||
from .audit import AuditLogger, audit_context
|
||||
from .bank_stats_cache import BankStatsCache
|
||||
from .bank_stats_cache import BankStatsCache, DistributedBankStatsCache
|
||||
from .db import DatabaseBackend, create_database_backend
|
||||
from .db_budget import budgeted_operation
|
||||
from .llm_interface import ProviderRateLimitResetError
|
||||
@@ -1322,14 +1322,21 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
regex_defense.set_context(self._ext_ctx)
|
||||
self._memory_defense = regex_defense
|
||||
|
||||
# Cache for get_bank_stats — short TTL + concurrent-loader coalescing.
|
||||
# The query joins memory_links to memory_units and can be a multi-second
|
||||
# parallel scan on large banks; a single polling client used to be able
|
||||
# to pin the primary by issuing several concurrent calls.
|
||||
self._bank_stats_cache = BankStatsCache(
|
||||
ttl_seconds=config.bank_stats_cache_ttl_seconds,
|
||||
max_entries=config.bank_stats_cache_max_entries,
|
||||
)
|
||||
# Cache for get_bank_stats — the query aggregates over memory_links /
|
||||
# unit_entities and can be a multi-second scan on large banks. On
|
||||
# PostgreSQL we back it with the shared bank_stats_cache table so one
|
||||
# worker's computation serves all workers (and survives restarts);
|
||||
# Oracle keeps the per-process in-memory cache.
|
||||
if self._database_backend_type == "postgresql":
|
||||
self._bank_stats_cache: BankStatsCache | DistributedBankStatsCache = DistributedBankStatsCache(
|
||||
backend=self._backend,
|
||||
ttl_seconds=config.bank_stats_cache_ttl_seconds,
|
||||
)
|
||||
else:
|
||||
self._bank_stats_cache = BankStatsCache(
|
||||
ttl_seconds=config.bank_stats_cache_ttl_seconds,
|
||||
max_entries=config.bank_stats_cache_max_entries,
|
||||
)
|
||||
|
||||
@property
|
||||
def audit_logger(self) -> AuditLogger:
|
||||
@@ -9739,13 +9746,15 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
force_refresh: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""Get statistics about memory nodes and links for a bank.
|
||||
|
||||
Results are served from a short-TTL per-process cache so a polling
|
||||
client cannot drive the link/unit aggregations multiple times per
|
||||
second; concurrent misses on the same bank are coalesced onto a
|
||||
single in-flight loader.
|
||||
Results are served from a short-TTL cache (a shared table on PostgreSQL,
|
||||
per-process on Oracle) so a polling client cannot drive the link/unit
|
||||
aggregations multiple times per second. Pass ``force_refresh=True`` to
|
||||
bypass the cached value and recompute (the fresh result also refreshes
|
||||
the cache for subsequent callers).
|
||||
"""
|
||||
await self._authenticate_tenant(request_context)
|
||||
if self._operation_validator:
|
||||
@@ -9761,6 +9770,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
schema,
|
||||
bank_id,
|
||||
lambda: self._compute_bank_stats(bank_id),
|
||||
force_refresh=force_refresh,
|
||||
)
|
||||
|
||||
async def _compute_bank_stats(self, bank_id: str) -> dict[str, Any]:
|
||||
|
||||
@@ -88,7 +88,10 @@ async def test_backup_tables_covers_entire_schema(backup_test_schema):
|
||||
await conn.close()
|
||||
|
||||
# alembic_version is migration bookkeeping, not data — never backed up.
|
||||
schema_tables = {r["table_name"] for r in rows} - {"alembic_version"}
|
||||
# bank_stats_cache is a derived TTL cache of get_bank_stats results: it has no
|
||||
# FK to banks (so the restore cascade never touches it) and repopulates itself
|
||||
# on demand, so it is deliberately not backed up — a restore starts it cold.
|
||||
schema_tables = {r["table_name"] for r in rows} - {"alembic_version", "bank_stats_cache"}
|
||||
backup_tables = set(BACKUP_TABLES)
|
||||
|
||||
missing = schema_tables - backup_tables
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
"""Tests for the table-backed (cross-process) get_bank_stats cache.
|
||||
|
||||
On PostgreSQL the engine backs `get_bank_stats` with the `bank_stats_cache`
|
||||
table (`DistributedBankStatsCache`) instead of a per-process dict, so one
|
||||
worker's computation is shared with every other worker. These tests verify:
|
||||
|
||||
* the PG engine actually selects the distributed cache,
|
||||
* a computed result is written to the table and served from it on the next call,
|
||||
* invalidation deletes the row so the next call recomputes, and
|
||||
* an unreachable cache table degrades to computing without caching rather than
|
||||
failing the endpoint.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.engine.bank_stats_cache import DistributedBankStatsCache
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine, get_current_schema
|
||||
|
||||
_PINNED_TTL_SECONDS = 300.0
|
||||
|
||||
|
||||
async def _insert_memory(conn, bank_id: str, text: str, fact_type: str = "experience") -> uuid.UUID:
|
||||
mem_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO memory_units (id, bank_id, text, fact_type, event_date, created_at, updated_at, consolidated_at)
|
||||
VALUES ($1, $2, $3, $4, NOW(), NOW(), NOW(), NOW())
|
||||
""",
|
||||
mem_id,
|
||||
bank_id,
|
||||
text,
|
||||
fact_type,
|
||||
)
|
||||
return mem_id
|
||||
|
||||
|
||||
async def _ensure_bank(memory: MemoryEngine, bank_id: str, request_context: RequestContext) -> None:
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
|
||||
def _pin_distributed_cache(memory: MemoryEngine) -> DistributedBankStatsCache:
|
||||
cache = DistributedBankStatsCache(backend=memory._backend, ttl_seconds=_PINNED_TTL_SECONDS)
|
||||
memory._bank_stats_cache = cache
|
||||
return cache
|
||||
|
||||
|
||||
class TestDistributedBankStatsCache:
|
||||
@pytest.mark.asyncio
|
||||
async def test_pg_engine_selects_distributed_cache(self, memory: MemoryEngine):
|
||||
if memory._database_backend_type != "postgresql":
|
||||
pytest.skip("distributed cache is PostgreSQL-only")
|
||||
assert isinstance(memory._bank_stats_cache, DistributedBankStatsCache)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_result_is_written_and_served_from_table(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
if memory._database_backend_type != "postgresql":
|
||||
pytest.skip("distributed cache is PostgreSQL-only")
|
||||
|
||||
bank_id = f"test-dist-stats-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
await _insert_memory(conn, bank_id, "Alice loves hiking.")
|
||||
|
||||
_pin_distributed_cache(memory)
|
||||
try:
|
||||
first = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert first["node_counts"].get("experience") == 1
|
||||
|
||||
# The computed result was persisted to the shared table.
|
||||
async with pool.acquire() as conn:
|
||||
rows = await conn.fetchval("SELECT count(*) FROM bank_stats_cache WHERE bank_id = $1", bank_id)
|
||||
assert rows == 1
|
||||
|
||||
# Mutate the underlying data WITHOUT going through an invalidating
|
||||
# engine method — the long-TTL cache must serve the stale row.
|
||||
async with pool.acquire() as conn:
|
||||
await _insert_memory(conn, bank_id, "Bob enjoys cycling.")
|
||||
served = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert served["node_counts"].get("experience") == 1 # still cached
|
||||
|
||||
# Invalidating drops the row → next call recomputes the true count.
|
||||
await memory._bank_stats_cache.invalidate(get_current_schema(), bank_id)
|
||||
async with pool.acquire() as conn:
|
||||
rows = await conn.fetchval("SELECT count(*) FROM bank_stats_cache WHERE bank_id = $1", bank_id)
|
||||
assert rows == 0
|
||||
fresh = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert fresh["node_counts"].get("experience") == 2
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_force_refresh_bypasses_and_updates_cache(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
if memory._database_backend_type != "postgresql":
|
||||
pytest.skip("distributed cache is PostgreSQL-only")
|
||||
|
||||
bank_id = f"test-dist-stats-fresh-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
await _insert_memory(conn, bank_id, "Alice loves hiking.")
|
||||
|
||||
_pin_distributed_cache(memory)
|
||||
try:
|
||||
# Warm the cache, then mutate the data without invalidation.
|
||||
assert (await memory.get_bank_stats(bank_id, request_context=request_context))["node_counts"][
|
||||
"experience"
|
||||
] == 1
|
||||
async with pool.acquire() as conn:
|
||||
await _insert_memory(conn, bank_id, "Bob enjoys cycling.")
|
||||
|
||||
# A normal read is served the stale cached count...
|
||||
stale = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert stale["node_counts"]["experience"] == 1
|
||||
# ...but force_refresh recomputes the true count.
|
||||
fresh = await memory.get_bank_stats(bank_id, request_context=request_context, force_refresh=True)
|
||||
assert fresh["node_counts"]["experience"] == 2
|
||||
# The forced result also refreshed the cache for the next caller.
|
||||
served = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert served["node_counts"]["experience"] == 2
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_degrades_when_cache_table_unreachable(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
if memory._database_backend_type != "postgresql":
|
||||
pytest.skip("distributed cache is PostgreSQL-only")
|
||||
|
||||
bank_id = f"test-dist-stats-degrade-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
await _insert_memory(conn, bank_id, "Alice loves hiking.")
|
||||
|
||||
# Point the cache at a table that does not exist: reads and writes fail,
|
||||
# so it must fall back to computing the real result (no real table touched).
|
||||
cache = _pin_distributed_cache(memory)
|
||||
cache._qualified = lambda schema: '"public".bank_stats_cache_does_not_exist' # type: ignore[method-assign]
|
||||
try:
|
||||
stats = await memory.get_bank_stats(bank_id, request_context=request_context)
|
||||
assert stats["node_counts"].get("experience") == 1
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -165,6 +165,12 @@ async def test_full_api_workflow(api_client, test_bank_id):
|
||||
assert "total_nodes" in stats
|
||||
assert stats["total_nodes"] > 0
|
||||
|
||||
# ?refresh=true forces a fresh recompute, bypassing the cache; same shape.
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/stats?refresh=true")
|
||||
assert response.status_code == 200
|
||||
fresh_stats = response.json()
|
||||
assert fresh_stats["total_nodes"] == stats["total_nodes"]
|
||||
|
||||
# Verify bank list returns stats (fact_count, last_document_at)
|
||||
response = await api_client.get("/v1/default/banks")
|
||||
assert response.status_code == 200
|
||||
|
||||
@@ -153,7 +153,9 @@ impl ApiClient {
|
||||
|
||||
pub fn get_stats(&self, agent_id: &str, _verbose: bool) -> Result<AgentStats> {
|
||||
self.runtime.block_on(async {
|
||||
let response = self.client.get_agent_stats(agent_id, None).await?;
|
||||
// Third arg is the `refresh` query param (force fresh stats); the CLI
|
||||
// always reads the cached value, so pass None.
|
||||
let response = self.client.get_agent_stats(agent_id, None, None).await?;
|
||||
let value = response.into_inner();
|
||||
// Convert to JSON Value first, then parse into our type
|
||||
let json_value = serde_json::to_value(&value)?;
|
||||
|
||||
@@ -583,6 +583,19 @@ paths:
|
||||
title: Bank Id
|
||||
type: string
|
||||
style: simple
|
||||
- description: "Force a fresh recompute, bypassing the cached value (and refreshing\
|
||||
\ the cache)."
|
||||
explode: true
|
||||
in: query
|
||||
name: refresh
|
||||
required: false
|
||||
schema:
|
||||
default: false
|
||||
description: "Force a fresh recompute, bypassing the cached value (and refreshing\
|
||||
\ the cache)."
|
||||
title: Refresh
|
||||
type: boolean
|
||||
style: form
|
||||
- explode: false
|
||||
in: header
|
||||
name: authorization
|
||||
|
||||
@@ -540,9 +540,16 @@ type ApiGetAgentStatsRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *BanksAPIService
|
||||
bankId string
|
||||
refresh *bool
|
||||
authorization *string
|
||||
}
|
||||
|
||||
// Force a fresh recompute, bypassing the cached value (and refreshing the cache).
|
||||
func (r ApiGetAgentStatsRequest) Refresh(refresh bool) ApiGetAgentStatsRequest {
|
||||
r.refresh = &refresh
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiGetAgentStatsRequest) Authorization(authorization string) ApiGetAgentStatsRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
@@ -591,6 +598,12 @@ func (a *BanksAPIService) GetAgentStatsExecute(r ApiGetAgentStatsRequest) (*Bank
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
|
||||
if r.refresh != nil {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "refresh", r.refresh, "form", "")
|
||||
} else {
|
||||
var defaultValue bool = false
|
||||
r.refresh = &defaultValue
|
||||
}
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{}
|
||||
|
||||
|
||||
@@ -16,7 +16,7 @@ from pydantic import validate_call, Field, StrictFloat, StrictStr, StrictInt
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
from typing_extensions import Annotated
|
||||
|
||||
from pydantic import Field, StrictStr
|
||||
from pydantic import Field, StrictBool, StrictStr
|
||||
from typing import Optional
|
||||
from typing_extensions import Annotated
|
||||
from hindsight_client_api.models.add_background_request import AddBackgroundRequest
|
||||
@@ -1228,6 +1228,7 @@ class BanksApi:
|
||||
async def get_agent_stats(
|
||||
self,
|
||||
bank_id: StrictStr,
|
||||
refresh: Annotated[Optional[StrictBool], Field(description="Force a fresh recompute, bypassing the cached value (and refreshing the cache).")] = None,
|
||||
authorization: Optional[StrictStr] = None,
|
||||
_request_timeout: Union[
|
||||
None,
|
||||
@@ -1248,6 +1249,8 @@ class BanksApi:
|
||||
|
||||
:param bank_id: (required)
|
||||
:type bank_id: str
|
||||
:param refresh: Force a fresh recompute, bypassing the cached value (and refreshing the cache).
|
||||
:type refresh: bool
|
||||
:param authorization:
|
||||
:type authorization: str
|
||||
:param _request_timeout: timeout setting for this request. If one
|
||||
@@ -1274,6 +1277,7 @@ class BanksApi:
|
||||
|
||||
_param = self._get_agent_stats_serialize(
|
||||
bank_id=bank_id,
|
||||
refresh=refresh,
|
||||
authorization=authorization,
|
||||
_request_auth=_request_auth,
|
||||
_content_type=_content_type,
|
||||
@@ -1300,6 +1304,7 @@ class BanksApi:
|
||||
async def get_agent_stats_with_http_info(
|
||||
self,
|
||||
bank_id: StrictStr,
|
||||
refresh: Annotated[Optional[StrictBool], Field(description="Force a fresh recompute, bypassing the cached value (and refreshing the cache).")] = None,
|
||||
authorization: Optional[StrictStr] = None,
|
||||
_request_timeout: Union[
|
||||
None,
|
||||
@@ -1320,6 +1325,8 @@ class BanksApi:
|
||||
|
||||
:param bank_id: (required)
|
||||
:type bank_id: str
|
||||
:param refresh: Force a fresh recompute, bypassing the cached value (and refreshing the cache).
|
||||
:type refresh: bool
|
||||
:param authorization:
|
||||
:type authorization: str
|
||||
:param _request_timeout: timeout setting for this request. If one
|
||||
@@ -1346,6 +1353,7 @@ class BanksApi:
|
||||
|
||||
_param = self._get_agent_stats_serialize(
|
||||
bank_id=bank_id,
|
||||
refresh=refresh,
|
||||
authorization=authorization,
|
||||
_request_auth=_request_auth,
|
||||
_content_type=_content_type,
|
||||
@@ -1372,6 +1380,7 @@ class BanksApi:
|
||||
async def get_agent_stats_without_preload_content(
|
||||
self,
|
||||
bank_id: StrictStr,
|
||||
refresh: Annotated[Optional[StrictBool], Field(description="Force a fresh recompute, bypassing the cached value (and refreshing the cache).")] = None,
|
||||
authorization: Optional[StrictStr] = None,
|
||||
_request_timeout: Union[
|
||||
None,
|
||||
@@ -1392,6 +1401,8 @@ class BanksApi:
|
||||
|
||||
:param bank_id: (required)
|
||||
:type bank_id: str
|
||||
:param refresh: Force a fresh recompute, bypassing the cached value (and refreshing the cache).
|
||||
:type refresh: bool
|
||||
:param authorization:
|
||||
:type authorization: str
|
||||
:param _request_timeout: timeout setting for this request. If one
|
||||
@@ -1418,6 +1429,7 @@ class BanksApi:
|
||||
|
||||
_param = self._get_agent_stats_serialize(
|
||||
bank_id=bank_id,
|
||||
refresh=refresh,
|
||||
authorization=authorization,
|
||||
_request_auth=_request_auth,
|
||||
_content_type=_content_type,
|
||||
@@ -1439,6 +1451,7 @@ class BanksApi:
|
||||
def _get_agent_stats_serialize(
|
||||
self,
|
||||
bank_id,
|
||||
refresh,
|
||||
authorization,
|
||||
_request_auth,
|
||||
_content_type,
|
||||
@@ -1464,6 +1477,10 @@ class BanksApi:
|
||||
if bank_id is not None:
|
||||
_path_params['bank_id'] = bank_id
|
||||
# process the query parameters
|
||||
if refresh is not None:
|
||||
|
||||
_query_params.append(('refresh', refresh))
|
||||
|
||||
# process the header parameters
|
||||
if authorization is not None:
|
||||
_header_params['authorization'] = authorization
|
||||
|
||||
@@ -4631,7 +4631,14 @@ export type GetAgentStatsData = {
|
||||
*/
|
||||
bank_id: string;
|
||||
};
|
||||
query?: never;
|
||||
query?: {
|
||||
/**
|
||||
* Refresh
|
||||
*
|
||||
* Force a fresh recompute, bypassing the cached value (and refreshing the cache).
|
||||
*/
|
||||
refresh?: boolean;
|
||||
};
|
||||
url: "/v1/default/banks/{bank_id}/stats";
|
||||
};
|
||||
|
||||
|
||||
@@ -18,6 +18,7 @@ uv run perf-test --output results.json # save JSON results
|
||||
|-------|-----------------|
|
||||
| `retain` | Full retain pipeline with mock LLM: fact extraction callback, embedding generation, DB writes, entity linking |
|
||||
| `recall` | Pre-populated bank recall: 4-way parallel retrieval (semantic, BM25, graph, temporal), RRF fusion, percentile latency |
|
||||
| `stats` | `/stats` endpoint (`get_bank_stats`): uncached aggregation latency (node/link counts + entity rollup join) vs. cached latency, plus cache speedup. Runs with the result cache **disabled** (TTL=0) so the headline numbers are the real per-poll cost |
|
||||
|
||||
### Scale Configurations
|
||||
|
||||
@@ -27,6 +28,19 @@ uv run perf-test --output results.json # save JSON results
|
||||
| `small` | 200 | 200 | 20 | 4 |
|
||||
| `medium` | 1,000 | 1,000 | 50 | 8 |
|
||||
| `large` | 5,000 | 5,000 | 100 | 16 |
|
||||
| `huge` | — (`stats` only) | — | — | — |
|
||||
|
||||
The `huge` scale is a prod-simulation for the `stats` suite only: it bulk-loads
|
||||
~500k units and ~17.8M physical `memory_links` rows (semantic + temporal +
|
||||
caused_by) via COPY, plus a `unit_entities` set whose `LEAST(n-1, 10)` rollup
|
||||
reproduces the ~110.9k *derived* entity links (entity edges aren't stored). The
|
||||
other suites fall back to `large` sizing at this scale, so prefer
|
||||
`--suite stats --scale huge`. Bulk-load takes a few minutes; FK triggers and
|
||||
non-essential `memory_links` indexes are dropped during COPY and restored after.
|
||||
|
||||
```bash
|
||||
uv run perf-test --suite stats --scale huge
|
||||
```
|
||||
|
||||
### CI
|
||||
|
||||
|
||||
@@ -14,6 +14,7 @@ Usage:
|
||||
uv run perf-test --suite retain
|
||||
uv run perf-test --suite recall
|
||||
uv run perf-test --suite graph-maintenance
|
||||
uv run perf-test --suite stats
|
||||
|
||||
# Configurable scale
|
||||
uv run perf-test --scale tiny # ~10s, CI smoke test
|
||||
@@ -21,6 +22,9 @@ Usage:
|
||||
uv run perf-test --scale medium # ~2min
|
||||
uv run perf-test --scale large # ~10min
|
||||
|
||||
# Prod-simulation for the stats suite (~500k units / ~18M links, bulk-loaded)
|
||||
uv run perf-test --suite stats --scale huge
|
||||
|
||||
# Save results as JSON
|
||||
uv run perf-test --output results.json
|
||||
"""
|
||||
@@ -32,8 +36,8 @@ import statistics
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Callable
|
||||
|
||||
from rich.console import Console
|
||||
from rich.progress import BarColumn, MofNCompleteColumn, Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
|
||||
@@ -65,6 +69,7 @@ SCALES: dict[str, dict[str, int]] = {
|
||||
"recall_concurrency": 1,
|
||||
"consolidation_items": 20,
|
||||
"graph_maintenance_bank_size": 20,
|
||||
"stats_bank_size": 20,
|
||||
},
|
||||
"small": {
|
||||
"retain_items": 200,
|
||||
@@ -73,6 +78,7 @@ SCALES: dict[str, dict[str, int]] = {
|
||||
"recall_concurrency": 4,
|
||||
"consolidation_items": 200,
|
||||
"graph_maintenance_bank_size": 200,
|
||||
"stats_bank_size": 200,
|
||||
},
|
||||
"medium": {
|
||||
"retain_items": 1_000,
|
||||
@@ -81,6 +87,7 @@ SCALES: dict[str, dict[str, int]] = {
|
||||
"recall_concurrency": 8,
|
||||
"consolidation_items": 1_000,
|
||||
"graph_maintenance_bank_size": 1_000,
|
||||
"stats_bank_size": 1_000,
|
||||
},
|
||||
"large": {
|
||||
"retain_items": 5_000,
|
||||
@@ -94,6 +101,32 @@ SCALES: dict[str, dict[str, int]] = {
|
||||
# as an Index Scan on idx_mu_emb_*. medium (1k) stays in the exact-scan
|
||||
# regime, so the two scales cover both planner paths.
|
||||
"graph_maintenance_bank_size": 15_000,
|
||||
# Large, entity-dense bank so the unit_entities→memory_units rollup join
|
||||
# in _compute_bank_stats is exercised at a size where its cost shows.
|
||||
"stats_bank_size": 15_000,
|
||||
},
|
||||
# Prod-simulation scale for the `stats` suite only. The numbers mirror a real
|
||||
# deployed bank: ~500k units and ~17.8M *physical* memory_links rows
|
||||
# (semantic + temporal + caused_by — entity links are NOT stored, they are
|
||||
# derived at query time from unit_entities). At this size the real retain
|
||||
# pipeline is infeasible, so `stats` bulk-loads via COPY (see
|
||||
# _bulk_populate_stats_bank). The non-stats keys fall back to `large` sizing
|
||||
# so a full `--scale huge` run doesn't try to retain 500k items.
|
||||
"huge": {
|
||||
"retain_items": 5_000,
|
||||
"recall_bank_size": 5_000,
|
||||
"recall_iterations": 10,
|
||||
"recall_concurrency": 4,
|
||||
"consolidation_items": 5_000,
|
||||
"graph_maintenance_bank_size": 15_000,
|
||||
"stats_bank_size": 500_000, # unused by the bulk path; kept for key parity
|
||||
"stats_units": 500_000,
|
||||
"stats_semantic_links": 9_460_147,
|
||||
"stats_temporal_links": 8_344_084,
|
||||
"stats_caused_by_links": 30_015,
|
||||
# Derived entity-link total to reproduce (unit_entities rollup, capped at
|
||||
# LEAST(n-1, 10) per entity). Not stored as memory_links rows.
|
||||
"stats_entity_links": 110_881,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -199,6 +232,23 @@ class GraphMaintenanceResult:
|
||||
temporal_calls: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class StatsResult:
|
||||
bank_size: int
|
||||
concurrency: int
|
||||
total_units: int
|
||||
total_links: int
|
||||
total_entities: int
|
||||
# Cache-miss path — engine._compute_bank_stats, the raw aggregation queries
|
||||
# (node/link/ops/doc counts + the unit_entities→memory_units entity rollup).
|
||||
cold_latency: PercentileStats
|
||||
cold_throughput_per_sec: float
|
||||
# Cache-hit path — engine.get_bank_stats, what a polling client actually
|
||||
# sees once the short-TTL per-process cache is warm.
|
||||
warm_latency: PercentileStats
|
||||
cache_speedup: float
|
||||
|
||||
|
||||
@dataclass
|
||||
class SuiteResult:
|
||||
name: str
|
||||
@@ -209,6 +259,7 @@ class SuiteResult:
|
||||
recall: RecallResult | None = None
|
||||
consolidation: ConsolidationResult | None = None
|
||||
graph_maintenance: GraphMaintenanceResult | None = None
|
||||
stats: StatsResult | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -292,6 +343,240 @@ async def _populate_bank(engine: Any, bank_id: str, size: int, event_date: str |
|
||||
pass
|
||||
|
||||
|
||||
# Average entities mentioned per unit when bulk-loading the stats bank — mirrors
|
||||
# the handful of entities the mock fact callback attaches per unit during real
|
||||
# retain, so the unit_entities table reaches a prod-like row count.
|
||||
STATS_ENTITY_MENTIONS_PER_UNIT = 3
|
||||
|
||||
# All memory_links indexes are dropped before the bulk COPY and recreated after.
|
||||
# Loading ~18M rows into an unindexed table + one bulk index build is far faster
|
||||
# than paying per-row btree maintenance on every insert (the dominant cost of the
|
||||
# load). Every index is restored afterward, so the table is left in its normal
|
||||
# shape — including idx_memory_links_bank_id_link_type, which the /stats link
|
||||
# GROUP BY relies on for its index-only scan, and the unique index, so this is
|
||||
# safe to run even against a shared database.
|
||||
_STATS_BULK_DROP_INDEXES = (
|
||||
"idx_memory_links_unique",
|
||||
"idx_memory_links_from_unit",
|
||||
"idx_memory_links_to_unit",
|
||||
"idx_memory_links_entity",
|
||||
"idx_memory_links_bank_id_link_type",
|
||||
)
|
||||
|
||||
|
||||
async def _bulk_populate_stats_bank(
|
||||
engine: Any,
|
||||
bank_id: str,
|
||||
*,
|
||||
units: int,
|
||||
semantic_links: int,
|
||||
temporal_links: int,
|
||||
caused_by_links: int,
|
||||
entity_link_target: int,
|
||||
) -> None:
|
||||
"""Bulk-load a prod-scale bank for the stats suite via COPY (no LLM/embeddings).
|
||||
|
||||
Rows are shaped exactly as the real retain pipeline leaves them so
|
||||
_compute_bank_stats sees a faithful workload:
|
||||
|
||||
* memory_units — `units` rows, NULL embeddings (stats never reads them), a
|
||||
realistic experience/world/observation fact_type mix.
|
||||
* memory_links — `semantic + temporal + caused_by` *physical* rows. Entity
|
||||
links are deliberately NOT inserted here: on the deployed schema the
|
||||
entity_id column is 100% NULL and the entity total is derived at query
|
||||
time. Pairs are generated collision-free so the unique index never trips.
|
||||
* entities / unit_entities — sized so the LEAST(n-1, 10) rollup in
|
||||
_compute_bank_stats reconstructs ~`entity_link_target` entity links: each
|
||||
of `entity_link_target` "shared" entities is attached to exactly 2 units
|
||||
(→ 1 derived link each), with the remaining mentions as singletons (→ 0)
|
||||
to reach a prod-like ~`units * STATS_ENTITY_MENTIONS_PER_UNIT` row count.
|
||||
|
||||
FK triggers on memory_links/unit_entities are disabled during the load (this
|
||||
is a throwaway perf bank), so COPY doesn't validate ~35M edge endpoints.
|
||||
"""
|
||||
from hindsight_api.engine.memory_engine import get_current_schema
|
||||
|
||||
pool = await engine._get_pool()
|
||||
schema = get_current_schema()
|
||||
|
||||
# Bulk maintenance (multi-million-row COPY, unique-index rebuild, VACUUM)
|
||||
# scales with the *whole* table, not just this bank, and can exceed the
|
||||
# engine's OLTP command_timeout when other large banks already exist. Give
|
||||
# these one-off fixture ops their own generous ceiling so a pre-existing
|
||||
# bank can't trip an (empty-message) asyncio.TimeoutError mid-load.
|
||||
bulk_timeout = 7200.0
|
||||
|
||||
unit_ids = [uuid.uuid4() for _ in range(units)]
|
||||
base_date = datetime(2024, 1, 1, tzinfo=timezone.utc)
|
||||
|
||||
def _fact_type(i: int) -> str:
|
||||
m = i % 20
|
||||
if m < 11:
|
||||
return "experience" # 55%
|
||||
if m < 18:
|
||||
return "world" # 35%
|
||||
return "observation" # 10%
|
||||
|
||||
def unit_records():
|
||||
for i, uid in enumerate(unit_ids):
|
||||
# Spread event_date across ~a year so date indexes look realistic.
|
||||
yield (uid, bank_id, f"perf stats unit {i}", base_date + timedelta(minutes=i % 525_600), _fact_type(i))
|
||||
|
||||
# Entity plan: `entity_link_target` shared entities (2 units each → 1 derived
|
||||
# link) + singleton entities padding out to the target unit_entities count.
|
||||
shared_entities = max(0, entity_link_target)
|
||||
total_unit_entities = max(2 * shared_entities, units * STATS_ENTITY_MENTIONS_PER_UNIT)
|
||||
singleton_entities = total_unit_entities - 2 * shared_entities
|
||||
shared_eids = [uuid.uuid4() for _ in range(shared_entities)]
|
||||
singleton_eids = [uuid.uuid4() for _ in range(singleton_entities)]
|
||||
half = max(1, units // 2)
|
||||
|
||||
def entity_records():
|
||||
for k, eid in enumerate(shared_eids):
|
||||
yield (eid, f"perf_entity_s_{k}", bank_id)
|
||||
for k, eid in enumerate(singleton_eids):
|
||||
yield (eid, f"perf_entity_x_{k}", bank_id)
|
||||
|
||||
def unit_entity_records():
|
||||
for j, eid in enumerate(shared_eids):
|
||||
yield (unit_ids[j % units], eid)
|
||||
yield (unit_ids[(j + half) % units], eid)
|
||||
for k, eid in enumerate(singleton_eids):
|
||||
yield (unit_ids[k % units], eid)
|
||||
|
||||
def link_records(link_type: str, count: int):
|
||||
# Collision-free distinct (from, to) pairs: spread `from` across leading
|
||||
# units, `to` across the rest skipping self. link_type is part of the
|
||||
# unique index, so semantic/temporal/caused_by never collide with each
|
||||
# other. COUNT(*) GROUP BY link_type is indifferent to the distribution.
|
||||
span = units - 1
|
||||
for k in range(count):
|
||||
f = k // span
|
||||
r = k % span
|
||||
t = r if r < f else r + 1
|
||||
yield (unit_ids[f], unit_ids[t], link_type, bank_id)
|
||||
|
||||
def _q(table: str) -> str:
|
||||
return f'"{schema}".{table}' if schema else table
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
TimeElapsedColumn(),
|
||||
console=console,
|
||||
) as progress:
|
||||
progress.add_task(
|
||||
f"Bulk-loading {units:,} units, "
|
||||
f"{semantic_links + temporal_links + caused_by_links:,} links, "
|
||||
f"{total_unit_entities:,} unit-entities…"
|
||||
)
|
||||
async with pool.acquire() as conn:
|
||||
# Every pooled connection carries a server-side statement_timeout
|
||||
# (the engine's runaway-query safety net, default 600s). The unique-
|
||||
# index rebuild and VACUUM over millions of rows legitimately exceed
|
||||
# it, so disable it on this connection for the load and restore it
|
||||
# after (PG cancels with "canceling statement due to statement
|
||||
# timeout" otherwise — which is exactly how the first 18M run died).
|
||||
prev_stmt_timeout = await conn.fetchval("SHOW statement_timeout")
|
||||
await conn.execute("SET statement_timeout = 0")
|
||||
await conn.copy_records_to_table(
|
||||
"entities",
|
||||
records=entity_records(),
|
||||
columns=["id", "canonical_name", "bank_id"],
|
||||
schema_name=schema,
|
||||
timeout=bulk_timeout,
|
||||
)
|
||||
await conn.copy_records_to_table(
|
||||
"memory_units",
|
||||
records=unit_records(),
|
||||
columns=["id", "bank_id", "text", "event_date", "fact_type"],
|
||||
schema_name=schema,
|
||||
timeout=bulk_timeout,
|
||||
)
|
||||
|
||||
dropped: list[str] = []
|
||||
triggers_disabled = False
|
||||
try:
|
||||
for idx in _STATS_BULK_DROP_INDEXES:
|
||||
await conn.execute(f"DROP INDEX IF EXISTS {_q(idx)}", timeout=bulk_timeout)
|
||||
dropped.append(idx)
|
||||
try:
|
||||
await conn.execute(f"ALTER TABLE {_q('memory_links')} DISABLE TRIGGER ALL", timeout=bulk_timeout)
|
||||
await conn.execute(f"ALTER TABLE {_q('unit_entities')} DISABLE TRIGGER ALL", timeout=bulk_timeout)
|
||||
triggers_disabled = True
|
||||
except Exception as exc: # noqa: BLE001 — best effort; fall back to validated COPY
|
||||
console.print(f" [yellow]Could not disable FK triggers ({exc}); COPY will validate FKs[/yellow]")
|
||||
|
||||
for link_type, count in (
|
||||
("semantic", semantic_links),
|
||||
("temporal", temporal_links),
|
||||
("caused_by", caused_by_links),
|
||||
):
|
||||
if count > 0:
|
||||
await conn.copy_records_to_table(
|
||||
"memory_links",
|
||||
records=link_records(link_type, count),
|
||||
columns=["from_unit_id", "to_unit_id", "link_type", "bank_id"],
|
||||
schema_name=schema,
|
||||
timeout=bulk_timeout,
|
||||
)
|
||||
await conn.copy_records_to_table(
|
||||
"unit_entities",
|
||||
records=unit_entity_records(),
|
||||
columns=["unit_id", "entity_id"],
|
||||
schema_name=schema,
|
||||
timeout=bulk_timeout,
|
||||
)
|
||||
finally:
|
||||
if triggers_disabled:
|
||||
await conn.execute(f"ALTER TABLE {_q('memory_links')} ENABLE TRIGGER ALL", timeout=bulk_timeout)
|
||||
await conn.execute(f"ALTER TABLE {_q('unit_entities')} ENABLE TRIGGER ALL", timeout=bulk_timeout)
|
||||
# Recreate the dropped indexes so the bank matches the prod schema
|
||||
# (the planner sees the same index set when the query runs).
|
||||
if "idx_memory_links_unique" in dropped:
|
||||
await conn.execute(
|
||||
f"CREATE UNIQUE INDEX IF NOT EXISTS idx_memory_links_unique ON {_q('memory_links')} "
|
||||
"(from_unit_id, to_unit_id, link_type, "
|
||||
"COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid))",
|
||||
timeout=bulk_timeout,
|
||||
)
|
||||
if "idx_memory_links_from_unit" in dropped:
|
||||
await conn.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_memory_links_from_unit ON {_q('memory_links')} (from_unit_id)",
|
||||
timeout=bulk_timeout,
|
||||
)
|
||||
if "idx_memory_links_to_unit" in dropped:
|
||||
await conn.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_memory_links_to_unit ON {_q('memory_links')} (to_unit_id)",
|
||||
timeout=bulk_timeout,
|
||||
)
|
||||
if "idx_memory_links_entity" in dropped:
|
||||
await conn.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_memory_links_entity ON {_q('memory_links')} (entity_id)",
|
||||
timeout=bulk_timeout,
|
||||
)
|
||||
if "idx_memory_links_bank_id_link_type" in dropped:
|
||||
# The index the /stats link-count GROUP BY uses — restore it
|
||||
# so the measured query plans against the prod index set.
|
||||
await conn.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_memory_links_bank_id_link_type "
|
||||
f"ON {_q('memory_links')} (bank_id, link_type)",
|
||||
timeout=bulk_timeout,
|
||||
)
|
||||
# VACUUM ANALYZE (not just ANALYZE): refresh planner stats AND set the
|
||||
# visibility map. Without the latter, the link_type COUNT can't use an
|
||||
# index-only scan (every tuple needs a heap visibility check), so a
|
||||
# fresh COPY measures a full heap scan that prod — where autovacuum
|
||||
# keeps the vm set — never pays. On an 18M-row bank this is the
|
||||
# difference between a ~1.6s seq scan and a ~1.0s index-only scan.
|
||||
await conn.execute(f"VACUUM (ANALYZE) {_q('memory_units')}", timeout=bulk_timeout)
|
||||
await conn.execute(f"VACUUM (ANALYZE) {_q('memory_links')}", timeout=bulk_timeout)
|
||||
await conn.execute(f"VACUUM (ANALYZE) {_q('unit_entities')}", timeout=bulk_timeout)
|
||||
# Restore the connection's statement_timeout before it returns to the
|
||||
# pool so later acquirers keep the runaway-query safety net.
|
||||
await conn.execute(f"SET statement_timeout = '{prev_stmt_timeout}'")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Suite: retain
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1139,6 +1424,201 @@ async def run_graph_maintenance_suite(scale_cfg: dict[str, int]) -> SuiteResult:
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Suite: stats
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def run_stats_suite(scale_cfg: dict[str, int]) -> SuiteResult:
|
||||
"""Measure the /stats endpoint (get_bank_stats) on a populated bank.
|
||||
|
||||
/stats is polled by the control plane and CLI, so it gets hammered far more
|
||||
often than retain/recall. Each cache miss runs a handful of bank-scoped
|
||||
aggregations — node counts by fact_type, link counts by link_type, and a
|
||||
unit_entities→memory_units rollup to reconstruct the entity-link total — so
|
||||
the cost grows with unit/entity density.
|
||||
|
||||
The suite runs with the per-process result cache **disabled** (TTL=0) so the
|
||||
headline numbers are the true uncached aggregation — the latency a server
|
||||
configured with ``HINDSIGHT_API_BANK_STATS_CACHE_TTL_SECONDS=0`` actually
|
||||
serves on every poll. It then re-enables the cache for a second pass to
|
||||
quantify what that cache buys:
|
||||
|
||||
* uncached — cache off, every call pays the full aggregation (the metric the
|
||||
"disable the cache" decision is about).
|
||||
* cached — cache on and warmed, what a steady polling client would see.
|
||||
|
||||
Population is either the real retain pipeline (small ``stats_bank_size``
|
||||
scales) or, for the prod-simulation ``huge`` scale, a direct COPY bulk-load
|
||||
(``_bulk_populate_stats_bank``) since retain can't reach ~500k units / ~18M
|
||||
links in reasonable time.
|
||||
"""
|
||||
from hindsight_api.engine.bank_stats_cache import BankStatsCache
|
||||
from hindsight_api.engine.schema import fq_table
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
bulk = "stats_semantic_links" in scale_cfg
|
||||
iterations = scale_cfg["recall_iterations"]
|
||||
concurrency = scale_cfg["recall_concurrency"]
|
||||
bank_id = f"perf-stats-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
if bulk:
|
||||
bank_size = scale_cfg["stats_units"]
|
||||
descr = (
|
||||
f"units={bank_size:,} links≈"
|
||||
f"{scale_cfg['stats_semantic_links'] + scale_cfg['stats_temporal_links'] + scale_cfg['stats_caused_by_links']:,}"
|
||||
)
|
||||
else:
|
||||
bank_size = scale_cfg["stats_bank_size"]
|
||||
descr = f"bank_size={bank_size:,}"
|
||||
|
||||
console.print(
|
||||
f"\n[bold cyan]Suite: stats[/bold cyan] "
|
||||
f"{descr} iterations={iterations} concurrency={concurrency} bank={bank_id}"
|
||||
)
|
||||
|
||||
engine = _build_engine(disable_observations=True)
|
||||
await engine.initialize()
|
||||
|
||||
# Disable the result cache on the server so the uncached pass measures the
|
||||
# real aggregation (mirrors HINDSIGHT_API_BANK_STATS_CACHE_TTL_SECONDS=0).
|
||||
def _disable_stats_cache() -> None:
|
||||
engine._bank_stats_cache = BankStatsCache(ttl_seconds=0, max_entries=0)
|
||||
|
||||
_disable_stats_cache()
|
||||
|
||||
if bulk:
|
||||
await _bulk_populate_stats_bank(
|
||||
engine,
|
||||
bank_id,
|
||||
units=scale_cfg["stats_units"],
|
||||
semantic_links=scale_cfg["stats_semantic_links"],
|
||||
temporal_links=scale_cfg["stats_temporal_links"],
|
||||
caused_by_links=scale_cfg["stats_caused_by_links"],
|
||||
entity_link_target=scale_cfg["stats_entity_links"],
|
||||
)
|
||||
else:
|
||||
# Real retain pipeline so units, links and entities all exist — the
|
||||
# entity rollup join is the part that scales with the bank.
|
||||
await _populate_bank(engine, bank_id, bank_size)
|
||||
|
||||
request_context = RequestContext()
|
||||
pool = await engine._get_pool()
|
||||
|
||||
units_row = await pool.fetchrow(
|
||||
f"SELECT COUNT(*) AS count FROM {fq_table('memory_units')} WHERE bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
links_row = await pool.fetchrow(
|
||||
f"SELECT COUNT(*) AS count FROM {fq_table('memory_links')} WHERE bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
# Derived entity-link total — the same unit_entities rollup _compute_bank_stats
|
||||
# reports as link_counts["entity"] (entity edges aren't stored; each entity
|
||||
# contributes LEAST(units_sharing - 1, 10) links). This is the prod-comparable
|
||||
# "entity" number, not the raw distinct-entity count.
|
||||
entities_row = await pool.fetchrow(
|
||||
f"""
|
||||
WITH per_entity AS (
|
||||
SELECT ue.entity_id, COUNT(*) AS n
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
JOIN {fq_table("memory_units")} mu ON mu.id = ue.unit_id
|
||||
WHERE mu.bank_id = $1
|
||||
GROUP BY ue.entity_id
|
||||
)
|
||||
SELECT COALESCE(SUM(LEAST(n - 1, 10)), 0)::bigint AS count FROM per_entity
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
total_units = int(units_row["count"]) if units_row else 0
|
||||
total_links = int(links_row["count"]) if links_row else 0
|
||||
total_entities = int(entities_row["count"]) if entities_row else 0
|
||||
|
||||
async def _run_batches(label: str, call: "Callable[[], Any]") -> list[float]:
|
||||
durations: list[float] = []
|
||||
|
||||
async def time_one() -> float:
|
||||
t0 = time.perf_counter()
|
||||
await call()
|
||||
return time.perf_counter() - t0
|
||||
|
||||
remaining = iterations
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
MofNCompleteColumn(),
|
||||
TimeElapsedColumn(),
|
||||
console=console,
|
||||
) as progress:
|
||||
task = progress.add_task(label, total=iterations)
|
||||
while remaining > 0:
|
||||
batch = min(concurrency, remaining)
|
||||
durations.extend(await asyncio.gather(*[time_one() for _ in range(batch)]))
|
||||
remaining -= batch
|
||||
progress.advance(task, batch)
|
||||
return durations
|
||||
|
||||
# Uncached: cache is disabled, so every get_bank_stats pays the full
|
||||
# aggregation — the endpoint latency with the cache turned off.
|
||||
cold_durations = await _run_batches(
|
||||
"Running stats (cache disabled / uncached)…",
|
||||
lambda: engine.get_bank_stats(bank_id, request_context=request_context),
|
||||
)
|
||||
# Cached: enable + warm the cache, then measure what a polling client sees.
|
||||
engine._bank_stats_cache = BankStatsCache(ttl_seconds=300.0, max_entries=16)
|
||||
await engine.get_bank_stats(bank_id, request_context=request_context) # prime
|
||||
warm_durations = await _run_batches(
|
||||
"Running stats (cache enabled / warm)…",
|
||||
lambda: engine.get_bank_stats(bank_id, request_context=request_context),
|
||||
)
|
||||
_disable_stats_cache() # leave the engine as the suite found it
|
||||
|
||||
await engine.delete_bank(bank_id=bank_id, request_context=request_context)
|
||||
await engine.close()
|
||||
|
||||
cold_stats = PercentileStats.from_samples(cold_durations)
|
||||
warm_stats = PercentileStats.from_samples(warm_durations)
|
||||
cold_total = sum(cold_durations)
|
||||
cold_throughput = iterations / (cold_total / concurrency) if cold_total > 0 else 0
|
||||
cache_speedup = cold_stats.p50 / warm_stats.p50 if warm_stats.p50 > 0 else 0
|
||||
|
||||
stats_result = StatsResult(
|
||||
bank_size=bank_size,
|
||||
concurrency=concurrency,
|
||||
total_units=total_units,
|
||||
total_links=total_links,
|
||||
total_entities=total_entities,
|
||||
cold_latency=cold_stats,
|
||||
cold_throughput_per_sec=round(cold_throughput, 2),
|
||||
warm_latency=warm_stats,
|
||||
cache_speedup=round(cache_speedup, 1),
|
||||
)
|
||||
|
||||
table = Table(title="Bank Stats Latency (cache disabled)")
|
||||
table.add_column("Metric", style="cyan")
|
||||
table.add_column("Value", style="green", justify="right")
|
||||
table.add_row(
|
||||
"Units / physical links / entity links (derived)",
|
||||
f"{total_units:,} / {total_links:,} / {total_entities:,}",
|
||||
)
|
||||
table.add_row("Iterations / concurrency", f"{iterations} / {concurrency}")
|
||||
table.add_row("Uncached throughput", f"{cold_throughput:.2f} calls/s")
|
||||
table.add_row("Uncached mean", f"{cold_stats.mean:.3f}s")
|
||||
table.add_row("Uncached p50 / p95 / p99", f"{cold_stats.p50:.3f}s / {cold_stats.p95:.3f}s / {cold_stats.p99:.3f}s")
|
||||
table.add_row("Cached mean", f"{warm_stats.mean:.4f}s")
|
||||
table.add_row("Cached p50 / p95 / p99", f"{warm_stats.p50:.4f}s / {warm_stats.p95:.4f}s / {warm_stats.p99:.4f}s")
|
||||
table.add_row("Cache speedup (p50)", f"{cache_speedup:.1f}x")
|
||||
console.print(table)
|
||||
|
||||
return SuiteResult(
|
||||
name="stats",
|
||||
duration_seconds=round(cold_total, 3),
|
||||
success=True,
|
||||
stats=stats_result,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registry and orchestrator
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1150,6 +1630,7 @@ SUITES = {
|
||||
"recall-temporal": run_recall_temporal_suite,
|
||||
"consolidation": run_consolidation_suite,
|
||||
"graph-maintenance": run_graph_maintenance_suite,
|
||||
"stats": run_stats_suite,
|
||||
}
|
||||
|
||||
|
||||
@@ -1423,6 +1904,38 @@ def _print_summary(results: PerfTestResults) -> None:
|
||||
"",
|
||||
)
|
||||
|
||||
if suite.stats:
|
||||
st = suite.stats
|
||||
table.add_row(
|
||||
suite.name,
|
||||
status,
|
||||
"uncached latency",
|
||||
f"mean={st.cold_latency.mean:.3f}s",
|
||||
f"{st.cold_latency.p50:.3f}s",
|
||||
f"{st.cold_latency.p95:.3f}s",
|
||||
f"{st.cold_latency.p99:.3f}s",
|
||||
)
|
||||
table.add_row(
|
||||
"",
|
||||
"",
|
||||
"cached latency",
|
||||
f"mean={st.warm_latency.mean:.4f}s",
|
||||
f"{st.warm_latency.p50:.4f}s",
|
||||
f"{st.warm_latency.p95:.4f}s",
|
||||
f"{st.warm_latency.p99:.4f}s",
|
||||
)
|
||||
table.add_row("", "", "uncached throughput", f"{st.cold_throughput_per_sec} calls/s", "", "", "")
|
||||
table.add_row("", "", "cache speedup (p50)", f"{st.cache_speedup}x", "", "", "")
|
||||
table.add_row(
|
||||
"",
|
||||
"",
|
||||
"units/phys-links/entity-links",
|
||||
f"{st.total_units:,} / {st.total_links:,} / {st.total_entities:,}",
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
)
|
||||
|
||||
if not suite.success:
|
||||
table.add_row(suite.name, status, "error", suite.error or "unknown", "", "", "")
|
||||
|
||||
|
||||
@@ -880,6 +880,18 @@
|
||||
"title": "Bank Id"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "refresh",
|
||||
"in": "query",
|
||||
"required": false,
|
||||
"schema": {
|
||||
"type": "boolean",
|
||||
"description": "Force a fresh recompute, bypassing the cached value (and refreshing the cache).",
|
||||
"default": false,
|
||||
"title": "Refresh"
|
||||
},
|
||||
"description": "Force a fresh recompute, bypassing the cached value (and refreshing the cache)."
|
||||
},
|
||||
{
|
||||
"name": "authorization",
|
||||
"in": "header",
|
||||
|
||||
@@ -880,6 +880,18 @@
|
||||
"title": "Bank Id"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "refresh",
|
||||
"in": "query",
|
||||
"required": false,
|
||||
"schema": {
|
||||
"type": "boolean",
|
||||
"description": "Force a fresh recompute, bypassing the cached value (and refreshing the cache).",
|
||||
"default": false,
|
||||
"title": "Refresh"
|
||||
},
|
||||
"description": "Force a fresh recompute, bypassing the cached value (and refreshing the cache)."
|
||||
},
|
||||
{
|
||||
"name": "authorization",
|
||||
"in": "header",
|
||||
|
||||
Reference in New Issue
Block a user