qa/query: JIT-hydrate doc chunks from bucket blobs on Tier B consumer

Adds hydrate_doc_jit(conn, doc_root, backend): for every chunk in this
document whose local content IS NULL, fetch blobs/<leaf_hash> from the
configured cold backend, verify hash, cache into the row. _load_doc_text
and _load_doc_chunks call it transparently before reading.

Env-gated: ARBORIST_JIT_CHUNKS=1 opts in (so non-JIT environments stay
a pure no-op); the backend comes from the standard ARBORIST_COLD_*
config. With this wired in, a Tier B consumer can clone metadata-only
shards and the query path transparently fetches answer chunks JIT — no
caller code changes.
This commit is contained in:
russell@unturf.com 2026-05-30 06:51:04 -04:00
parent 5b8e36987b
commit 7ada42823f
No known key found for this signature in database
2 changed files with 80 additions and 0 deletions

View file

@ -386,6 +386,56 @@ def clone_from_bucket(
# ---------- JIT chunk fetch (Tier B consumer) -------------------------------
def hydrate_doc_jit(
conn: sqlite3.Connection,
document_root: str,
backend: ObjectStoreBackend,
*,
fetch_workers: int = 16,
) -> int:
"""Fetch any chunks for ``document_root`` whose ``content`` is NULL
in the local shard from ``blobs/<leaf_hash>`` on the bucket and
cache them back into the row. No-op if every chunk already has
content (the cheap WHERE count short-circuits). Returns the number
of chunks fetched.
Used by Tier B consumers ('clone --just-enough' produced metadata-
only shards): the query path calls this just before reading chunk
text for a hit, so the rest of the pipeline ('SELECT content FROM
chunks WHERE content IS NOT NULL') sees the same locally-resident
shape it would for a full clone."""
from arborist.merkle import hash_leaf
rows = conn.execute(
"SELECT chunk_id, leaf_hash FROM chunks "
"WHERE document_root = ? AND content IS NULL",
(document_root,),
).fetchall()
if not rows:
return 0
def _fetch(pair: tuple[int, str]) -> tuple[int, bytes]:
chunk_id, leaf_hash = pair
key = f"{BLOB_PREFIX}/{leaf_hash[:2]}/{leaf_hash[2:]}"
body = backend.get(key)
if hash_leaf(body).hex() != leaf_hash:
raise ValueError(
f"chunk {chunk_id}: bucket blob hash mismatch (leaf_hash {leaf_hash[:12]})"
)
return chunk_id, body
pairs = [(r[0], r[1]) for r in rows]
n_workers = max(1, min(fetch_workers, len(pairs)))
with ThreadPoolExecutor(max_workers=n_workers) as ex:
fetched = list(ex.map(_fetch, pairs))
with conn:
conn.executemany(
"UPDATE chunks SET content = ? WHERE chunk_id = ?",
[(body, cid) for cid, body in fetched],
)
return len(fetched)
def fetch_chunk_jit(
conn: sqlite3.Connection,
chunk_id: int,

View file

@ -1850,8 +1850,37 @@ def _rerank_by_body_coverage(
return hits
def _maybe_jit_hydrate(shard_path: str, document_root: str) -> None:
"""If ARBORIST_JIT_CHUNKS=1 and a cold backend is configured
(ARBORIST_COLD_BUCKET / _ENDPOINT_URL), fetch any chunks for this
document whose ``content`` is NULL from ``blobs/<hash>`` on the
bucket and cache them locally. No-op otherwise. Lets a Tier B
consumer (metadata-only clone) serve queries by pulling answer
chunks from S3 just in time."""
import os
if os.environ.get("ARBORIST_JIT_CHUNKS", "0") != "1":
return
try:
from arborist.cli import _make_cold_backend
from arborist.cold_clone import hydrate_doc_jit
backend = _make_cold_backend()
except Exception:
return
conn = connect(shard_path)
try:
n = hydrate_doc_jit(conn, document_root, backend)
if n:
print(
f"[jit] hydrated {n} chunks for {document_root[:12]} from bucket",
file=sys.stderr,
)
finally:
conn.close()
def _load_doc_text(shard_path: str, document_root: str) -> str | None:
"""Concatenate all hot chunks of a document. Returns None if cold or missing."""
_maybe_jit_hydrate(shard_path, document_root)
conn = connect(shard_path)
try:
rows = conn.execute(
@ -1880,6 +1909,7 @@ def _load_doc_chunks(
paragraph that supports a claim instead of lazy-anchoring the
whole article.
"""
_maybe_jit_hydrate(shard_path, document_root)
conn = connect(shard_path)
try:
rows = conn.execute(