Skip to content

Add recall harness CI progress output #175

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ajroetker merged 2 commits into from
Jun 5, 2026
Merged

Add recall harness CI progress output #175

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
690c5e4
Add recall harness CI progress output
ajroetker Jun 5, 2026
62c497a
Harden resolution e2e deadline diagnostics
ajroetker Jun 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
123 changes: 110 additions & 13 deletions zig/bench/vectors/recall_harness.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,22 @@ const Config = struct {
per_query_only: bool = false,
};

const Heartbeat = struct {
stop: std.atomic.Value(bool) = .init(false),

fn run(self: *Heartbeat) void {
var elapsed_ms: u64 = 0;
while (!self.stop.load(.acquire)) {
sleepMs(1_000);
elapsed_ms += 1_000;
if (elapsed_ms >= 30_000 and !self.stop.load(.acquire)) {
std.debug.print("recall_harness_heartbeat alive\n", .{});
elapsed_ms = 0;
}
}
}
};

const RecallCase = recall_cases.RecallCase;

pub fn main(init: std.process.Init) !void {
Expand All @@ -55,6 +71,13 @@ pub fn main(init: std.process.Init) !void {
const alloc = gpa_state.allocator();

const cfg = try parseArgs(init.minimal.args);
var heartbeat = Heartbeat{};
const heartbeat_thread = try std.Thread.spawn(.{ .stack_size = 256 * 1024 }, Heartbeat.run, .{&heartbeat});
defer {
heartbeat.stop.store(true, .release);
heartbeat_thread.join();
}

var ok = true;
if (cfg.per_query_metric != null and cfg.per_query_only) {
_ = try runSuite(init.io, alloc, cfg, "hbc-per-query", &recall_cases.hbc_cases, runHBCCase);
Expand Down Expand Up @@ -141,6 +164,10 @@ fn runSuite(
const dataset_path = try joinConvertedDatasetPath(alloc, cfg.dataset_dir, case.dataset);
defer alloc.free(dataset_path);

std.debug.print(
"recall_harness_case_start suite={s} dataset={s} randomize={any} count={d} topk={d}\n",
.{ suite_name, case.dataset, case.randomize, case.count, case.top_k },
);
const actual = try runner(io, alloc, cfg, dataset_path, case);
const case_ok = compareMetrics(case, actual);
all_ok = all_ok and case_ok;
Expand Down Expand Up @@ -191,31 +218,46 @@ fn convertedDatasetName(alloc: std.mem.Allocator, gob_name: []const u8) ![]u8 {

fn runQuantizerCase(io: std.Io, alloc: std.mem.Allocator, cfg: Config, dataset_path: []const u8, case: RecallCase) !common.MetricStats {
_ = cfg;
std.debug.print("recall_harness_load_start suite=quantizer dataset={s} path={s}\n", .{ case.dataset, dataset_path });
var loaded = try common.loadVectorSet(io, alloc, dataset_path);
defer loaded.deinit(alloc);
std.debug.print(
"recall_harness_load_done suite=quantizer dataset={s} dims={d} count={d}\n",
.{ case.dataset, loaded.dims, loaded.count },
);

var working = try common.cloneSet(alloc, loaded.asSet());
defer working.deinit(alloc);

const split = try common.splitDataset(working.asSet(), case.count);
std.debug.print(
"recall_harness_split suite=quantizer dataset={s} data={d} queries={d}\n",
.{ case.dataset, split.data.count, split.queries.count },
);
if (case.randomize) {
std.debug.print("recall_harness_randomize_start suite=quantizer dataset={s}\n", .{case.dataset});
try common.applyRandomTransformInPlace(alloc, split.data, 42);
try common.applyRandomTransformInPlace(alloc, split.queries, 42);
std.debug.print("recall_harness_randomize_done suite=quantizer dataset={s}\n", .{case.dataset});
}

return .{
.euclidean = 100.0 * try calculateQuantizerRecallMetric(alloc, split, case.top_k, .l2_squared),
.inner_product = 100.0 * try calculateQuantizerRecallMetric(alloc, split, case.top_k, .inner_product),
.cosine = 100.0 * try calculateQuantizerRecallMetric(alloc, split, case.top_k, .cosine),
};
const euclidean = 100.0 * try calculateQuantizerRecallMetric(alloc, case.dataset, split, case.top_k, .l2_squared);
const inner_product = 100.0 * try calculateQuantizerRecallMetric(alloc, case.dataset, split, case.top_k, .inner_product);
const cosine = 100.0 * try calculateQuantizerRecallMetric(alloc, case.dataset, split, case.top_k, .cosine);
return .{ .euclidean = euclidean, .inner_product = inner_product, .cosine = cosine };
}

fn calculateQuantizerRecallMetric(
alloc: std.mem.Allocator,
dataset_label: []const u8,
split: common.SplitDataset,
top_k: usize,
metric: vec.DistanceMetric,
) !f64 {
std.debug.print(
"recall_harness_metric_start suite=quantizer dataset={s} metric={s} data={d} queries={d}\n",
.{ dataset_label, @tagName(metric), split.data.count, split.queries.count },
);
var data_owned = try common.cloneSet(alloc, split.data);
defer data_owned.deinit(alloc);
var query_owned = try common.cloneSet(alloc, split.queries);
Expand Down Expand Up @@ -249,6 +291,13 @@ fn calculateQuantizerRecallMetric(

var recall_sum: f64 = 0;
for (0..queries.count) |query_idx| {
const query_number = query_idx + 1;
if (shouldLogQueryProgress(query_number, queries.count)) {
std.debug.print(
"recall_harness_query_progress suite=quantizer dataset={s} metric={s} query={d}/{d}\n",
.{ dataset_label, @tagName(metric), query_number, queries.count },
);
}
const query = queries.atConst(query_idx);
try quantizer.estimateDistancesWithScratch(&quantized, query, estimated, error_bounds, &scratch);
for (prediction, 0..) |*slot, i| slot.* = i;
Expand All @@ -258,23 +307,36 @@ fn calculateQuantizerRecallMetric(
defer alloc.free(truth);
recall_sum += common.calculateRecall(prediction[0..top_k], truth);
}
std.debug.print(
"recall_harness_metric_done suite=quantizer dataset={s} metric={s} recall={d:.4}\n",
.{ dataset_label, @tagName(metric), recall_sum / @as(f64, @floatFromInt(queries.count)) },
);
return recall_sum / @as(f64, @floatFromInt(queries.count));
}

fn runHBCCase(io: std.Io, alloc: std.mem.Allocator, cfg: Config, dataset_path: []const u8, case: RecallCase) !common.MetricStats {
std.debug.print("recall_harness_load_start suite=hbc dataset={s} path={s}\n", .{ case.dataset, dataset_path });
var loaded = try common.loadVectorSet(io, alloc, dataset_path);
defer loaded.deinit(alloc);
std.debug.print(
"recall_harness_load_done suite=hbc dataset={s} dims={d} count={d}\n",
.{ case.dataset, loaded.dims, loaded.count },
);

const loaded_set = loaded.asSet();
const split = try common.splitDataset(loaded_set, case.count);
std.debug.print(
"recall_harness_split suite=hbc dataset={s} data={d} queries={d}\n",
.{ case.dataset, split.data.count, split.queries.count },
);

if (cfg.dump_query_index) |query_index| {
if (cfg.dump_randomize) |want_randomize| {
if (case.randomize != want_randomize) {
return .{
.euclidean = 100.0 * try calculateHBCRecallMetric(alloc, split, case.top_k, case.randomize, .l2_squared, cfg.bulk_build, cfg.centroid_only_routing),
.inner_product = 100.0 * try calculateHBCRecallMetric(alloc, split, case.top_k, case.randomize, .inner_product, cfg.bulk_build, cfg.centroid_only_routing),
.cosine = 100.0 * try calculateHBCRecallMetric(alloc, split, case.top_k, case.randomize, .cosine, cfg.bulk_build, cfg.centroid_only_routing),
.euclidean = 100.0 * try calculateHBCRecallMetric(alloc, case.dataset, split, case.top_k, case.randomize, .l2_squared, cfg.bulk_build, cfg.centroid_only_routing),
.inner_product = 100.0 * try calculateHBCRecallMetric(alloc, case.dataset, split, case.top_k, case.randomize, .inner_product, cfg.bulk_build, cfg.centroid_only_routing),
.cosine = 100.0 * try calculateHBCRecallMetric(alloc, case.dataset, split, case.top_k, case.randomize, .cosine, cfg.bulk_build, cfg.centroid_only_routing),
};
}
}
Expand All @@ -292,29 +354,44 @@ fn runHBCCase(io: std.Io, alloc: std.mem.Allocator, cfg: Config, dataset_path: [
}
}

return .{
.euclidean = 100.0 * try calculateHBCRecallMetric(alloc, split, case.top_k, case.randomize, .l2_squared, cfg.bulk_build, cfg.centroid_only_routing),
.inner_product = 100.0 * try calculateHBCRecallMetric(alloc, split, case.top_k, case.randomize, .inner_product, cfg.bulk_build, cfg.centroid_only_routing),
.cosine = 100.0 * try calculateHBCRecallMetric(alloc, split, case.top_k, case.randomize, .cosine, cfg.bulk_build, cfg.centroid_only_routing),
};
const euclidean = 100.0 * try calculateHBCRecallMetric(alloc, case.dataset, split, case.top_k, case.randomize, .l2_squared, cfg.bulk_build, cfg.centroid_only_routing);
const inner_product = 100.0 * try calculateHBCRecallMetric(alloc, case.dataset, split, case.top_k, case.randomize, .inner_product, cfg.bulk_build, cfg.centroid_only_routing);
const cosine = 100.0 * try calculateHBCRecallMetric(alloc, case.dataset, split, case.top_k, case.randomize, .cosine, cfg.bulk_build, cfg.centroid_only_routing);
return .{ .euclidean = euclidean, .inner_product = inner_product, .cosine = cosine };
}

fn calculateHBCRecallMetric(
alloc: std.mem.Allocator,
dataset_label: []const u8,
split: common.SplitDataset,
top_k: usize,
randomize: bool,
metric: vec.DistanceMetric,
bulk_build: bool,
centroid_only_routing: bool,
) !f64 {
std.debug.print(
"recall_harness_metric_start suite=hbc dataset={s} randomize={any} metric={s} data={d} queries={d}\n",
.{ dataset_label, randomize, @tagName(metric), split.data.count, split.queries.count },
);
var built = try buildHBCIndex(alloc, split, randomize, metric, bulk_build, centroid_only_routing);
defer built.deinit();
std.debug.print(
"recall_harness_build_done suite=hbc dataset={s} randomize={any} metric={s}\n",
.{ dataset_label, randomize, @tagName(metric) },
);
const data = built.data();
const queries = built.queries();

var recall_sum: f64 = 0;
for (0..queries.count) |query_idx| {
const query_number = query_idx + 1;
if (shouldLogQueryProgress(query_number, queries.count)) {
std.debug.print(
"recall_harness_query_progress suite=hbc dataset={s} randomize={any} metric={s} query={d}/{d}\n",
.{ dataset_label, randomize, @tagName(metric), query_number, queries.count },
);
}
const query = queries.atConst(query_idx);
var results = try built.idx.search(query, top_k);
defer results.deinit();
Expand All @@ -330,9 +407,29 @@ fn calculateHBCRecallMetric(
defer alloc.free(truth);
recall_sum += common.calculateRecall(prediction, truth);
}
std.debug.print(
"recall_harness_metric_done suite=hbc dataset={s} randomize={any} metric={s} recall={d:.4}\n",
.{ dataset_label, randomize, @tagName(metric), recall_sum / @as(f64, @floatFromInt(queries.count)) },
);
return recall_sum / @as(f64, @floatFromInt(queries.count));
}

fn shouldLogQueryProgress(query_number: usize, total_queries: usize) bool {
return query_number == 1 or query_number == total_queries or query_number % 25 == 0;
}

fn sleepMs(ms: u64) void {
var req = std.posix.timespec{
.sec = @intCast(ms / std.time.ms_per_s),
.nsec = @intCast((ms % std.time.ms_per_s) * std.time.ns_per_ms),
};
while (true) switch (std.posix.errno(std.posix.system.nanosleep(&req, &req))) {
.SUCCESS => return,
.INTR => continue,
else => return,
};
}

const BuiltHBC = struct {
tp: TestPath,
data_owned: common.OwnedVectorSet,
Expand Down
37 changes: 31 additions & 6 deletions zig/e2e/antfly/test_resolution.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
from test_scaling import MultiNodeScalingCluster

AUTOGRAPH_E2E_TIMEOUT_S = 115.0
AUTOGRAPH_CLUSTER_STARTUP_TIMEOUT_S = 115.0
AUTOGRAPH_E2E_TEARDOWN_TIMEOUT_S = 5.0
POLL_INTERVAL_S = 0.5
POLL_REQUEST_TIMEOUT_S = 5.0
Expand Down Expand Up @@ -99,14 +100,14 @@ def resolution_cluster():
pytest.skip(f"Antfly binary not found: {binary} (set ANTFLY_BIN)")
if Path(binary).name != "antfly":
pytest.skip("distributed autograph e2e requires the antfly binary")
deadline = _Deadline(AUTOGRAPH_E2E_TIMEOUT_S)
startup_deadline = _Deadline(AUTOGRAPH_CLUSTER_STARTUP_TIMEOUT_S)
cluster = MultiNodeScalingCluster(
binary,
initial_data_node_count=3,
startup_deadline_at=deadline.expires_at,
startup_deadline_at=startup_deadline.expires_at,
)
try:
yield cluster, deadline
yield cluster
finally:
cluster.stop(timeout_s=AUTOGRAPH_E2E_TEARDOWN_TIMEOUT_S)

Expand Down Expand Up @@ -173,12 +174,32 @@ def lookup(self, table: str, key: str, *, timeout: float = 10.0) -> dict | None:
def query_table(self, table: str, payload: dict, *, timeout: float = 30.0) -> dict:
return self._check(self.s.post(f"{self.url}/tables/{table}/query", json=payload, timeout=timeout))

def diagnostic(self) -> str:
parts: list[str] = []
for label, path in (
("documents table", "/tables/documents"),
("relations graph index", "/tables/documents/indexes/relations_graph"),
("entities table", "/tables/entities"),
):
try:
response = self.s.get(f"{self.url}{path}", timeout=5)
parts.append(f"[{label}] {response.status_code} {response.text[:4000]}")
except requests.RequestException as exc:
parts.append(f"[{label}] unavailable: {exc!r}")
parts.append(f"[metadata snapshot]\n{self._server.metadata_snapshot_diagnostic()}")
parts.append(f"[logs]\n{self._server.debug_logs()}")
return "\n".join(parts)


class _Deadline:
def __init__(self, timeout_s: float):
self.timeout_s = timeout_s
self.started_at = time.monotonic()
self.expires_at = time.monotonic() + timeout_s

def elapsed(self) -> float:
return max(0.0, time.monotonic() - self.started_at)

def remaining(self) -> float:
return max(0.0, self.expires_at - time.monotonic())

Expand Down Expand Up @@ -222,7 +243,8 @@ def _wait_for_entities(api: _Api, expected_names: dict[str, str], *, deadline: _

raise AssertionError(
f"entities were not promoted within {deadline.timeout_s}s "
f"(pending={sorted(pending)!r}, last={last!r}, last_error={last_error!r})"
f"(elapsed={deadline.elapsed():.1f}s, pending={sorted(pending)!r}, "
f"last={last!r}, last_error={last_error!r})\n{api.diagnostic()}"
)


Expand Down Expand Up @@ -310,12 +332,15 @@ def _wait_for_mention_hydration(
deadline.sleep()
raise AssertionError(
f"mention graph did not hydrate promoted entities within {deadline.timeout_s}s "
f"(start_node={start_node!r}, expected={expected_names!r}, last={last!r}, last_error={last_error!r})"
f"(elapsed={deadline.elapsed():.1f}s, start_node={start_node!r}, "
f"expected={expected_names!r}, last={last!r}, last_error={last_error!r})\n"
f"{api.diagnostic()}"
)


def test_multinode_autograph_resolves_promotes_and_hydrates_entities(resolution_cluster):
cluster, deadline = resolution_cluster
cluster = resolution_cluster
deadline = _Deadline(AUTOGRAPH_E2E_TIMEOUT_S)
api = _Api(cluster.data_api_urls[0], cluster)

# Entities live in their own table (own shard group); documents are spread
Expand Down
Loading