lanes: one fixed code for every wait and refusal (#2135)

zig/build/server.zig

@@ -79,6 +79,9 @@ pub fn toolHostTests(b: *std.Build, parser: *std.Build.Module, server: *std.Buil
     const profile_line = b.addRunArtifact(b.addTest(.{ .root_module = server, .filters = &.{"profile line:"} }));
     profile_line.has_side_effects = true;
     b.step("test-profile-line-host", "Profile start-up line scoping against the tree's real backend names").dependOn(&profile_line.step);
+    const metrics = b.addRunArtifact(b.addTest(.{ .root_module = server, .filters = &.{ "scrape", "load object", "wait reasons:" } }));
+    metrics.has_side_effects = true;
+    b.step("test-load-metrics-host", "Load ledger scrape and /v1/load carry bounded wait and refusal reasons").dependOn(&metrics.step);
     const admitted = b.addRunArtifact(b.addTest(.{ .root_module = server, .filters = &.{"admission"} }));
     admitted.has_side_effects = true;
     b.step("test-admission-host", "Opt-in in-flight caps and per-key rates: limits, 429s over sockets, metrics").dependOn(&admitted.step);

zig/src/core/engine_api.zig

@@ -16,6 +16,7 @@ pub const EmbeddingBatch = lanes.embedding.Batch;
 pub const embedding_vector = @import("embedding_vector.zig");
 pub const load_ledger = @import("load_ledger.zig");
 pub const kv_ledger = @import("kv_ledger.zig");
+pub const wait_reason = @import("wait_reason.zig");
 
 /// One image of a request (``Request.images``): the tower's input and the placeholder rows its output replaces.
 pub const Image = lanes.stream.Image;

zig/src/core/lane_halt_test.zig

@@ -46,7 +46,7 @@ fn haltCase(release_at_stop: ?*bool) !void {
         fn admit(ctx: *anyopaque, ask: api.Rounds.Ask, _: *[]const u8) api.Rounds.Verdict {
             const p = self(ctx);
             for (p.seqs.items) |x| if (!x.decoding) return .wait;
-            p.seqs.append(gpa, .{ .key = ask.key, .n = ask.prompt.len }) catch return .refuse;
+            p.seqs.append(gpa, .{ .key = ask.key, .n = ask.prompt.len }) catch return .{ .refuse = .backend };
             return .admit;
         }
         fn admitted(_: *anyopaque) anyerror!void {}

zig/src/core/lane_host.zig

@@ -30,6 +30,7 @@ const rounds = @import("lane_rounds.zig");
 const scoring = @import("lane_host_score.zig");
 const lane_stream = @import("lane_stream.zig");
 const lane_halt = @import("lane_halt.zig");
+const reasons = @import("wait_reason.zig");
 
 pub const LaneHost = struct {
     gpa: Allocator,
@@ -91,6 +92,7 @@ pub const LaneHost = struct {
     load: load.State = .{}, // this host's part of the server's load ledger (lane_host_load.zig)
     capacity: ?load.Source = null, // how the engine holds KV, set before start; null: no capacity group
     kv: ?@import("kv_ledger.zig").Ledger = null, // Info.ledger's budget or the paged pool's (lane_host_ledger.zig)
+    refusals: reasons.Refusals = .{}, // refused requests by fixed reason, under mutex
 
     pub const ThreadInit = struct { ctx: *anyopaque, run: *const fn (ctx: *anyopaque) anyerror!void };
 
@@ -191,26 +193,7 @@ pub const LaneHost = struct {
     }
 
     fn submitFn(ctx: *anyopaque, id: Id, request: *const Request, sink: Sink) api.SubmitError!void {
-        const h = self(ctx);
-        if (@import("lane_prompt_only.zig").refuses(h, request)) return error.InvalidRequest;
-        const job = h.gpa.create(Job) catch return error.Busy;
-        job.* = .{ .host = h, .id = id, .request = request, .sink = sink, .submitted = std.Io.Clock.awake.now(h.io).toNanoseconds() };
-        h.mutex.lockUncancelable(h.io);
-        defer h.unlock();
-        const unavailable = h.closing or (if (h.health) |source| !source.read(source.ctx) else false);
-        if (unavailable or h.stale or pressure.refuses(h)) {
-            h.gpa.destroy(job);
-            return if (unavailable) error.Closed else if (h.stale) error.Stale else error.MemoryPressure;
-        }
-        // foreground before background, each in arrival order (the Python job queue's priority)
-        var at = h.queued.items.len;
-        while (!request.background and at > 0 and h.queued.items[at - 1].request.background) at -= 1;
-        h.queued.insert(h.gpa, at, job) catch {
-            h.gpa.destroy(job);
-            return error.Busy;
-        };
-        h.load.started += 1;
-        h.wake.signal(h.io);
+        return @import("lane_host_submit.zig").submit(self(ctx), id, request, sink);
     }
 
     fn cancelFn(ctx: *anyopaque, id: Id) void {
@@ -345,6 +328,7 @@ pub const LaneHost = struct {
             return .none;
         };
         if (ledger.never(h, h.queued.items[at])) { // past the whole budget alone: no release will ever make it fit
+            h.refusals.note(.prompt_too_large);
             const job = h.queued.orderedRemove(at);
             h.unlock();
             h.finish(job, .failed, ledger.too_big);
@@ -359,6 +343,7 @@ pub const LaneHost = struct {
         job.seq = h.admissions;
         h.admissions += 1;
         h.admitted.append(h.gpa, job) catch {
+            h.refusals.note(.engine_busy);
             h.unlock();
             h.finish(job, .failed, "out of memory");
             return .failed;
@@ -407,7 +392,10 @@ pub const LaneHost = struct {
         return switch (e) {
             error.Cancelled => h.cancel(job),
             error.LaneMemory => ends.requeue(h, job),
-            else => ends.dropFailed(h, job, e),
+            else => blk: {
+                h.refused(.backend_prompt);
+                break :blk ends.dropFailed(h, job, e);
+            },
         };
     }
 
@@ -492,6 +480,13 @@ pub const LaneHost = struct {
         return x.text(x.ctx, e) orelse @errorName(e);
     }
 
+    /// One refused request, by its fixed reason; the caller does not hold the lock.
+    pub fn refused(h: *LaneHost, why: reasons.Refusal) void {
+        h.lock();
+        defer h.unlock();
+        h.refusals.note(why);
+    }
+
     pub fn drop(h: *LaneHost, job: *Job, message: []const u8) bool {
         h.remove(job);
         if (job.started and !job.stream.finished) h.core.discard(&job.stream);

zig/src/core/lane_host_end.zig

@@ -65,6 +65,7 @@ pub fn deliverPending(h: *LaneHost) void {
 /// A round the backend failed: the failure hook hears the error and every admitted stream gets its name.
 pub fn failRound(h: *LaneHost, err: anyerror) void {
     failed(h, err);
+    h.refused(.backend_round);
     failAll(h, @errorName(err));
 }
 
@@ -123,7 +124,7 @@ pub fn requeue(h: *LaneHost, job: *Job) bool {
     if (waits) {
         h.load.requeued += 1;
         h.memory_lanes = others; // under the lock: a publish reads it
-    }
+    } else h.refusals.note(.lane_memory_alone);
     h.unlock();
     if (!waits) {
         h.finish(job, .failed, if (others > 0) "out of memory" else alone_words);

zig/src/core/lane_host_fetch.zig

@@ -43,7 +43,7 @@ pub fn queue(h: *LaneHost) void {
 }
 
 /// Whether queued `job` waits for its read: still going and inside the bound (the lock held).
-pub fn held(h: *LaneHost, job: *const Job) bool {
+pub fn held(h: *const LaneHost, job: *const Job) bool {
     const r = (h.cache orelse return false).disk orelse return false;
     const f = r.fetch orelse return false;
     if (job.fetch.writing) |name| {

zig/src/core/lane_host_load.zig

@@ -195,7 +195,11 @@ pub fn yielded(h: *LaneHost) void {
 /// The `requests` group from the host's lists; the caller holds the host mutex.
 fn requests(h: *const LaneHost) ledger.Requests {
     var waiting_tokens: u64 = 0;
-    for (h.queued.items) |job| waiting_tokens += job.request.prompt.len;
+    var waits: @import("wait_reason.zig").Waits = .{};
+    for (h.queued.items) |job| {
+        waiting_tokens += job.request.prompt.len;
+        waits.note(waitReason(h, job));
+    }
     const left = filling(h);
     return .{
         .running = @intCast(h.admitted.items.len),
@@ -206,9 +210,24 @@ fn requests(h: *const LaneHost) ledger.Requests {
         .started = h.load.started,
         .requeued = h.load.requeued,
         .yields = h.load.yields,
+        .waiting_reasons = waits.counts,
+        .refusals = h.refusals,
     };
 }
 
+/// Why a queued request is not running: one fixed code, in the order admission decides (lock held).
+fn waitReason(h: *const LaneHost, job: *const Job) @import("wait_reason.zig").Wait {
+    if (h.stale) return .engine_unavailable;
+    if (job.prefix.leader != null) return .prefix_cut;
+    if (@import("lane_host_fetch.zig").held(h, job)) return .state_read;
+    if (@import("lane_host_pressure.zig").holds(h)) return .memory_pressure;
+    const lanes_now = @min(h.info_.lanes, h.memory_lanes orelse h.info_.lanes);
+    if (h.admitted.items.len >= lanes_now) return .engine_busy;
+    if (@import("lane_host_ledger.zig").priced(h)) return if (@import("lane_host_ledger.zig").fits(h, job)) .engine_busy else .ledger_bytes;
+    if (h.info_.admit) |a| return if (a.fits(a.ctx, @intCast(job.request.prompt.len), job.request.max_tokens)) .engine_busy else .pool_room;
+    return .engine_busy;
+}
+
 /// Prompt tokens left in the foreground passes filling now: 0 for none, null when the backend reports no progress.
 fn filling(h: *const LaneHost) ?u64 {
     var left: u64 = 0;

zig/src/core/lane_host_load_test.zig

@@ -9,16 +9,35 @@ const Id = api.Id;
 const Event = api.Event;
 const Request = api.Request;
 const Engine = api.Engine;
+const reasons = api.wait_reason;
+
+/// A pool where prompts of 64 tokens or more fit only when nothing else runs.
+const Pool = struct {
+    fn fits(_: *anyopaque, prompt: u32, _: u32) bool {
+        return prompt < 64;
+    }
+};
+
+fn cost(_: *anyopaque, prompt: u32, max_tokens: u32) u64 {
+    return @as(u64, prompt + max_tokens) * 4;
+}
 
 /// How a request finished, as the host reports it.
 const Box = struct {
     mutex: std.Io.Mutex = .init,
     done: ?Reason = null,
+    tokens: usize = 0,
     fn event(ctx: *anyopaque, _: Id, e: *const Event) void {
         const b: *@This() = @ptrCast(@alignCast(ctx));
         b.mutex.lockUncancelable(std.testing.io);
         defer b.mutex.unlock(std.testing.io);
         if (e.* == .finished) b.done = e.finished.reason;
+        if (e.* == .tokens) b.tokens += e.tokens.len;
+    }
+    fn started(b: *@This()) bool {
+        b.mutex.lockUncancelable(std.testing.io);
+        defer b.mutex.unlock(std.testing.io);
+        return b.tokens > 0;
     }
     fn wait(b: *@This()) Reason {
         while (true) {
@@ -410,3 +429,133 @@ test "a request holding a lane lowers the free lanes" {
     try std.testing.expectEqual(@as(?u32, 1), watch.free);
     try std.testing.expectEqual(@as(u32, 2), ledger.read().capacity.?.lanes_free);
 }
+
+test "waits by reason: a pool-room head and a lane-full queue each carry one fixed code" {
+    var rig: Rig = undefined;
+    try rig.init(.{ .lanes = 3, .admit = .{ .ctx = @ptrCast(&rig), .fits = Pool.fits } });
+    defer rig.deinit();
+    var ledger: Ledger = .init(std.testing.io);
+    ledger.attach(rig.host.engine());
+    const short = [_]u32{ 3, 1, 4 };
+    const long: [64]u32 = @splat(7);
+    var a: Box = .{};
+    var c: Box = .{};
+    var d: Box = .{};
+    var e: Box = .{};
+    var f: Box = .{};
+    const rs: Request = .{ .prompt = &short, .max_tokens = 512 };
+    const rl: Request = .{ .prompt = &long, .max_tokens = 4 };
+    try rig.host.engine().submit(1, &rs, a.sink());
+    try rig.host.engine().submit(2, &rl, c.sink());
+    try rig.host.engine().submit(3, &rs, d.sink());
+    try rig.host.start();
+    // one short prompt runs, a second goes beside it, and the long head waits with a lane still free
+    while (!d.started()) std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {};
+    var q = ledger.read().requests.?;
+    try std.testing.expectEqual(@as(u32, 1), q.waiting);
+    try std.testing.expectEqual(@as(u32, 1), q.waiting_reasons[@backingInt(reasons.Wait.pool_room)]);
+    try std.testing.expectEqual(@as(u32, 0), q.waiting_reasons[@backingInt(reasons.Wait.engine_busy)]);
+    // two more short prompts: one takes the last lane, the rest wait for a lane, the long head included
+    try rig.host.engine().submit(4, &rs, e.sink());
+    try rig.host.engine().submit(5, &rs, f.sink());
+    for (0..2000) |_| {
+        var status: api.Status = .{};
+        rig.host.engine().status(&status, &.{});
+        if (status.running == 3) break;
+        std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {};
+    }
+    q = ledger.read().requests.?;
+    try std.testing.expectEqual(@as(u32, 2), q.waiting);
+    try std.testing.expectEqual(@as(u32, 0), q.waiting_reasons[@backingInt(reasons.Wait.pool_room)]);
+    try std.testing.expectEqual(q.waiting, q.waiting_reasons[@backingInt(reasons.Wait.engine_busy)]);
+    try std.testing.expectEqual(q.waiting, (reasons.Waits{ .counts = q.waiting_reasons }).total());
+    _ = a.wait();
+    _ = c.wait();
+    _ = d.wait();
+    _ = e.wait();
+    _ = f.wait();
+    ledger.detach(false);
+}
+
+test "repeated prompt failures release every lane and count one fixed refusal each" {
+    var rig: Rig = undefined;
+    try rig.init(.{ .lanes = 2, .ledger = .{ .ctx = @ptrCast(&rig), .total = 1 << 20, .cost = cost } });
+    defer rig.deinit();
+    rig.target.memory_lanes = 0; // every prompt pass finds no lane memory, even alone
+    var ledger: Ledger = .init(std.testing.io);
+    ledger.attach(rig.host.engine());
+    const prompt = [_]u32{ 2, 7, 1, 8 };
+    var boxes: [3]Box = .{ .{}, .{}, .{} };
+    for (&boxes, 0..) |*b, i| {
+        const request: Request = .{ .prompt = &prompt, .max_tokens = 8 };
+        try rig.host.engine().submit(@intCast(i + 1), &request, b.sink());
+    }
+    try rig.host.start();
+    for (&boxes) |*b| try std.testing.expectEqual(Reason.failed, b.wait());
+    rig.host.lock(); // the last finish's publish has run by the time its event lands
+    rig.host.unlock();
+    const q = ledger.read().requests.?;
+    try std.testing.expectEqual(@as(u32, 0), q.waiting);
+    try std.testing.expectEqual(@as(u32, 0), q.running);
+    try std.testing.expectEqual(@as(u64, 3), q.refusals.count(.lane_memory_alone));
+    try std.testing.expect(rig.host.kv.?.idle()); // every promised lease came back
+    // memory returns: one request runs to its end and leaves the lane free
+    rig.target.memory_lanes = null;
+    var ok: Box = .{};
+    const request: Request = .{ .prompt = &prompt, .max_tokens = 8 };
+    try rig.host.engine().submit(9, &request, ok.sink());
+    _ = ok.wait();
+    rig.host.lock();
+    rig.host.unlock();
+    var status: api.Status = .{};
+    rig.host.engine().status(&status, &.{});
+    try std.testing.expectEqual(@as(u32, 0), status.running);
+    try std.testing.expectEqual(@as(u32, 0), status.waiting);
+    try std.testing.expect(rig.host.kv.?.idle());
+    ledger.detach(false);
+}
+
+test "page-backed waits: publishing a queued request never asks for an absent per-request price" {
+    var rig: Rig = undefined;
+    try rig.init(.{ .lanes = 1 });
+    defer rig.deinit();
+    const kv = @import("kv_ledger.zig");
+    const Pages = struct {
+        fn read(_: *const anyopaque) kv.Pages {
+            return .{ .total = 4, .available = 4, .reclaimable = 0, .unit = 64 };
+        }
+        fn detach(_: *anyopaque) void {}
+    };
+    rig.host.kv = .ofPages(.{ .ctx = &rig, .read = Pages.read, .detach = Pages.detach });
+    var ledger: Ledger = .init(std.testing.io);
+    ledger.attach(rig.host.engine());
+    const prompt = [_]u32{ 2, 7, 1 };
+    const request: Request = .{ .prompt = &prompt, .max_tokens = 4 };
+    var box: Box = .{};
+    try rig.host.engine().submit(1, &request, box.sink());
+    try rig.host.start();
+    _ = box.wait();
+    ledger.detach(false);
+}
+
+test "capacity refusal codes: closed and stale keep their errors and share engine_unavailable" {
+    var rig: Rig = undefined;
+    try rig.init(.{ .lanes = 1 });
+    defer rig.deinit();
+    var ledger: Ledger = .init(std.testing.io);
+    ledger.attach(rig.host.engine());
+    const prompt = [_]u32{ 2, 7, 1 };
+    const request: Request = .{ .prompt = &prompt, .max_tokens = 4 };
+    var box: Box = .{};
+    rig.host.stale = true;
+    try std.testing.expectError(error.Stale, rig.host.engine().submit(1, &request, box.sink()));
+    rig.host.stale = false;
+    rig.host.closing = true;
+    try std.testing.expectError(error.Closed, rig.host.engine().submit(2, &request, box.sink()));
+    const q = ledger.read().requests.?;
+    try std.testing.expectEqual(@as(u64, 2), q.refusals.count(.engine_unavailable));
+    try std.testing.expectEqual(@as(u64, 0), q.refusals.count(.engine_busy));
+    try std.testing.expectEqual(@as(u32, 0), q.waiting);
+    rig.host.closing = false;
+    ledger.detach(false);
+}

zig/src/core/lane_host_submit.zig

@@ -0,0 +1,36 @@
+//! Request submission to the lane host: the same queue order and refusal conditions, with bounded reason counts.
+const std = @import("std");
+const api = @import("engine_api.zig");
+const LaneHost = @import("lane_host.zig").LaneHost;
+const Job = LaneHost.Job;
+const pressure = @import("lane_host_pressure.zig");
+
+pub fn submit(h: *LaneHost, id: api.Id, request: *const api.Request, sink: api.Sink) api.SubmitError!void {
+    if (@import("lane_prompt_only.zig").refuses(h, request)) {
+        h.refused(.invalid_request);
+        return error.InvalidRequest;
+    }
+    const job = h.gpa.create(Job) catch {
+        h.refused(.engine_busy);
+        return error.Busy;
+    };
+    job.* = .{ .host = h, .id = id, .request = request, .sink = sink, .submitted = std.Io.Clock.awake.now(h.io).toNanoseconds() };
+    h.mutex.lockUncancelable(h.io);
+    defer h.unlock();
+    const unavailable = h.closing or (if (h.health) |source| !source.read(source.ctx) else false);
+    if (unavailable or h.stale or pressure.refuses(h)) {
+        h.refusals.note(if (unavailable or h.stale) .engine_unavailable else .memory_pressure);
+        h.gpa.destroy(job);
+        return if (unavailable) error.Closed else if (h.stale) error.Stale else error.MemoryPressure;
+    }
+    // foreground before background, each in arrival order (the Python job queue's priority)
+    var at = h.queued.items.len;
+    while (!request.background and at > 0 and h.queued.items[at - 1].request.background) at -= 1;
+    h.queued.insert(h.gpa, at, job) catch {
+        h.refusals.note(.engine_busy);
+        h.gpa.destroy(job);
+        return error.Busy;
+    };
+    h.load.started += 1;
+    h.wake.signal(h.io);
+}

zig/src/core/lane_host_tests.zig

@@ -10,6 +10,7 @@ test {
     _ = @import("lane_init_test.zig");
     _ = @import("lane_failure_review_test.zig");
     _ = @import("lane_host_first_test.zig");
+    _ = @import("wait_reason.zig");
 }
 
 test {

zig/src/core/lane_job.zig

@@ -68,6 +68,7 @@ pub const Job = struct {
         const store = job.host.cache orelse return;
         if (!job.started) return store.release(e);
         job.stream.reuse.saved = null; // the backend restores before its first chunk; a later keep may free it
+        if (job.stream.reuse_failed) job.host.refused(.restore_refused);
         const ok = !job.stream.reuse_failed and job.stream.cached == e.at;
         if (job.stream.reuse_failed or ok) store.resumedWithReason(e, job.cacheKey(), if (ok) null else job.stream.reuse_error orelse error.RestoreRejected) else store.release(e);
     }