lanes: the lane thread takes the server's compiled grammar, no second compile

zig/src/core/engine_api.zig

@@ -48,6 +48,12 @@ pub const Structure = struct {
     after: ?u32 = null,
     /// The server's compiler, which has compiled `text` once already (the engine's compile is then a cache hit).
     compiler: ?*grammar.Compiler = null,
+    /// That compile, which the engine reads instead of compiling again; the server frees it after `finished`.
+    compiled: ?grammar.Compiled = null,
+
+    pub fn free(st: Structure) void {
+        if (st.compiled) |g| g.free();
+    }
 };
 
 pub const grammar = lanes.grammar;

zig/src/core/lane_host.zig

@@ -392,6 +392,7 @@ pub const LaneHost = struct {
     /// A reply's grammar from the server's compiler (Info.structures engines get requests that carry one).
     pub fn constraint(h: *LaneHost, st: api.Structure) !*lanes.grammar.Constraint {
         const c = st.compiler orelse return error.NoGrammarCompiler;
+        if (st.compiled) |g| return c.constrained(h.gpa, g, st.after); // the server's compile, without the lock
         return c.constraint(h.gpa, h.io, @fromBackingInt(@intCast(@backingInt(st.kind))), st.text, st.after);
     }
 

zig/src/core/lane_host_grammar_test.zig

@@ -0,0 +1,70 @@
+//! A request's grammar the server compiled reaches the lane thread as is: no compile there, the same masked reply.
+const std = @import("std");
+const lanes = @import("lanes");
+const api = @import("engine_api.zig");
+const fake = lanes.grammar_fake;
+const a = std.testing.allocator;
+
+const Box = struct {
+    reason: api.Reason = .failed,
+    tokens: std.ArrayList(u32) = .empty,
+    done: std.atomic.Value(bool) = .init(false),
+    fn event(ctx: *anyopaque, _: api.Id, e: *const api.Event) void {
+        const b: *Box = @ptrCast(@alignCast(ctx));
+        switch (e.*) {
+            .tokens => |t| b.tokens.appendSlice(a, t) catch {},
+            .finished => |f| {
+                b.reason = f.reason;
+                b.done.store(true, .release);
+            },
+            else => {},
+        }
+    }
+    fn wait(b: *Box) !api.Reason {
+        const began = std.Io.Clock.awake.now(std.testing.io).toNanoseconds();
+        while (!b.done.load(.acquire)) {
+            if (std.Io.Clock.awake.now(std.testing.io).toNanoseconds() - began > 2 * std.time.ns_per_s) return error.RequestTimedOut;
+            std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {};
+        }
+        return b.reason;
+    }
+};
+
+/// One reply under the request's structure; returns its tokens (owned) after `finished`.
+fn reply(host: *api.LaneHost, id: api.Id, st: api.Structure) ![]u32 {
+    const prompt = [_]u32{ 1, 2, 3, 4, 5, 6 };
+    const request: api.Request = .{ .prompt = &prompt, .max_tokens = 8, .structure = st };
+    var box: Box = .{};
+    try host.engine().submit(id, &request, .{ .ctx = &box, .event = Box.event });
+    try std.testing.expectEqual(api.Reason.length, try box.wait());
+    return box.tokens.toOwnedSlice(a);
+}
+
+test "a compiled grammar handed in compiles nothing on the lane thread and masks the reply as a compile does" {
+    const old = fake.install();
+    defer _ = lanes.grammar.swapLibrary(old);
+    var c = fake.compiler();
+    var cfg = try lanes.Config.init(a, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7);
+    defer cfg.deinit(a);
+    var target: lanes.fake.Fake = .{ .gpa = a };
+    defer target.deinit();
+    var clock: lanes.fake.FixedClock = .{};
+    var core = lanes.Engine.init(a, &cfg, target.backend(), clock.clock());
+    defer core.deinit();
+    var host = api.LaneHost.init(a, std.testing.io, &core, .{ .lanes = 1 });
+    try host.start();
+    defer host.stop();
+    var err: [512]u8 = undefined;
+    const text = "{\"type\":\"object\"}";
+    const g = try c.compile(std.testing.io, .json_schema, text, &err); // the server's compile
+    defer g.free();
+    const handed = try reply(&host, 1, .{ .kind = .json_schema, .text = text, .compiler = &c, .compiled = g });
+    defer a.free(handed);
+    try std.testing.expectEqual(@as(usize, 1), fake.compiles);
+    try std.testing.expectEqual(@as(usize, 0), fake.frees); // the handle stays the server's
+    const compiled = try reply(&host, 2, .{ .kind = .json_schema, .text = text, .compiler = &c });
+    defer a.free(compiled);
+    try std.testing.expectEqual(@as(usize, 2), fake.compiles); // without a handle the lane thread compiles, as before
+    try std.testing.expectEqual(@as(usize, 8), handed.len);
+    try std.testing.expectEqualSlices(u32, compiled, handed);
+}

zig/src/core/lane_host_tests.zig

@@ -14,6 +14,7 @@ test {
 
 test {
     _ = @import("lane_host_score_exit_test.zig");
+    _ = @import("lane_host_grammar_test.zig");
 }
 
 test {

zig/src/core/lanes/grammar.zig

@@ -50,6 +50,17 @@ fn load(io: std.Io, a: Allocator) ?Fns {
     return f;
 }
 
+/// libtfgrammar's entry points, for a test's stand-in library (lanes/grammar_fake.zig).
+pub const Library = Fns;
+
+/// Tests only: swap in a stand-in library and get back the one it replaces.
+pub fn swapLibrary(f: ?Fns) ?Fns {
+    comptime std.debug.assert(builtin.is_test);
+    const old = lib;
+    lib = f;
+    return old;
+}
+
 /// The library's file name, for the refusal when it cannot be loaded.
 pub const library = if (builtin.os.tag == .macos) "libtfgrammar.dylib" else "libtfgrammar.so";
 
@@ -89,6 +100,11 @@ pub const Compiler = struct {
         var err: [512]u8 = undefined;
         const g = try c.compile(io, kind, text, &err);
         defer g.free();
+        return c.constrained(gpa, g, after);
+    }
+
+    /// A fresh reply's grammar from one compiled already: no compile and no lock; `g` stays the caller's.
+    pub fn constrained(c: *const Compiler, gpa: Allocator, g: Compiled, after: ?u32) !*Constraint {
         const m = try g.matcher();
         errdefer m.free();
         const out = try gpa.create(Constraint);
@@ -481,4 +497,5 @@ test "a chain grows past its last kept row with the grammar's own drafts, never
 
 test {
     _ = @import("grammar_error_test.zig");
+    _ = @import("grammar_handoff_test.zig");
 }

zig/src/core/lanes/grammar_fake.zig

@@ -0,0 +1,102 @@
+//! A stand-in libtfgrammar for tests: each grammar allows its own token pattern, and compiles and frees are counted.
+const std = @import("std");
+const grammar = @import("grammar.zig");
+
+pub var compiles: usize = 0;
+pub var frees: usize = 0;
+
+/// Grammar `g` at step `k` allows token t when (t + g + k) % 3 == 0 (so masks tell grammars and steps apart).
+const State = struct { g: usize, k: usize };
+
+fn compile(_: *anyopaque, kind: c_int, text: [*]const u8, len: usize, err: [*]u8, _: usize) callconv(.c) ?*anyopaque {
+    if (std.mem.eql(u8, text[0..len], "bad")) {
+        err[0] = 0;
+        return null;
+    }
+    compiles += 1;
+    return @ptrFromInt(1 + @as(usize, @intCast(kind)) * 1000 + len);
+}
+
+fn free(_: *anyopaque) callconv(.c) void {
+    frees += 1;
+}
+
+fn matcher(g: *anyopaque) callconv(.c) ?*anyopaque {
+    const m = std.heap.page_allocator.create(State) catch return null;
+    m.* = .{ .g = @intFromPtr(g), .k = 0 };
+    return m;
+}
+
+fn matcherFree(m: *anyopaque) callconv(.c) void {
+    std.heap.page_allocator.destroy(@as(*State, @ptrCast(@alignCast(m))));
+}
+
+fn accept(p: *anyopaque, token: i32) callconv(.c) c_int {
+    const m: *State = @ptrCast(@alignCast(p));
+    if ((@as(usize, @intCast(token)) + m.g + m.k) % 3 != 0) return 0;
+    m.k += 1;
+    return 1;
+}
+
+fn rollback(p: *anyopaque, n: c_int) callconv(.c) c_int {
+    const m: *State = @ptrCast(@alignCast(p));
+    m.k -= @intCast(n);
+    return 0;
+}
+
+fn terminated(_: *anyopaque) callconv(.c) c_int {
+    return 0;
+}
+
+fn fill(p: *anyopaque, words: [*]i32, n: c_int) callconv(.c) c_int {
+    const m: *State = @ptrCast(@alignCast(p));
+    const out: [*]u32 = @ptrCast(words);
+    for (0..@intCast(n)) |w| {
+        var bits: u32 = 0;
+        for (0..32) |b| if ((w * 32 + b + m.g + m.k) % 3 == 0) {
+            bits |= @as(u32, 1) << @intCast(b);
+        };
+        out[w] = bits;
+    }
+    return 0;
+}
+
+fn drafts(_: *anyopaque, _: [*]i32, _: c_int) callconv(.c) c_int {
+    return 0;
+}
+
+fn unused(_: *anyopaque) callconv(.c) void {}
+
+fn open(_: [*]const u8, _: usize, _: c_int, _: [*]const i32, _: c_int, _: [*]u8, _: usize) callconv(.c) ?*anyopaque {
+    return null;
+}
+
+fn wordCount(_: *anyopaque) callconv(.c) c_int {
+    return 2;
+}
+
+const library: grammar.Library = .{
+    .open = open,
+    .close = unused,
+    .words = wordCount,
+    .compile = compile,
+    .free = free,
+    .matcher = matcher,
+    .matcher_free = matcherFree,
+    .accept = accept,
+    .rollback = rollback,
+    .terminated = terminated,
+    .fill = fill,
+    .drafts = drafts,
+};
+
+/// The stand-in installed (give the returned library back to grammar.swapLibrary after) and a two-word compiler.
+pub fn install() ?grammar.Library {
+    compiles = 0;
+    frees = 0;
+    return grammar.swapLibrary(library);
+}
+
+pub fn compiler() grammar.Compiler {
+    return .{ .f = library, .h = @ptrFromInt(8), .words = 2 };
+}

zig/src/core/lanes/grammar_handoff_test.zig

@@ -0,0 +1,55 @@
+//! A grammar compiled once (the server's) gives the engine's constraint without a second compile, mask for mask.
+const std = @import("std");
+const grammar = @import("grammar.zig");
+const fake = @import("grammar_fake.zig");
+
+/// The masks of `steps` tokens, each step following the first token its mask allows.
+fn walk(a: std.mem.Allocator, c: *grammar.Constraint, steps: usize) ![]u32 {
+    var out: std.ArrayList(u32) = .empty;
+    for (0..steps) |_| {
+        const bits = (try c.next(a)).?;
+        try out.appendSlice(a, bits);
+        const t: u32 = for (0..64) |t| {
+            if (bits[t / 32] >> @intCast(t % 32) & 1 == 1) break @intCast(t);
+        } else unreachable;
+        try std.testing.expect(try c.follow(t));
+    }
+    return out.toOwnedSlice(a);
+}
+
+test "a compiled grammar's constraint compiles nothing and walks the masks a fresh compile walks" {
+    const gpa = std.testing.allocator;
+    var arena = std.heap.ArenaAllocator.init(gpa);
+    defer arena.deinit();
+    const old = fake.install();
+    defer _ = grammar.swapLibrary(old);
+    var c = fake.compiler();
+    var err: [512]u8 = undefined;
+    const g = try c.compile(std.testing.io, .json_schema, "{\"type\":\"object\"}", &err);
+    try std.testing.expectEqual(@as(usize, 1), fake.compiles);
+    const handed = try c.constrained(gpa, g, null);
+    defer handed.deinit(gpa);
+    try std.testing.expectEqual(@as(usize, 1), fake.compiles); // no second compile, so no lock either
+    const fresh = try c.constraint(gpa, std.testing.io, .json_schema, "{\"type\":\"object\"}", null);
+    defer fresh.deinit(gpa);
+    try std.testing.expectEqual(@as(usize, 2), fake.compiles);
+    try std.testing.expectEqual(@as(usize, 1), fake.frees); // the fresh compile is freed, the handed one stays the caller's
+    try std.testing.expectEqualSlices(u32, try walk(arena.allocator(), fresh, 7), try walk(arena.allocator(), handed, 7));
+    g.free();
+}
+
+test "another grammar's masks differ, so equal walks are not vacuous" {
+    const gpa = std.testing.allocator;
+    var arena = std.heap.ArenaAllocator.init(gpa);
+    defer arena.deinit();
+    const old = fake.install();
+    defer _ = grammar.swapLibrary(old);
+    var c = fake.compiler();
+    const x = try c.constraint(gpa, std.testing.io, .regex, "a+", null);
+    defer x.deinit(gpa);
+    const y = try c.constraint(gpa, std.testing.io, .regex, "ab+", null);
+    defer y.deinit(gpa);
+    const wx = try walk(arena.allocator(), x, 3);
+    const wy = try walk(arena.allocator(), y, 3);
+    try std.testing.expect(!std.mem.eql(u32, wx, wy));
+}

zig/src/core/lanes/lanes.zig

@@ -34,6 +34,7 @@ pub const drafted = @import("drafted.zig");
 pub const recorder = @import("recorder.zig");
 pub const logprob = @import("logprob.zig");
 pub const grammar = @import("grammar.zig");
+pub const grammar_fake = @import("grammar_fake.zig"); // a stand-in libtfgrammar for tests
 
 pub const Engine = engine.Engine;
 pub const Config = config.Config;

zig/src/server/chat.zig

@@ -133,7 +133,7 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa
         (priority == null and @import("server.zig").namedBackground(srv.config, f));
     const preparing = !background;
     if (preparing) _ = srv.preparing.fetchAdd(1, .acq_rel);
-    errdefer release(srv, preparing);
+    errdefer srv.release(preparing, null);
     if (gone.check()) return error.Cancelled;
     var thinking = flag(f, "enable_thinking") orelse srv.config.enable_thinking;
     if (input.prompt != null) thinking = false;
@@ -179,6 +179,7 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa
         .cache_key = rendered.cache_key,
     };
     try srv.checkFeatures(cx, f, input.tools.len > 0, thinking, rendered.ids, input.tools, &request);
+    errdefer srv.release(false, request);
     if (thinking) {
         const budget_field = f.get("thinking_budget");
         const budget: i64 = if (budget_field != null and budget_field.?.truthy()) budget_field.?.int64() orelse 0 else srv.config.thinking_budget;
@@ -209,7 +210,7 @@ pub fn generate(srv: *Server, cx: *Cx, prepared: Prepared, sink: ?Sink, gone: an
     const effort = prepared.effort;
     const stops_opt = prepared.stops;
     var preparing = prepared.preparing;
-    defer release(srv, preparing);
+    defer srv.release(preparing, request); // after the engine's finished: the errdefer below drains first
     if (request.background) {
         while (true) {
             const now = nowNs(io);
@@ -236,7 +237,7 @@ pub fn generate(srv: *Server, cx: *Cx, prepared: Prepared, sink: ?Sink, gone: an
         error.InvalidRequest => cx.fail(.request, "the request does not match the loaded model", .{}),
         error.Stale => cx.fail(.server, "{s}", .{api.stale_message}),
     };
-    release(srv, preparing); // a background request waits only while a foreground one prepares
+    srv.release(preparing, null); // a background request waits only while a foreground one prepares
     preparing = false;
     var gen: Generation = .{ .srv = srv, .a = a, .box = &box, .id = id, .reply_id = input.id, .seed = if (request.sampling) |v| v.seed else null, .sink = sink, .thinking = thinking or reply_text.isChannel(srv.markers), .stops = .{ .strings = stops_opt.strings }, .ignore_eos = stops_opt.ignore_eos, .max_tokens = request.max_tokens, .tools = input.tools, .token_ids = prepared.token_ids, .lp = if (request.logprobs != null) .{} else null };
     defer srv.noteRequest(prepared.prompt_len, box.cached orelse 0, gen.collected.items.len, box.stats.drafted, box.stats.accepted, box.stats.rounds, received, gen.first_ns, gen.last_ns, box.stats.prefill_seconds);
@@ -268,11 +269,6 @@ pub fn run(srv: *Server, cx: *Cx, input: Input, sink: ?Sink, gone: anytype) Fail
     return @import("compact.zig").run(srv, cx, input, sink, gone);
 }
 
-/// Give back a foreground request's preparing count: at its submit, or when it is never generated.
-pub fn release(srv: *Server, preparing: bool) void {
-    if (preparing) _ = srv.preparing.fetchSub(1, .acq_rel);
-}
-
 /// The token loop and the text it streams.
 const Generation = struct {
     srv: *Server,

zig/src/server/compact.zig

@@ -107,7 +107,7 @@ pub fn prepare(srv: *Server, cx: *Cx, input: chat.Input, gone: anytype) chat.Fai
     if (chat.prepare(srv, cx, input, gone)) |prepared| {
         if (!over(prepared.prompt_len, wanted(srv, input), window, at))
             return .{ .prepared = prepared, .stamp = .{ .window = window, .used = prepared.prompt_len } };
-        chat.release(srv, prepared.preparing);
+        srv.release(prepared.preparing, prepared.request);
     } else |e| if (!(e == error.Refused and cx.kind == .context_length)) return e;
     return rebuild(srv, cx, input, gone, at, input.messages.array);
 }
@@ -265,7 +265,7 @@ fn oneNote(srv: *Server, cx: *Cx, gone: anytype, task: []const u8, prev: []const
         };
         const room = srv.info.context_window - @as(i64, @intCast(fitted.prompt_len));
         if (room < 1) {
-            chat.release(srv, fitted.preparing);
+            srv.release(fitted.preparing, fitted.request);
             return .{ .text = prev, .tokens = 0 };
         }
         var out = fitted;

zig/src/server/cortheon/loop.zig

@@ -26,7 +26,7 @@ pub fn whole(t: *Turn, srv: *Server, cx: *errors.Cx, input: *chat.Input, gone: a
 pub fn stream(t: *Turn, srv: *Server, cx: *errors.Cx, input: *chat.Input, first: chat.Prepared, sink: chat.Sink, gone: anytype) chat.Failure!chat.Reply {
     var prepared = first;
     if (t.research) {
-        chat.release(srv, prepared.preparing);
+        srv.release(prepared.preparing, prepared.request);
         try research.run(t, srv, cx, input, sink, gone);
         prepared = try prepare(srv, cx, input.*, gone);
     }

zig/src/server/imprint_routes.zig

@@ -34,7 +34,7 @@ pub fn post(srv: *Server, conn: *Conn, a: Allocator, op: api.ImprintOp) void {
     const body = http_body.readJson(conn, &cx) catch return refuse(conn, a, &cx);
     const input = openai.chatInput(srv, &cx, body) catch return refuse(conn, a, &cx);
     const p = chat.prepare(srv, &cx, input, Never{}) catch return refuse(conn, a, &cx);
-    defer chat.release(srv, p.preparing);
+    defer srv.release(p.preparing, p.request);
     const ask = make(srv, op, p.request) catch return conn.sendError(a, 500, "out of memory", null) catch {};
     answer(srv, conn, a, ask);
 }

zig/src/server/openai.zig

@@ -520,7 +520,7 @@ const Run = struct {
             break :blk ready.prepared;
         };
         var handed = false; // generate gives the preparing count back from here on
-        defer if (!handed) chat.release(r.srv, prepared.preparing);
+        defer if (!handed) r.srv.release(prepared.preparing, prepared.request);
         var sink_state: StreamSink = .{ .run = r, .tools = tools, .prompt_len = prepared.prompt_len };
         handed = true;
         const sink: chat.Sink = .{ .ctx = &sink_state, .call = StreamSink.call, .begin = StreamSink.begin };