homelab/stack
1
0
Fork 0
Code Issues 58 Pull Requests 1 Actions Packages Projects 1 Releases Wiki Activity

fix(llm): per-call DuckDB cursor for the cached replica handle; lock the cache; concurrency test (refs P48)

Browse Source 
A /chat stream is a sync generator, so Starlette runs it in a threadpool
worker: concurrent turns shared the one cached DuckDBPyConnection and
read each other's result sets ("not enough values to unpack"), which the
broad except swallowed as "valuation evidence skipped" — the table just
vanished. Each call now runs on its own cursor off the shared handle and
closes it in a finally, and the cache check-and-open is under a lock so
a race on the first open cannot leak a connection.
This commit is contained in:
kert
2026-09-08 23:36:01 -04:00
parent 456aee26bb
commit 418fab29c3
2 changed files with 131 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
src/llm/evidence.py
View File

@@ -12,6 +12,7 @@ from __future__ import annotations
import logging
import os
import threading
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Sequence
@@ -98,6 +99,7 @@ class ValuationEvidence:
#: so the handle stays valid on the old inode until we notice the new
#: mtime and reopen.
_REPLICA: tuple[str, int, Any] | None = None
_REPLICA_LOCK = threading.Lock()
def _mtime(path: str) -> int:
@@ -112,18 +114,19 @@ def _connect(path: str) -> Any:
replica file has been republished since we opened it."""
global _REPLICA
mtime = _mtime(path)
if _REPLICA is not None:
cached_path, cached_mtime, con = _REPLICA
if (cached_path, cached_mtime) == (path, mtime):
return con
_REPLICA = None
try:
con.close()
except Exception as e: # noqa: BLE001 — a stale handle must not break the chat
log.warning("stale replica handle not closed: %s", e)
con = duckdb.connect(path, read_only=True)
_REPLICA = (path, mtime, con)
return con
with _REPLICA_LOCK: # two chat threads racing the first open would leak one
if _REPLICA is not None:
cached_path, cached_mtime, con = _REPLICA
if (cached_path, cached_mtime) == (path, mtime):
return con
_REPLICA = None
try:
con.close()
except Exception as e: # noqa: BLE001 — a stale handle must not break the chat
log.warning("stale replica handle not closed: %s", e)
con = duckdb.connect(path, read_only=True)
_REPLICA = (path, mtime, con)
return con
def _replica_path(cfg: LlmConfig) -> str:
@@ -141,7 +144,10 @@ def valuation_evidence(question: str, cfg: LlmConfig) -> ValuationEvidence | Non
return None
path = _replica_path(cfg)
try:
con = _connect(path)
# A chat turn runs in a threadpool worker, and DuckDB gives each
# thread its own cursor off the shared connection — without one,
# concurrent turns read each other's result sets.
cur = _connect(path).cursor()
except Exception as e: # noqa: BLE001 — duckdb raises several types
log.warning("valuation evidence skipped (%s): %s", path, e)
return None
@@ -149,10 +155,12 @@ def valuation_evidence(question: str, cfg: LlmConfig) -> ValuationEvidence | Non
# crowd the excerpts out of the prompt.
years = cfg.valuation_years if len(det.codes) <= 3 else 2
try:
rows, unpriced = valuation(con, list(det.codes), years=years)
rows, unpriced = valuation(cur, list(det.codes), years=years)
except Exception as e: # noqa: BLE001
log.warning("valuation evidence skipped (query): %s", e)
return None
finally:
cur.close() # the cached parent connection stays open
return ValuationEvidence(det.codes, det.families, tuple(rows), tuple(unpriced))

115
tests/llm/test_evidence.py
View File

@@ -2,8 +2,11 @@
from __future__ import annotations
from concurrent.futures import ThreadPoolExecutor
from dataclasses import replace
from unittest.mock import MagicMock, patch
import duckdb
import pytest
from llm import evidence
@@ -32,6 +35,14 @@ CFG = LlmConfig(
)
@pytest.fixture(autouse=True)
def _clear_replica_cache():
"""The replica handle is process-global; no test may inherit another's."""
evidence._REPLICA = None
yield
evidence._REPLICA = None
def _row(**kw) -> ValuationRow:
base = dict(
code="G0556",
@@ -161,12 +172,6 @@ class TestEvidenceText:
class TestValuationEvidence:
@pytest.fixture(autouse=True)
def _clear_replica_cache(self):
evidence._REPLICA = None
yield
evidence._REPLICA = None
def test_none_when_no_codes(self):
assert (
valuation_evidence("what did commenters say about telehealth?", CFG) is None
@@ -223,6 +228,104 @@ class TestValuationEvidence:
assert mock_val.call_args.kwargs == {"years": 2}
RVU_COLS = (
"hcpcs VARCHAR, mod VARCHAR, description VARCHAR, status_code VARCHAR, "
"work_rvu DOUBLE, non_fac_pe_rvu DOUBLE, fac_pe_rvu DOUBLE, mp_rvu DOUBLE, "
"non_fac_total DOUBLE, fac_total DOUBLE, conv_factor DOUBLE, year INTEGER"
)
PROPOSED_COLS = (
"hcpcs VARCHAR, mod VARCHAR, description VARCHAR, status_code VARCHAR, "
"work_rvu DOUBLE, non_fac_pe_rvu DOUBLE, fac_pe_rvu DOUBLE, mp_rvu DOUBLE, "
"cms_rule_id VARCHAR"
)
class TestConcurrentChats:
"""/chat streams run in Starlette's threadpool, so several turns share
the cached replica handle. Without a per-call cursor their result sets
cross-talk and the broad except swallows it as "no valuation"."""
@pytest.fixture
def replica(self, tmp_path):
db = tmp_path / "aco.ro.duckdb"
con = duckdb.connect(str(db))
con.execute("CREATE SCHEMA pfs")
con.execute(f"CREATE TABLE pfs.rvu ({RVU_COLS})")
con.execute(f"CREATE TABLE pfs.rvu_proposed ({PROPOSED_COLS})")
# two unfamilied codes with different row counts, so a swapped
# result set is visible as a wrong count
con.executemany(
"INSERT INTO pfs.rvu VALUES (?,?,?,?,?,?,?,?,?,?,?,?)",
[
(
"99213",
None,
"Office visit",
"A",
1.0,
1.0,
0.5,
0.05,
2.05,
1.55,
30.0,
2026,
),
*[
(
"99214",
None,
"Office visit",
"A",
1.5,
1.0,
0.5,
0.05,
2.55,
2.05,
30.0,
y,
)
for y in (2024, 2025, 2026)
],
],
)
con.close()
return db
def test_parallel_turns_share_one_handle_without_cross_talk(self, replica):
cfg = replace(CFG, duckdb_replica=str(replica))
real_connect, opens = duckdb.connect, []
def counting_connect(*a, **kw):
opens.append(a)
return real_connect(*a, **kw)
expected = {"99213": 1, "99214": 3}
questions = ["What does 99213 pay?", "What does 99214 pay?"]
def one(i):
q = questions[i % 2]
ev = valuation_evidence(q, cfg)
assert ev is not None, "valuation dropped under concurrency"
(code,) = ev.codes
return code, len(ev.rows)
with patch("llm.evidence.duckdb.connect", counting_connect):
with ThreadPoolExecutor(max_workers=8) as pool:
results = list(pool.map(one, range(160)))
assert all(n == expected[code] for code, n in results), sorted(set(results))
assert len(opens) == 1 # one handle, shared
def test_cursor_is_closed_but_the_handle_is_kept(self, replica):
cfg = replace(CFG, duckdb_replica=str(replica))
assert valuation_evidence("99213?", cfg) is not None
assert evidence._REPLICA is not None
# the cached parent still answers — it was never closed
assert valuation_evidence("99214?", cfg) is not None
def _engine(rows):
engine = MagicMock()
conn = engine.begin.return_value.__enter__.return_value