from __future__ import annotations
import hashlib
import importlib.metadata
import json
import os
import platform
import subprocess
from datetime import datetime, timezone
from pathlib import Path
def add_protocol_arguments(parser):
parser.add_argument(
"--partition", choices=["dev", "test", "exploratory"], default="dev"
)
parser.add_argument("--split-seed", type=int, default=20260924)
parser.add_argument(
"--freeze-config",
type=Path,
help="write one selected dev operating point and exit without evaluation",
)
parser.add_argument(
"--frozen-config",
type=Path,
help="required for test; all settings and corpus content must match",
)
def _digest(value):
return hashlib.sha256(
json.dumps(
value,
sort_keys=True,
separators=(",", ":"),
ensure_ascii=False,
allow_nan=False,
).encode()
).hexdigest()
def slice_digest(docs, queries, qrels):
h = hashlib.sha256()
for d in sorted(docs, key=lambda x: x.doc_id):
h.update(
bytes.fromhex(_digest(["doc", d.doc_id, getattr(d, "title", ""), d.text]))
)
for q in sorted(queries, key=lambda x: x.query_id):
h.update(
bytes.fromhex(
_digest(["query", q.query_id, q.text, qrels.get(q.query_id, {})])
)
)
return h.hexdigest()
def query_split(queries, seed):
ids = [q.query_id for q in queries]
if len(ids) != len(set(ids)):
raise ValueError("duplicate query IDs in benchmark slice")
ordered = sorted(ids, key=lambda qid: (_digest([seed, qid]), qid))
cut = len(ordered) // 2
return sorted(ordered[:cut]), sorted(ordered[cut:])
def file_digest(path):
h = hashlib.sha256()
with Path(path).open("rb") as f:
for block in iter(lambda: f.read(1024 * 1024), b""):
h.update(block)
return h.hexdigest()
def source_digest(workspace):
h = hashlib.sha256()
paths = sorted((workspace / "crates").glob("*/src/**/*.rs"))
paths += sorted((workspace / "crates").glob("*/Cargo.toml"))
paths += sorted((workspace / "crates/annex-multivector/benchmark").glob("*.py"))
paths += [workspace / "Cargo.lock", workspace / "Cargo.toml"]
for path in paths:
h.update(str(path.relative_to(workspace)).encode() + b"\0")
h.update(bytes.fromhex(file_digest(path)))
return h.hexdigest()
def _linux_cpu_info(path=Path("/proc/cpuinfo")):
try:
fields = dict(
line.split(":", 1)
for line in path.read_text().split("\n\n", 1)[0].splitlines()
if ":" in line
)
except OSError:
return {}
fields = {key.strip(): value.strip() for key, value in fields.items()}
flags = fields.get("flags", fields.get("Features", "")).split()
return {
"processor": fields.get(
"model name", fields.get("Processor", platform.processor())
),
"cpu_features": sorted(set(flags)),
}
def environment_settings():
packages = {}
for name in ("numpy", "ir-datasets", "pylate", "torch", "qdrant-client", "lancedb"):
try:
packages[name] = importlib.metadata.version(name)
except importlib.metadata.PackageNotFoundError:
packages[name] = None
hardware = {"processor": platform.processor()}
if platform.system() == "Darwin":
hardware.update(
processor=subprocess.check_output(
["sysctl", "-n", "machdep.cpu.brand_string"], text=True
).strip(),
memory_bytes=int(
subprocess.check_output(["sysctl", "-n", "hw.memsize"], text=True)
),
)
elif platform.system() == "Linux":
hardware.update(_linux_cpu_info())
if hasattr(os, "sysconf"):
hardware["memory_bytes"] = os.sysconf("SC_PAGE_SIZE") * os.sysconf(
"SC_PHYS_PAGES"
)
return {
"hardware": hardware,
"packages": packages,
"python": platform.python_version(),
"machine": platform.machine(),
"threads": {
k: os.environ.get(k)
for k in ("RAYON_NUM_THREADS", "OMP_NUM_THREADS", "MKL_NUM_THREADS")
},
"rustflags": os.environ.get("RUSTFLAGS"),
"os": platform.platform(),
"cpu_count": os.cpu_count(),
"runtime_controls": {
k: v
for k, v in sorted(os.environ.items())
if k.startswith(("VECTORDB_", "ANNEX_"))
},
}
def validate_protocol_args(args, operating_points):
if args.freeze_config and args.frozen_config:
raise ValueError("choose --freeze-config or --frozen-config")
if args.partition == "test" and not args.frozen_config:
raise ValueError("test requires --frozen-config; tune only on dev")
if args.partition == "test" and operating_points != 1:
raise ValueError("test evaluation cannot sweep operating points")
if args.freeze_config and (args.partition != "dev" or operating_points != 1):
raise ValueError("freeze exactly one operating point from dev")
if args.frozen_config and args.partition != "test":
raise ValueError("--frozen-config is used only with --partition test")
def prepare_protocol(args, docs, queries, qrels, settings, *, operating_points):
if not docs or not queries:
raise ValueError("benchmark slice must contain documents and evaluable queries")
if len({d.doc_id for d in docs}) != len(docs):
raise ValueError("duplicate document IDs in benchmark slice")
if set(qrels) != {q.query_id for q in queries}:
raise ValueError("every benchmark query must have retained qrels")
validate_protocol_args(args, operating_points)
dev_ids, test_ids = query_split(queries, args.split_seed)
if args.partition != "exploratory" and (not dev_ids or not test_ids):
raise ValueError("need at least two evaluable queries for disjoint dev/test")
settings = json.loads(json.dumps(settings, allow_nan=False))
contract = {
"protocol_version": 1,
"split_seed": args.split_seed,
"slice": {
"dataset": args.dataset,
"sampling": args.sampling,
"sample_seed": args.sample_seed,
"limit_docs": args.limit_docs,
"limit_queries": args.limit_queries,
"documents": len(docs),
"evaluable_queries_before_split": len(queries),
"content_sha256": slice_digest(docs, queries, qrels),
},
"dev_query_ids": dev_ids,
"test_query_ids": test_ids,
"settings": settings,
}
contract_hash = _digest(contract)
if args.freeze_config:
artifact = {
"created_at": datetime.now(timezone.utc).isoformat(),
"contract_sha256": contract_hash,
"contract": contract,
}
args.freeze_config.parent.mkdir(parents=True, exist_ok=True)
with args.freeze_config.open("x", encoding="utf-8") as f:
json.dump(artifact, f, sort_keys=True, indent=2, allow_nan=False)
f.write("\n")
if args.frozen_config:
artifact = json.loads(args.frozen_config.read_text())
if artifact.get("contract_sha256") != _digest(artifact.get("contract")):
raise ValueError("frozen configuration digest mismatch")
if artifact["contract"] != contract:
raise ValueError(
"test settings, source, dataset content, or split differ from frozen dev configuration"
)
ids = set(dev_ids if args.partition == "dev" else test_ids)
selected = (
list(queries)
if args.partition == "exploratory"
else [q for q in queries if q.query_id in ids]
)
selected_qrels = {q.query_id: qrels[q.query_id] for q in selected}
return (
selected,
selected_qrels,
{
**contract,
"partition": args.partition,
"contract_sha256": contract_hash,
"evaluated_query_ids": [q.query_id for q in selected],
"frozen_config": str(args.frozen_config) if args.frozen_config else None,
"label": "held-out frozen-config test"
if args.partition == "test"
else "development exploration",
},
)