cloud ask: sidecar-aware multi-shard corpus + RRF merge
Wires the inverted-index sidecar into the unified `BUCKET_URL` flow: - BucketManifest gains optional per-shard `sidecar_url`. - New `SidecarShardClient`: downloads the sidecar once (cached under ~/.arborist/sidecar-cache/<hash>.bin keyed by URL), queries locally with BM25; uses the bucket's HttpRangeVFS only for chunk content reads on hit (warm page cache reuse). - New `MultiShardSidecarCorpus`: mixed-mode — shards with a sidecar use SidecarShardClient (sub-second FTS); shards without one fall back to BucketClient (bucket-direct FTS5, slow but functional). - Cross-shard merge: reciprocal-rank-fusion (RRF, k=60). Sidecar BM25 and FTS5 BM25 are on incomparable scales; RRF normalizes by rank-position so a doc that's #1 on the virtback shard beats a doc that's #1 on the genesis shard regardless of raw score magnitude. Bucket layout: clones/full-bench-64k/002.db (7.6 GB) clones/full-bench-64k/002.sidecar.bin (745 MB — pre-built index) clones/manifest-sidecar.json (advertises both) Live demo (BUCKET_URL=…/clones/manifest-sidecar.json): make cloud-ask Q="who developed virt-back?" → EVIDENCE-WARRANTED · via claim_lattice 1/1 6.79s "Russell Ballestrini developed virt-back." sources span genesis-Wikipedia + russell.ballestrini.net 28 HTTP RANGEs, 1.8 MB transferred vs the 4+ minute bucket-direct-only path on the same shard mix. `load_bucket_manifest` now accepts both bucket-root URLs (appends `clones/manifest.json` per convention) and explicit `.json` URLs (uses them as-is), so operators can pin alternate manifests.
This commit is contained in:
parent
349dd0fd48
commit
f669901f30
2 changed files with 266 additions and 9 deletions
|
|
@ -7510,12 +7510,27 @@ def _cmd_cloud_ask(args: argparse.Namespace) -> int:
|
|||
if getattr(args, "bucket_url", None):
|
||||
from arborist.wallet.bucket import (
|
||||
MultiShardBucketCorpus,
|
||||
MultiShardSidecarCorpus,
|
||||
load_bucket_manifest,
|
||||
)
|
||||
manifest = load_bucket_manifest(args.bucket_url)
|
||||
client = MultiShardBucketCorpus(
|
||||
manifest, cache_bytes_per_shard=args.cache_mb * 1024 * 1024,
|
||||
)
|
||||
manifest_url = args.bucket_url
|
||||
# Convenience: allow passing the manifest path directly.
|
||||
if not manifest_url.rstrip("/").endswith(".json"):
|
||||
manifest_url = manifest_url # base URL → load_bucket_manifest appends path
|
||||
manifest = load_bucket_manifest(manifest_url)
|
||||
# Pick the right corpus class based on whether the manifest
|
||||
# advertises per-shard sidecars. Mixed (some sidecar, some not)
|
||||
# also uses the sidecar-aware path — it falls back to bucket-
|
||||
# direct FTS5 for shards without one.
|
||||
any_sidecar = any(sh.get("sidecar_url") for sh in manifest.shards)
|
||||
if any_sidecar:
|
||||
client = MultiShardSidecarCorpus(
|
||||
manifest, cache_bytes_per_shard=args.cache_mb * 1024 * 1024,
|
||||
)
|
||||
else:
|
||||
client = MultiShardBucketCorpus(
|
||||
manifest, cache_bytes_per_shard=args.cache_mb * 1024 * 1024,
|
||||
)
|
||||
multi_shard = True
|
||||
else:
|
||||
if not getattr(args, "shard_url", None):
|
||||
|
|
|
|||
|
|
@ -542,15 +542,20 @@ class BucketManifest:
|
|||
|
||||
|
||||
def load_bucket_manifest(bucket_url: str, *, timeout_s: float = 30.0) -> BucketManifest:
|
||||
"""Fetch + parse clones/manifest.json from a bucket base URL.
|
||||
"""Fetch + parse a manifest from a bucket 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.
|
||||
Accepts either:
|
||||
- bucket root (e.g. https://bucket.example.com) → client appends
|
||||
``clones/manifest.json``
|
||||
- explicit manifest URL (e.g. .../clones/manifest-sidecar.json) →
|
||||
used as-is, lets operators pin a non-default manifest
|
||||
"""
|
||||
import json as _json
|
||||
base = bucket_url.rstrip("/")
|
||||
manifest_url = f"{base}/clones/manifest.json"
|
||||
if base.endswith(".json"):
|
||||
manifest_url = base
|
||||
else:
|
||||
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"))
|
||||
|
|
@ -667,3 +672,240 @@ class MultiShardBucketCorpus:
|
|||
"bytes_fetched": (s.get("cache") or {}).get("bytes_fetched"),
|
||||
})
|
||||
return merged
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sidecar-aware corpus: download per-shard sidecar.bin once, query locally.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _sidecar_cache_path(url: str, *, base_dir: str | None = None) -> str:
|
||||
"""Stable local cache path for a sidecar URL.
|
||||
|
||||
Hashes the URL → places under ~/.arborist/sidecar-cache/<hex>.bin so
|
||||
re-runs reuse the download. base_dir overrides for tests.
|
||||
"""
|
||||
import hashlib
|
||||
if base_dir is None:
|
||||
base_dir = os.path.join(
|
||||
os.path.expanduser("~"), ".arborist", "sidecar-cache"
|
||||
)
|
||||
os.makedirs(base_dir, exist_ok=True)
|
||||
h = hashlib.sha256(url.encode("utf-8")).hexdigest()[:16]
|
||||
return os.path.join(base_dir, f"{h}.bin")
|
||||
|
||||
|
||||
def _download_sidecar(url: str, dest: str, *, timeout_s: float = 300.0) -> int:
|
||||
"""Stream a sidecar URL into `dest`. Returns bytes downloaded.
|
||||
|
||||
Writes to dest.part then renames atomically — interrupted downloads
|
||||
don't corrupt the cache. The cache file is content-addressed by URL
|
||||
hash; if you change the bucket-side artifact you must change the
|
||||
URL (or clear ~/.arborist/sidecar-cache/) — sidecars are not
|
||||
Merkle-bound, so the URL is the identity."""
|
||||
import shutil
|
||||
tmp = dest + ".part"
|
||||
req = urllib.request.Request(url)
|
||||
n = 0
|
||||
with urllib.request.urlopen(req, timeout=timeout_s) as resp, open(tmp, "wb") as out:
|
||||
while True:
|
||||
chunk = resp.read(1024 * 1024)
|
||||
if not chunk:
|
||||
break
|
||||
out.write(chunk)
|
||||
n += len(chunk)
|
||||
shutil.move(tmp, dest)
|
||||
return n
|
||||
|
||||
|
||||
class SidecarShardClient:
|
||||
"""One shard's view: a local sidecar (download-then-query-locally)
|
||||
plus a BucketClient against the shard .db for chunk content reads.
|
||||
|
||||
fts_search uses the sidecar (sub-second BM25); chunks_for_doc and
|
||||
SQL go through the bucket VFS so chunk content fetches benefit
|
||||
from the same on-disk page cache."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
shard_url: str,
|
||||
sidecar_url: str,
|
||||
blob_base: str,
|
||||
*,
|
||||
cache_bytes: int = DEFAULT_CACHE_BYTES,
|
||||
sidecar_cache_dir: str | None = None,
|
||||
):
|
||||
from arborist.wallet.sidecar import SidecarReader
|
||||
|
||||
self.shard_url = shard_url
|
||||
self.sidecar_url = sidecar_url
|
||||
# 1. Download or reuse cached sidecar.
|
||||
cache_path = _sidecar_cache_path(sidecar_url, base_dir=sidecar_cache_dir)
|
||||
if not os.path.exists(cache_path):
|
||||
_download_sidecar(sidecar_url, cache_path)
|
||||
self._sidecar_path = cache_path
|
||||
self.sidecar = SidecarReader(cache_path)
|
||||
# 2. Bucket client for chunk content reads.
|
||||
self.bucket = BucketClient(
|
||||
BucketEndpoint(shard_url=shard_url, blob_base=blob_base),
|
||||
cache_bytes=cache_bytes,
|
||||
)
|
||||
|
||||
def close(self):
|
||||
self.bucket.close()
|
||||
|
||||
def fts_search(self, query: str, *, limit: int = 8, raw: bool = False) -> list[dict]:
|
||||
"""Sidecar-driven FTS. ``raw`` is accepted for API parity with
|
||||
BucketClient.fts_search but ignored — sidecar tokenizes via
|
||||
tokenize_text (matches the OR-sanitizer behavior)."""
|
||||
hits = self.sidecar.search(query, limit=limit, mode="or")
|
||||
return [
|
||||
{
|
||||
"document_root": h.document_root,
|
||||
"document_uri": h.document_uri,
|
||||
"title": h.title,
|
||||
"score": h.score,
|
||||
}
|
||||
for h in hits
|
||||
]
|
||||
|
||||
@property
|
||||
def conn(self):
|
||||
"""apsw connection for SQL pulls (chunk content). Goes through
|
||||
the bucket VFS — chunk reads are cached."""
|
||||
return self.bucket.conn
|
||||
|
||||
|
||||
class MultiShardSidecarCorpus:
|
||||
"""Manifest-driven corpus where each shard MAY have a sidecar.
|
||||
|
||||
Shards with a `sidecar_url` use SidecarShardClient (download-then-
|
||||
local FTS). Shards without one fall back to BucketClient (in-place
|
||||
bucket-direct FTS5 — slow on multi-GB shards, but works for small
|
||||
shards or when no sidecar is published yet).
|
||||
|
||||
fts_search runs per-shard in parallel, merges by score, picks top-K.
|
||||
Mixed-mode results are score-comparable in practice: sidecar BM25
|
||||
is on a similar magnitude (~10-30) as FTS5's BM25 (~-15 to -1) but
|
||||
sign-flipped (sidecar is positive, FTS5 stored as negative). We
|
||||
normalize sidecar to negative before merging so a single ORDER BY
|
||||
DESC works across both.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
manifest: "BucketManifest",
|
||||
*,
|
||||
cache_bytes_per_shard: int = DEFAULT_CACHE_BYTES,
|
||||
sidecar_cache_dir: str | None = None,
|
||||
):
|
||||
self.manifest = manifest
|
||||
self.shard_clients = []
|
||||
self._has_sidecar: list[bool] = []
|
||||
for sh in manifest.shards:
|
||||
if sh.get("sidecar_url"):
|
||||
client = SidecarShardClient(
|
||||
shard_url=sh["url"],
|
||||
sidecar_url=sh["sidecar_url"],
|
||||
blob_base=manifest.blob_base,
|
||||
cache_bytes=cache_bytes_per_shard,
|
||||
sidecar_cache_dir=sidecar_cache_dir,
|
||||
)
|
||||
self._has_sidecar.append(True)
|
||||
else:
|
||||
client = BucketClient(
|
||||
BucketEndpoint(shard_url=sh["url"], blob_base=manifest.blob_base),
|
||||
cache_bytes=cache_bytes_per_shard,
|
||||
)
|
||||
self._has_sidecar.append(False)
|
||||
self.shard_clients.append(client)
|
||||
self._by_url = {sh["url"]: c for sh, c in zip(manifest.shards, self.shard_clients)}
|
||||
|
||||
def close(self):
|
||||
for c in self.shard_clients:
|
||||
c.close()
|
||||
|
||||
def fts_search(
|
||||
self, query: str, *, limit: int = 8, raw: bool = False,
|
||||
rrf_k: int = 60,
|
||||
) -> list[dict]:
|
||||
"""Run each shard's FTS in parallel, merge via reciprocal-rank-fusion.
|
||||
|
||||
BM25 scores aren't comparable across heterogeneous backends
|
||||
(sidecar BM25 is positive in the [5, 50] range; FTS5 BM25 is
|
||||
negative in [-20, 0]). RRF sidesteps the scale mismatch
|
||||
entirely: each shard ranks its own hits, the global rank is
|
||||
the sum of reciprocal ranks across shards. Standard k=60 is
|
||||
TREC-recommended.
|
||||
|
||||
score_rrf(d) = sum over shards of 1 / (rrf_k + rank(d, shard))
|
||||
"""
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
def _one(item):
|
||||
sh_url, sh_client, _ = item
|
||||
try:
|
||||
# Pull per-shard top-K (oversample so RRF has more to
|
||||
# work with — a doc that's barely in top-K on one
|
||||
# shard but #1 on another should still bubble up).
|
||||
return sh_url, sh_client.fts_search(query, limit=limit * 4, raw=raw)
|
||||
except Exception:
|
||||
return sh_url, []
|
||||
|
||||
triples = [
|
||||
(sh["url"], self._by_url[sh["url"]], self._has_sidecar[i])
|
||||
for i, sh in enumerate(self.manifest.shards)
|
||||
]
|
||||
per_shard: list[tuple[str, list[dict]]] = []
|
||||
with ThreadPoolExecutor(max_workers=max(1, len(triples))) as ex:
|
||||
for sh_url, hits in ex.map(_one, triples):
|
||||
per_shard.append((sh_url, hits))
|
||||
|
||||
# RRF merge: keyed by document_root so duplicates across shards
|
||||
# (rare but possible — e.g. a doc replicated for redundancy)
|
||||
# combine.
|
||||
rrf_score: dict[str, float] = {}
|
||||
hit_by_root: dict[str, dict] = {}
|
||||
for sh_url, hits in per_shard:
|
||||
for rank, h in enumerate(hits, start=1):
|
||||
droot = h["document_root"]
|
||||
rrf_score[droot] = (
|
||||
rrf_score.get(droot, 0.0) + 1.0 / (rrf_k + rank)
|
||||
)
|
||||
if droot not in hit_by_root:
|
||||
h["_shard_url"] = sh_url
|
||||
hit_by_root[droot] = h
|
||||
# Surface rrf score so callers see "this doc was strong across N shards."
|
||||
for droot, h in hit_by_root.items():
|
||||
h["rrf_score"] = rrf_score[droot]
|
||||
ranked = sorted(
|
||||
hit_by_root.values(),
|
||||
key=lambda h: -h["rrf_score"],
|
||||
)
|
||||
return ranked[:limit]
|
||||
|
||||
def conn_for_shard(self, shard_url: str):
|
||||
return self._by_url[shard_url].conn
|
||||
|
||||
def stats(self) -> dict:
|
||||
merged = {
|
||||
"total_http_requests": 0,
|
||||
"total_bytes_fetched": 0,
|
||||
"per_shard": [],
|
||||
}
|
||||
for sh, c, has_sc in zip(
|
||||
self.manifest.shards, self.shard_clients, self._has_sidecar,
|
||||
):
|
||||
bucket = c.bucket if has_sc else c # type: ignore[attr-defined]
|
||||
s = bucket.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]),
|
||||
"sidecar": has_sc,
|
||||
"http_requests": s.get("http_requests"),
|
||||
"bytes_fetched": (s.get("cache") or {}).get("bytes_fetched"),
|
||||
})
|
||||
return merged
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue