russell@unturf.com
d43714a503
Replaces the batched chunk-pack phase with per-chunk content-addressed blob uploads to `blobs/<hash[:2]>/<hash[2:]>`. The metadata pack still ships (small, fast to restore), but consumers no longer have to pull multi-GB chunk packs to get queryable: `cold unpack --mode just-enough` + `ARBORIST_JIT_CHUNKS=1` fetches single chunks on cache miss. Producer (`_stream_jit_blobs` in evict.py): - ThreadPoolExecutor with bounded queue (workers*4) keeps memory flat across millions of chunks - HEAD-checks object_size for idempotent re-upload - Mutually exclusive with chunk packs — manifest's `chunk_pack_hashes` is empty in JIT mode (consumer reads that as "JIT-only") Consumer (`hydrate_doc_jit` in cold_clone.py + `_maybe_jit_hydrate` in qa/query.py): - Detects both content shapes that need JIT: NULL (Tier B raw-clone) and zeroblob placeholders (just-enough pack restore, per #53). Discriminator is first-byte = NUL — zstd-framed bodies start with 0x28, plain UTF-8 prose never has leading NUL. - Same placeholder filter applied to chunk-read sites in qa/query.py so partial hydrate doesn't surface zero-bytes content into the LLM context. Test (`TestJitBlobsPackMode` in tests/test_cold_unpack_routed.py): - End-to-end push → just-enough hydrate → JIT-fetch → content matches original byte-for-byte through `unpack_chunk`. Docs (cold-object-store.md): - Hard-invariant #1 updated: bucket holds packs by default; `blobs/` and `clones/` are opt-in prefixes for the JIT and Tier-A flows. - New "Three consumer modes" section: full-pack vs JIT-blobs vs raw-clone comparison table + operator decision tree.
444 lines
16 KiB
Python
444 lines
16 KiB
Python
"""Tests for #000067 M-aware cold-pack hydration.
|
|
|
|
Build a tiny corpus → pack into a MemoryBackend → hydrate into M=4
|
|
fresh target shards via the routed path. Assert every doc lands on
|
|
``shard_for_document(document_root, M)`` and corpus-wide tables
|
|
consolidate to target 0.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from pathlib import Path
|
|
from typing import Iterator
|
|
|
|
import pytest
|
|
|
|
from arborist.cold_object import MemoryBackend
|
|
from arborist.document import Document, shard_for_document
|
|
from arborist.evict import (
|
|
hydrate_from_metadata_pack_routed,
|
|
push_pack,
|
|
)
|
|
from arborist.ingest import ingest_source
|
|
from arborist.source import Source
|
|
from arborist.store import (
|
|
SCHEMA_SQL,
|
|
connect,
|
|
get_corpus_shard_count,
|
|
)
|
|
|
|
|
|
class _FakeSource(Source):
|
|
source_type = "html"
|
|
|
|
def __init__(self, docs: list[Document]):
|
|
self.docs = docs
|
|
|
|
def iter_documents(self) -> Iterator[Document]:
|
|
yield from self.docs
|
|
|
|
|
|
def _make_doc(uri: str, seed: str) -> Document:
|
|
content = (
|
|
f"This is doc {uri} seed {seed}. " * 10
|
|
+ f"Body text {seed} with tokens enough for chunking. " * 30
|
|
)
|
|
return Document(
|
|
uri=uri, content=content, source_type="html",
|
|
title=uri.split("/")[-1],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def producer_shard(tmp_path: Path) -> Path:
|
|
"""A small producer corpus packed into one shard."""
|
|
db = tmp_path / "producer.db"
|
|
conn = connect(db)
|
|
try:
|
|
docs = [
|
|
_make_doc(f"https://prod/doc-{i}", seed=f"s0d{i}")
|
|
for i in range(20)
|
|
]
|
|
ingest_source(conn, _FakeSource(docs))
|
|
finally:
|
|
conn.close()
|
|
return db
|
|
|
|
|
|
def _make_target_shards(target_dir: Path, M: int) -> list[sqlite3.Connection]:
|
|
targets: list[sqlite3.Connection] = []
|
|
for i in range(M):
|
|
path = target_dir / f"{i:03d}.db"
|
|
c = sqlite3.connect(str(path))
|
|
c.row_factory = sqlite3.Row
|
|
c.executescript(SCHEMA_SQL)
|
|
c.execute("PRAGMA foreign_keys = OFF")
|
|
targets.append(c)
|
|
return targets
|
|
|
|
|
|
class TestRoutedHydrate:
|
|
def test_pack_then_hydrate_routed(
|
|
self, producer_shard: Path, tmp_path: Path
|
|
):
|
|
# Pack the producer shard into a MemoryBackend.
|
|
backend = MemoryBackend()
|
|
with sqlite3.connect(str(producer_shard)) as src_conn:
|
|
src_conn.row_factory = sqlite3.Row
|
|
push_result = push_pack(
|
|
src_conn, backend,
|
|
document_root=None,
|
|
max_chunks=None,
|
|
max_pack_bytes=10 * 1024 ** 2,
|
|
allow_license_class="unknown",
|
|
)
|
|
# Identify the metadata pack hash.
|
|
metadata_pack_hash = None
|
|
for p in push_result.get("packs", []):
|
|
if p.get("kind") == "metadata":
|
|
metadata_pack_hash = p["pack_hash"]
|
|
break
|
|
if metadata_pack_hash is None:
|
|
metadata_pack_hash = push_result["packs"][0]["pack_hash"]
|
|
|
|
# Hydrate into a fresh 4-shard target.
|
|
target_dir = tmp_path / "tgt"
|
|
target_dir.mkdir()
|
|
M = 4
|
|
targets = _make_target_shards(target_dir, M)
|
|
try:
|
|
hydrate_result = hydrate_from_metadata_pack_routed(
|
|
targets, backend, metadata_pack_hash, M=M, mode="full",
|
|
)
|
|
finally:
|
|
for t in targets:
|
|
t.close()
|
|
assert hydrate_result["status"] == "hydrated"
|
|
assert hydrate_result["M"] == M
|
|
|
|
# Verify: every doc on the producer must land on its hash-routed
|
|
# target shard.
|
|
prod_roots: dict[str, int] = {}
|
|
c = sqlite3.connect(f"file:{producer_shard}?mode=ro", uri=True)
|
|
try:
|
|
for (root,) in c.execute("SELECT document_root FROM documents"):
|
|
prod_roots[root] = shard_for_document(root, M)
|
|
finally:
|
|
c.close()
|
|
|
|
for tidx in range(M):
|
|
tdb = target_dir / f"{tidx:03d}.db"
|
|
c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True)
|
|
try:
|
|
actual = {
|
|
r[0] for r in c.execute("SELECT document_root FROM documents")
|
|
}
|
|
finally:
|
|
c.close()
|
|
for root, expected_idx in prod_roots.items():
|
|
if expected_idx == tidx:
|
|
assert root in actual, (
|
|
f"doc {root[:8]}… expected on shard {tidx} but missing"
|
|
)
|
|
else:
|
|
assert root not in actual, (
|
|
f"doc {root[:8]}… on shard {tidx} but hashes to {expected_idx}"
|
|
)
|
|
|
|
def test_corpus_shard_count_set_on_routed_hydrate(
|
|
self, producer_shard: Path, tmp_path: Path
|
|
):
|
|
"""CLI sets corpus_shard_count meta on every target; here the
|
|
test directly calls the helper to confirm the meta plumbing
|
|
works from the routed entry point."""
|
|
backend = MemoryBackend()
|
|
with sqlite3.connect(str(producer_shard)) as src_conn:
|
|
src_conn.row_factory = sqlite3.Row
|
|
push_result = push_pack(
|
|
src_conn, backend,
|
|
document_root=None,
|
|
max_chunks=None,
|
|
max_pack_bytes=10 * 1024 ** 2,
|
|
allow_license_class="unknown",
|
|
)
|
|
metadata_pack_hash = next(
|
|
p["pack_hash"] for p in push_result["packs"]
|
|
if p.get("kind") == "metadata"
|
|
)
|
|
|
|
target_dir = tmp_path / "tgt"
|
|
target_dir.mkdir()
|
|
M = 4
|
|
targets = _make_target_shards(target_dir, M)
|
|
try:
|
|
hydrate_from_metadata_pack_routed(
|
|
targets, backend, metadata_pack_hash, M=M, mode="full",
|
|
)
|
|
# Simulate the CLI's post-hydrate meta stamp.
|
|
from arborist.store import set_corpus_shard_count
|
|
for t in targets:
|
|
with t:
|
|
set_corpus_shard_count(t, M)
|
|
finally:
|
|
for t in targets:
|
|
t.close()
|
|
for i in range(M):
|
|
tdb = target_dir / f"{i:03d}.db"
|
|
c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True)
|
|
try:
|
|
assert get_corpus_shard_count(c) == M
|
|
finally:
|
|
c.close()
|
|
|
|
|
|
class TestFtsPackRoutingRegression:
|
|
"""Regression test for the v4 bench bug (#52): _pull_fts_pack_into_targets
|
|
used to INSERT shadow tables into every target, cross-contaminating
|
|
FTS indexes across shards. After the fix it should INSERT only into
|
|
the ONE target that owns the chunks the fts pack references.
|
|
|
|
Assertion: per-shard chunks_fts row count equals per-shard chunks row
|
|
count (no cross-target leakage).
|
|
"""
|
|
|
|
def test_fts_pack_only_on_owning_target(
|
|
self, producer_shard: Path, tmp_path: Path
|
|
):
|
|
from arborist.evict import push_pack
|
|
backend = MemoryBackend()
|
|
with sqlite3.connect(str(producer_shard)) as src_conn:
|
|
src_conn.row_factory = sqlite3.Row
|
|
push_result = push_pack(
|
|
src_conn, backend,
|
|
document_root=None,
|
|
max_chunks=None,
|
|
max_pack_bytes=10 * 1024 ** 2,
|
|
allow_license_class="unknown",
|
|
include_fts=True,
|
|
)
|
|
metadata_pack_hash = next(
|
|
p["pack_hash"] for p in push_result["packs"]
|
|
if p.get("kind") == "metadata"
|
|
)
|
|
|
|
target_dir = tmp_path / "tgt"
|
|
target_dir.mkdir()
|
|
M = 4
|
|
targets = _make_target_shards(target_dir, M)
|
|
try:
|
|
hydrate_from_metadata_pack_routed(
|
|
targets, backend, metadata_pack_hash, M=M, mode="full",
|
|
)
|
|
finally:
|
|
for t in targets:
|
|
t.close()
|
|
|
|
# Regression check: the v4 bug INSERT'd the fts shadow tables
|
|
# into EVERY target shard. After the fix, only ONE target gets
|
|
# the FTS data (the target whose chunks the fts pack references).
|
|
# The test fixture is single-shard so the producer's pack contains
|
|
# FTS data for all docs, but only one consumer target hash-routes
|
|
# to ownership of the chunks the FTS data references. The other
|
|
# 3 targets must have ZERO chunks_fts data (no cross-contamination).
|
|
per_shard = []
|
|
for i in range(M):
|
|
tdb = target_dir / f"{i:03d}.db"
|
|
c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True)
|
|
try:
|
|
chunks = c.execute("SELECT COUNT(*) FROM chunks").fetchone()[0]
|
|
fts_data = c.execute("SELECT COUNT(*) FROM chunks_fts_data").fetchone()[0]
|
|
fts_docsize = c.execute("SELECT COUNT(*) FROM chunks_fts_docsize").fetchone()[0]
|
|
per_shard.append({"chunks": chunks, "fts_data": fts_data, "fts_docsize": fts_docsize})
|
|
finally:
|
|
c.close()
|
|
# Exactly one shard should hold FTS data; the rest must be empty.
|
|
# (FTS5 shadow tables are segment-based — not subsettable per-row
|
|
# — so the owning shard receives ALL of the fts pack's rows even
|
|
# if its own chunks count is smaller.)
|
|
with_fts = [s for s in per_shard if s["fts_docsize"] > 0]
|
|
without_fts = [s for s in per_shard if s["fts_docsize"] == 0]
|
|
assert len(with_fts) == 1, (
|
|
f"expected exactly 1 target with FTS data; got {len(with_fts)} — "
|
|
"v4 cross-target contamination bug back"
|
|
)
|
|
# chunks_fts_data has 1-2 default FTS5 internal config/structure
|
|
# rows on a freshly-created virtual table; the per-document signal
|
|
# lives in chunks_fts_docsize which is what we asserted above.
|
|
|
|
|
|
class TestPreSizedChunks:
|
|
"""#53: producer dumps a synthetic _content_size column for chunks;
|
|
consumer INSERTs with right-sized BLOB placeholders so phase 2
|
|
UPDATE replaces same-size content in place — no row growth, no
|
|
page splits, no ext4 metadata-journal events per UPDATE.
|
|
"""
|
|
|
|
def test_chunks_content_pre_sized_after_metadata_restore(
|
|
self, producer_shard: Path, tmp_path: Path
|
|
):
|
|
from arborist.evict import push_pack
|
|
backend = MemoryBackend()
|
|
with sqlite3.connect(str(producer_shard)) as src_conn:
|
|
src_conn.row_factory = sqlite3.Row
|
|
push_result = push_pack(
|
|
src_conn, backend,
|
|
document_root=None,
|
|
max_chunks=None,
|
|
max_pack_bytes=10 * 1024 ** 2,
|
|
allow_license_class="unknown",
|
|
include_fts=True,
|
|
)
|
|
metadata_pack_hash = next(
|
|
p["pack_hash"] for p in push_result["packs"]
|
|
if p.get("kind") == "metadata"
|
|
)
|
|
|
|
# Pull ONLY the metadata pack (just-enough hydration). Verify
|
|
# chunks rows land with right-sized BLOB content (NOT NULL).
|
|
target_dir = tmp_path / "tgt"
|
|
target_dir.mkdir()
|
|
M = 4
|
|
targets = _make_target_shards(target_dir, M)
|
|
try:
|
|
hydrate_from_metadata_pack_routed(
|
|
targets, backend, metadata_pack_hash, M=M, mode="just-enough",
|
|
)
|
|
finally:
|
|
for t in targets:
|
|
t.close()
|
|
# Find the target with chunks rows; verify content is non-NULL
|
|
# BLOB of correct size.
|
|
for i in range(M):
|
|
tdb = target_dir / f"{i:03d}.db"
|
|
c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True)
|
|
try:
|
|
rows = c.execute(
|
|
"SELECT chunk_id, length(content) AS sz, content IS NOT NULL AS has "
|
|
"FROM chunks LIMIT 3"
|
|
).fetchall()
|
|
finally:
|
|
c.close()
|
|
if not rows:
|
|
continue
|
|
for cid, sz, has in rows:
|
|
assert has, f"chunk {cid} on shard {i} has NULL content (#53 not applied)"
|
|
assert sz > 0, f"chunk {cid} on shard {i} has zero-size placeholder"
|
|
|
|
|
|
class TestJitBlobsPackMode:
|
|
"""`cold pack --jit-blobs` skips chunk packs and uploads per-chunk
|
|
content-addressed blobs to `blobs/<hash[:2]>/<hash[2:]>`. A consumer
|
|
pulls the metadata pack via `cold unpack --mode just-enough` and
|
|
then JIT-fetches individual chunks from the bucket on demand.
|
|
"""
|
|
|
|
def test_jit_blobs_skips_chunk_packs_and_uploads_blobs(
|
|
self, producer_shard: Path, tmp_path: Path
|
|
):
|
|
from arborist.cold_clone import hydrate_doc_jit
|
|
from arborist.compress import unpack_chunk
|
|
from arborist.document import shard_for_document
|
|
|
|
backend = MemoryBackend()
|
|
with sqlite3.connect(str(producer_shard)) as src_conn:
|
|
src_conn.row_factory = sqlite3.Row
|
|
push_result = push_pack(
|
|
src_conn, backend,
|
|
document_root=None,
|
|
max_chunks=None,
|
|
max_pack_bytes=10 * 1024 ** 2,
|
|
allow_license_class="unknown",
|
|
jit_blobs=True,
|
|
)
|
|
|
|
kinds = [p.get("kind") for p in push_result.get("packs", [])]
|
|
assert "chunks" not in kinds, "jit_blobs should skip chunk packs"
|
|
assert kinds.count("metadata") == 1
|
|
assert push_result["jit_blobs"] is True
|
|
stats = push_result["jit_blob_stats"]
|
|
assert stats["put"] > 0
|
|
|
|
# Every chunk row's leaf_hash must have a blob in the bucket.
|
|
with sqlite3.connect(f"file:{producer_shard}?mode=ro", uri=True) as src:
|
|
src.row_factory = sqlite3.Row
|
|
chunks = src.execute(
|
|
"SELECT chunk_id, document_root, leaf_hash, content "
|
|
"FROM chunks WHERE content IS NOT NULL"
|
|
).fetchall()
|
|
for r in chunks:
|
|
key = f"blobs/{r['leaf_hash'][:2]}/{r['leaf_hash'][2:]}"
|
|
assert backend.head(key), f"missing blob for chunk {r['chunk_id']}"
|
|
|
|
# Hydrate metadata-only into M=4 fresh shards.
|
|
meta_hash = next(
|
|
p["pack_hash"] for p in push_result["packs"]
|
|
if p.get("kind") == "metadata"
|
|
)
|
|
target_dir = tmp_path / "tgt-jit"
|
|
target_dir.mkdir()
|
|
M = 4
|
|
targets = _make_target_shards(target_dir, M)
|
|
try:
|
|
hydrate_from_metadata_pack_routed(
|
|
targets, backend, meta_hash, M=M, mode="just-enough",
|
|
)
|
|
finally:
|
|
for t in targets:
|
|
t.close()
|
|
|
|
# Pick one doc; JIT-hydrate it on its owning target.
|
|
pick = chunks[0]
|
|
doc_root = pick["document_root"]
|
|
tidx = shard_for_document(doc_root, M)
|
|
tdb = target_dir / f"{tidx:03d}.db"
|
|
|
|
tconn = sqlite3.connect(str(tdb))
|
|
tconn.row_factory = sqlite3.Row
|
|
try:
|
|
n_fetched = hydrate_doc_jit(tconn, doc_root, backend)
|
|
assert n_fetched > 0, "JIT should fetch placeholder chunks"
|
|
|
|
tgt = {
|
|
r["leaf_hash"]: r["content"]
|
|
for r in tconn.execute(
|
|
"SELECT leaf_hash, content FROM chunks WHERE document_root = ?",
|
|
(doc_root,),
|
|
)
|
|
}
|
|
finally:
|
|
tconn.close()
|
|
|
|
prod = {
|
|
r["leaf_hash"]: r["content"]
|
|
for r in chunks
|
|
if r["document_root"] == doc_root
|
|
}
|
|
# Both sides must decode to identical text. The producer side
|
|
# may be compressed (BLOB) or plain (TEXT); the target side
|
|
# after JIT is always plain UTF-8 from the bucket blob.
|
|
for lh, prod_body in prod.items():
|
|
assert lh in tgt
|
|
assert unpack_chunk(tgt[lh]) == unpack_chunk(prod_body)
|
|
|
|
|
|
class TestRoutedHydrateValidation:
|
|
def test_M_mismatch_rejected(self, tmp_path: Path):
|
|
target_dir = tmp_path / "tgt"
|
|
target_dir.mkdir()
|
|
targets = _make_target_shards(target_dir, 4)
|
|
backend = MemoryBackend()
|
|
try:
|
|
with pytest.raises(ValueError, match="must match"):
|
|
hydrate_from_metadata_pack_routed(
|
|
targets, backend, "deadbeef" * 8, M=3, mode="just-enough",
|
|
)
|
|
finally:
|
|
for t in targets:
|
|
t.close()
|
|
|
|
def test_empty_targets_rejected(self, tmp_path: Path):
|
|
backend = MemoryBackend()
|
|
with pytest.raises(ValueError, match="at least one target"):
|
|
hydrate_from_metadata_pack_routed(
|
|
[], backend, "deadbeef" * 8, M=0, mode="just-enough",
|
|
)
|