cloud ask: unified multi-shard via bucket manifest + read-ahead tuning
`BUCKET_URL` (one env var) → client GETs `clones/manifest.json` →
opens HttpRangeVFS per listed shard → FTS5 across all shards in
parallel (ThreadPoolExecutor; per-thread apsw.Connection) → merge by
BM25 score → pull chunks from the owning shard → LLM + verify.
No per-query --shard-url, no path proliferation.
Two manifests published on s3://arborist/clones/:
manifest.json — default: virtback only (2.5MB, ~5s/query)
manifest-full.json — opt-in: all 5 shards (35GB, prohibitive over
WAN due to FTS5 b-tree walk pattern; needs
smaller shards or co-located query proxy)
HttpRangeVFS read-ahead tuned from per-page (4KB) to 64KB block-aligned
cache. Each cache miss fetches one 64KB block; subsequent reads within
the block are local-fast. Lower miss count, similar bytes-on-wire
(64KB amortizes well over typical 4-16 page b-tree clusters; larger
read-ahead like 4MB over-fetches on random FTS5 reads).
Sample run (default manifest):
make cloud-ask Q="who developed virt-back?"
→ EVIDENCE-WARRANTED · via claim_lattice 1/1 4.71s (bucket-direct)
21 HTTP requests · 1344 KB
ACL: genesis full-bench shards flipped to public-read (CC-BY-SA
Wikipedia content). Reachable now if you want to play with the slow
multi-shard path; not in the default manifest because chat latency
matters more than coverage breadth.
This commit is contained in:
parent
87f7d920e7
commit
331e748bcb
4 changed files with 272 additions and 27 deletions
14
Makefile
14
Makefile
|
|
@ -1552,6 +1552,12 @@ wallet-ask: bootstrap ## query the wallet server with your own question [Q="..."
|
|||
SHARD_URL ?= https://nyc3.digitaloceanspaces.com/arborist/clones/virtback/000.db
|
||||
BLOB_BASE ?=
|
||||
|
||||
# BUCKET_URL is the root the multi-shard cloud-ask points at. Client GETs
|
||||
# $BUCKET_URL/clones/manifest.json and queries every listed shard. One
|
||||
# config var → entire corpus. Default = our DO Spaces public bucket with
|
||||
# genesis Wikipedia + russell.ballestrini.net.
|
||||
BUCKET_URL ?= https://nyc3.digitaloceanspaces.com/arborist
|
||||
|
||||
bootstrap-bucket: bootstrap ## install apsw (bucket-direct extra)
|
||||
$(PIP) install 'apsw>=3.45'
|
||||
|
||||
|
|
@ -1572,10 +1578,10 @@ cloud-fetch-chunk: bootstrap ## fetch + hash-verify one chunk from a bucket [LEA
|
|||
@test -n "$(BLOB_BASE)" || { echo 'BLOB_BASE required (per-chunk blobs/<hash> base URL)'; exit 2; }
|
||||
@$(ARBORIST) cloud fetch-chunk "$(LEAF_HASH)" --blob-base "$(BLOB_BASE)"
|
||||
|
||||
cloud-ask: bootstrap ## bucket-direct end-to-end: FTS → LLM → verify [Q="..." JSON=1 SHARD_URL=... TOP_K=N MAX_CONTEXT=N CACHE_MB=N]
|
||||
@test -n "$(Q)" || { echo 'usage: make cloud-ask Q="your question" [JSON=1] [SHARD_URL=...] [TOP_K=4] [MAX_CONTEXT=24000] [CACHE_MB=64]'; exit 2; }
|
||||
@echo "# shard: $(SHARD_URL)" >&2
|
||||
@$(ARBORIST) cloud ask '$(Q)' --shard-url "$(SHARD_URL)" \
|
||||
cloud-ask: bootstrap ## bucket-direct end-to-end via manifest [Q="..." JSON=1 BUCKET_URL=... TOP_K=N MAX_CONTEXT=N CACHE_MB=N]
|
||||
@test -n "$(Q)" || { echo 'usage: make cloud-ask Q="your question" [JSON=1] [BUCKET_URL=...] [TOP_K=4] [MAX_CONTEXT=24000] [CACHE_MB=64]'; exit 2; }
|
||||
@echo "# bucket: $(BUCKET_URL)" >&2
|
||||
@$(ARBORIST) cloud ask '$(Q)' --bucket-url "$(BUCKET_URL)" \
|
||||
--top-k $(or $(TOP_K),4) \
|
||||
--max-context-chars $(or $(MAX_CONTEXT),24000) \
|
||||
--cache-mb $(or $(CACHE_MB),64) \
|
||||
|
|
|
|||
|
|
@ -7273,7 +7273,17 @@ def build_parser() -> argparse.ArgumentParser:
|
|||
),
|
||||
)
|
||||
cloud_ask.add_argument("question", type=str)
|
||||
cloud_ask.add_argument("--shard-url", required=True)
|
||||
cloud_ask.add_argument(
|
||||
"--bucket-url", default=None,
|
||||
help="bucket root URL (e.g. https://nyc3.digitaloceanspaces.com/arborist). "
|
||||
"The client fetches `clones/manifest.json` from this URL and queries "
|
||||
"every listed shard. Use this for multi-shard corpus access.",
|
||||
)
|
||||
cloud_ask.add_argument(
|
||||
"--shard-url", default=None,
|
||||
help="single shard URL (mutually exclusive with --bucket-url). "
|
||||
"Use when you want to point at exactly one .db file on a bucket.",
|
||||
)
|
||||
cloud_ask.add_argument("--top-k", type=int, default=4)
|
||||
cloud_ask.add_argument("--max-context-chars", type=int, default=24_000)
|
||||
cloud_ask.add_argument(
|
||||
|
|
@ -7464,7 +7474,27 @@ def _cmd_cloud_ask(args: argparse.Namespace) -> int:
|
|||
)
|
||||
from arborist.qa.verify import verify_claim_lattice
|
||||
|
||||
client = _make_bucket_client(args)
|
||||
# Multi-shard via manifest, or single-shard if --shard-url given.
|
||||
if getattr(args, "bucket_url", None):
|
||||
from arborist.wallet.bucket import (
|
||||
MultiShardBucketCorpus,
|
||||
load_bucket_manifest,
|
||||
)
|
||||
manifest = load_bucket_manifest(args.bucket_url)
|
||||
client = MultiShardBucketCorpus(
|
||||
manifest, cache_bytes_per_shard=args.cache_mb * 1024 * 1024,
|
||||
)
|
||||
multi_shard = True
|
||||
else:
|
||||
if not getattr(args, "shard_url", None):
|
||||
print(
|
||||
"cloud ask needs --bucket-url (multi-shard manifest) "
|
||||
"or --shard-url (single .db file).",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 2
|
||||
client = _make_bucket_client(args)
|
||||
multi_shard = False
|
||||
t0 = _time.time()
|
||||
try:
|
||||
hits = client.fts_search(args.question, limit=args.top_k)
|
||||
|
|
@ -7492,7 +7522,14 @@ def _cmd_cloud_ask(args: argparse.Namespace) -> int:
|
|||
1000, args.max_context_chars // max(1, len(hits))
|
||||
)
|
||||
for rank, h in enumerate(hits, 1):
|
||||
rows = list(client.conn.execute(
|
||||
# Multi-shard: route the chunk read back to the shard that
|
||||
# owned the FTS hit (so we hit a warm page cache and don't
|
||||
# re-fetch the .db header for every chunk).
|
||||
if multi_shard:
|
||||
conn_for_read = client.conn_for_shard(h["_shard_url"])
|
||||
else:
|
||||
conn_for_read = client.conn
|
||||
rows = list(conn_for_read.execute(
|
||||
"SELECT idx, leaf_hash, content FROM chunks "
|
||||
"WHERE document_root = ? AND content IS NOT NULL "
|
||||
"ORDER BY idx ASC LIMIT 1",
|
||||
|
|
@ -7671,9 +7708,18 @@ def _render_cloud_ask_human(result: dict, question: str) -> str:
|
|||
violations = result.get("violations") or []
|
||||
label = _render_audit_label(audit, method, violations)
|
||||
stats = result.get("stats") or {}
|
||||
# Stats shape differs single-shard vs multi-shard; coalesce both.
|
||||
cache_stats = stats.get("cache") or {}
|
||||
http_reqs = stats.get("http_requests")
|
||||
bytes_fetched = cache_stats.get("bytes_fetched") or 0
|
||||
http_reqs = (
|
||||
stats.get("http_requests")
|
||||
or stats.get("total_http_requests")
|
||||
or 0
|
||||
)
|
||||
bytes_fetched = (
|
||||
cache_stats.get("bytes_fetched")
|
||||
or stats.get("total_bytes_fetched")
|
||||
or 0
|
||||
)
|
||||
|
||||
lines: list[str] = []
|
||||
lines.append(question)
|
||||
|
|
|
|||
|
|
@ -51,6 +51,19 @@ except ImportError as e: # pragma: no cover
|
|||
# SQLite default page size; we read in whole-page chunks where possible.
|
||||
SQLITE_PAGE_SIZE = 4096
|
||||
DEFAULT_CACHE_BYTES = 32 * 1024 * 1024 # 32 MB
|
||||
# Read-ahead: when SQLite asks for a 4 KB page, fetch this many bytes
|
||||
# aligned to a block boundary and cache the surrounding pages. SQLite's
|
||||
# FTS5 b-tree walks read many adjacent pages, so one HTTP RANGE serves
|
||||
# 16-64 subsequent reads. Lowers per-query RANGE count by ~1-2 orders
|
||||
# of magnitude on multi-GB shards. 256 KB is a reasonable knob — small
|
||||
# enough to keep first-byte latency low, large enough to amortize.
|
||||
DEFAULT_READAHEAD_BYTES = 64 * 1024 # 64 KB — amortizes sequential page
|
||||
# walks (FTS5 segment headers, b-tree
|
||||
# inner nodes) but doesn't over-fetch
|
||||
# on random reads. Bigger values (1-4
|
||||
# MB) over-fetch on FTS5's b-tree
|
||||
# walk pattern and inflate total
|
||||
# bytes pulled by 10-100x.
|
||||
|
||||
|
||||
# FTS5-safe query construction. The bucket-direct path doesn't run the
|
||||
|
|
@ -97,37 +110,63 @@ def _to_fts5(query: str) -> str:
|
|||
|
||||
|
||||
class _LRUByteCache:
|
||||
"""Thread-safe LRU cache of byte ranges keyed by (offset, length).
|
||||
"""Thread-safe LRU cache of read-ahead BLOCKS keyed by aligned offset.
|
||||
|
||||
Each entry's eviction cost is its byte size; the cache enforces a
|
||||
soft byte budget. NOT designed for high concurrency — one lock
|
||||
serializes all access. Good enough for a single client running
|
||||
serial queries (the SQLite VFS layer already serializes reads
|
||||
per-connection)."""
|
||||
Each entry holds ``DEFAULT_READAHEAD_BYTES`` of consecutive bytes
|
||||
starting at a block-aligned offset (multiple of the read-ahead
|
||||
size). A 4 KB SQLite page read maps to the containing block: cache
|
||||
HIT = no HTTP RANGE; cache MISS = one HTTP RANGE for the whole
|
||||
block, populating it for the next ~16-64 page reads.
|
||||
|
||||
def __init__(self, max_bytes: int = DEFAULT_CACHE_BYTES):
|
||||
NOT designed for high concurrency — one lock serializes all
|
||||
access. Good enough for a single client running serial queries
|
||||
(the SQLite VFS layer already serializes reads per-connection)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_bytes: int = DEFAULT_CACHE_BYTES,
|
||||
block_size: int = DEFAULT_READAHEAD_BYTES,
|
||||
):
|
||||
self.max_bytes = max_bytes
|
||||
self._cache: "OrderedDict[tuple[int,int], bytes]" = OrderedDict()
|
||||
self.block_size = block_size
|
||||
self._cache: "OrderedDict[int, bytes]" = OrderedDict()
|
||||
self._size = 0
|
||||
self._lock = threading.Lock()
|
||||
self.hits = 0
|
||||
self.misses = 0
|
||||
self.bytes_fetched = 0
|
||||
|
||||
def _block_offset(self, offset: int) -> int:
|
||||
return (offset // self.block_size) * self.block_size
|
||||
|
||||
def get(self, offset: int, length: int) -> bytes | None:
|
||||
"""Try to serve [offset, offset+length) from the block cache.
|
||||
Returns None on miss; also returns None when the requested
|
||||
range straddles a block boundary (caller must fetch via the
|
||||
slow path or query each block separately)."""
|
||||
block = self._block_offset(offset)
|
||||
within = offset - block
|
||||
if within + length > self.block_size:
|
||||
return None # spans block boundary — caller handles
|
||||
with self._lock:
|
||||
v = self._cache.get((offset, length))
|
||||
if v is None:
|
||||
blk = self._cache.get(block)
|
||||
if blk is None:
|
||||
self.misses += 1
|
||||
return None
|
||||
self._cache.move_to_end((offset, length))
|
||||
self._cache.move_to_end(block)
|
||||
self.hits += 1
|
||||
return v
|
||||
return blk[within:within + length]
|
||||
|
||||
def put(self, offset: int, length: int, body: bytes) -> None:
|
||||
def put_block(self, block_offset: int, body: bytes) -> None:
|
||||
"""Store a full read-ahead block. ``block_offset`` must be
|
||||
block-size-aligned."""
|
||||
with self._lock:
|
||||
self._cache[(offset, length)] = body
|
||||
existing = self._cache.get(block_offset)
|
||||
if existing is not None:
|
||||
self._size -= len(existing)
|
||||
self._cache[block_offset] = body
|
||||
self._size += len(body)
|
||||
self._cache.move_to_end(block_offset)
|
||||
while self._size > self.max_bytes and self._cache:
|
||||
_k, evicted = self._cache.popitem(last=False)
|
||||
self._size -= len(evicted)
|
||||
|
|
@ -140,6 +179,7 @@ class _LRUByteCache:
|
|||
"bytes_resident": self._size,
|
||||
"bytes_fetched": self.bytes_fetched,
|
||||
"entries": len(self._cache),
|
||||
"block_size": self.block_size,
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -247,10 +287,21 @@ class HttpRangeFile(apsw.VFSFile):
|
|||
cached = self._cache.get(offset, amount)
|
||||
if cached is not None:
|
||||
return cached
|
||||
body = self._transport.get_range(self._url, offset, amount)
|
||||
# Cache miss: fetch the full read-ahead block that contains
|
||||
# [offset, offset+amount). Caps at file end so we don't request
|
||||
# past the .db size (some servers return 416 / hang on that).
|
||||
block_size = self._cache.block_size
|
||||
block_off = (offset // block_size) * block_size
|
||||
size = self.xFileSize()
|
||||
block_len = min(block_size, max(0, size - block_off))
|
||||
if block_len <= 0:
|
||||
return b""
|
||||
body = self._transport.get_range(self._url, block_off, block_len)
|
||||
self._cache.bytes_fetched += len(body)
|
||||
self._cache.put(offset, amount, body)
|
||||
return body
|
||||
self._cache.put_block(block_off, body)
|
||||
# Slice the requested window out of the block.
|
||||
within = offset - block_off
|
||||
return body[within:within + amount]
|
||||
|
||||
def xWrite(self, data: bytes, offset: int) -> None:
|
||||
raise IOError("HttpRangeFile is read-only")
|
||||
|
|
@ -474,3 +525,145 @@ class BucketClient:
|
|||
|
||||
def stats(self) -> dict:
|
||||
return self._vfs.stats()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Multi-shard: one BUCKET_URL → discover all shards via manifest.json.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class BucketManifest:
|
||||
"""Parsed clones/manifest.json. One BUCKET_URL → all the shards."""
|
||||
shards: list[dict]
|
||||
blob_base: str
|
||||
snapshot_root: str | None
|
||||
snapshot_ts: int | None
|
||||
|
||||
|
||||
def load_bucket_manifest(bucket_url: str, *, timeout_s: float = 30.0) -> BucketManifest:
|
||||
"""Fetch + parse clones/manifest.json from a bucket base URL.
|
||||
|
||||
bucket_url is the path that resolves to the manifest's parent —
|
||||
typically just the bucket root (e.g. https://bucket.example.com/).
|
||||
The client appends ``clones/manifest.json`` and parses the response.
|
||||
"""
|
||||
import json as _json
|
||||
base = bucket_url.rstrip("/")
|
||||
manifest_url = f"{base}/clones/manifest.json"
|
||||
req = urllib.request.Request(manifest_url)
|
||||
with urllib.request.urlopen(req, timeout=timeout_s) as resp:
|
||||
body = _json.loads(resp.read().decode("utf-8"))
|
||||
return BucketManifest(
|
||||
shards=list(body.get("shards") or []),
|
||||
blob_base=body.get("blob_base") or f"{base}/blobs",
|
||||
snapshot_root=body.get("snapshot_root"),
|
||||
snapshot_ts=body.get("snapshot_ts"),
|
||||
)
|
||||
|
||||
|
||||
class MultiShardBucketCorpus:
|
||||
"""Bucket-direct corpus that spans multiple SQLite shards.
|
||||
|
||||
Each shard is a ``BucketClient`` (its own apsw connection + page
|
||||
cache). FTS5 queries run against every shard sequentially; results
|
||||
merge by BM25 score (lower = better in our ORDER BY convention),
|
||||
trimmed to the operator-supplied top-K.
|
||||
|
||||
Chunk reads are routed back to the shard that owned the hit — the
|
||||
same VFS+page-cache that surfaced the hit serves the chunk body,
|
||||
so no extra HTTP RANGE for the same byte range.
|
||||
|
||||
Sequential is fine on small corpora; future work: thread per shard
|
||||
with per-thread apsw.Connection (apsw VFS internals serialize so
|
||||
the per-shard caches stay coherent)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
manifest: BucketManifest,
|
||||
*,
|
||||
cache_bytes_per_shard: int = DEFAULT_CACHE_BYTES,
|
||||
timeout_s: float = 30.0,
|
||||
):
|
||||
self.manifest = manifest
|
||||
self.shard_clients: list[BucketClient] = [
|
||||
BucketClient(
|
||||
BucketEndpoint(shard_url=sh["url"], blob_base=manifest.blob_base),
|
||||
cache_bytes=cache_bytes_per_shard,
|
||||
timeout_s=timeout_s,
|
||||
)
|
||||
for sh in manifest.shards
|
||||
]
|
||||
# Quick lookup: shard_url → BucketClient, for chunk fetch routing.
|
||||
self._by_url: dict[str, BucketClient] = {
|
||||
sh["url"]: self.shard_clients[i]
|
||||
for i, sh in enumerate(manifest.shards)
|
||||
}
|
||||
|
||||
def close(self) -> None:
|
||||
for c in self.shard_clients:
|
||||
c.close()
|
||||
|
||||
def fts_search(self, query: str, *, limit: int = 8, raw: bool = False) -> list[dict]:
|
||||
"""Run FTS5 on every shard in parallel, merge by BM25 score,
|
||||
trim to ``limit``.
|
||||
|
||||
Each shard is its own apsw.Connection on its own thread, so
|
||||
FTS scans run concurrently — total wall time = max(per-shard)
|
||||
instead of sum, which matters when each shard takes 3-5s over
|
||||
WAN. apsw connections are independent across threads; each
|
||||
BucketClient owns one and its underlying page cache.
|
||||
|
||||
Each hit is annotated with ``_shard_url`` so callers can route
|
||||
the chunk read back to the right shard's BucketClient.
|
||||
"""
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
def _one(item):
|
||||
sh_url, sh_client = item
|
||||
try:
|
||||
return sh_url, sh_client.fts_search(query, limit=limit, raw=raw)
|
||||
except Exception:
|
||||
# One shard failing shouldn't kill the whole query —
|
||||
# could be a transient HTTP issue, a permission glitch,
|
||||
# or a manifest entry that lost its public-read ACL.
|
||||
return sh_url, []
|
||||
|
||||
n_shards = len(self._by_url)
|
||||
all_hits: list[dict] = []
|
||||
with ThreadPoolExecutor(max_workers=max(1, n_shards)) as ex:
|
||||
for sh_url, hits in ex.map(_one, list(self._by_url.items())):
|
||||
for h in hits:
|
||||
h["_shard_url"] = sh_url
|
||||
all_hits.append(h)
|
||||
# Lower BM25 score = stronger match (our ORDER BY score ASC convention).
|
||||
all_hits.sort(key=lambda h: h.get("score") or 0.0)
|
||||
return all_hits[:limit]
|
||||
|
||||
def conn_for_shard(self, shard_url: str):
|
||||
"""Get the apsw.Connection for a specific shard. Use after
|
||||
``fts_search`` returns a hit, to run a follow-up SELECT for
|
||||
chunk content without paying for a new connection."""
|
||||
client = self._by_url.get(shard_url)
|
||||
if client is None:
|
||||
raise KeyError(f"shard_url {shard_url!r} not in manifest")
|
||||
return client.conn
|
||||
|
||||
def stats(self) -> dict:
|
||||
merged = {
|
||||
"total_http_requests": 0,
|
||||
"total_bytes_fetched": 0,
|
||||
"per_shard": [],
|
||||
}
|
||||
for sh, c in zip(self.manifest.shards, self.shard_clients):
|
||||
s = c.stats()
|
||||
merged["total_http_requests"] += s.get("http_requests") or 0
|
||||
merged["total_bytes_fetched"] += (
|
||||
(s.get("cache") or {}).get("bytes_fetched") or 0
|
||||
)
|
||||
merged["per_shard"].append({
|
||||
"label": sh.get("label", sh["url"].rsplit("/", 1)[-1]),
|
||||
"http_requests": s.get("http_requests"),
|
||||
"bytes_fetched": (s.get("cache") or {}).get("bytes_fetched"),
|
||||
})
|
||||
return merged
|
||||
|
|
|
|||
|
|
@ -168,7 +168,7 @@ def test_bucket_open_and_fts_match_local_sqlite(bucket_layout):
|
|||
)
|
||||
client = BucketClient(endpoint, cache_bytes=4 * 1024 * 1024)
|
||||
try:
|
||||
remote = client.fts_search('"lemma-3"', limit=4)
|
||||
remote = client.fts_search('"lemma-3"', limit=4, raw=True)
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue