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:
parent
5b8e36987b
commit
7ada42823f
2 changed files with 80 additions and 0 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue