From a3303c75569177788544bed303409969a0bf176c Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:28:21 +0700 Subject: [PATCH 01/43] CUDA keyed graph cache and device arena, both off until a family installs them. graph_cache.Cache captures one executable graph per caller-named key and replays it when the key comes back; ranks agree on every capture, eviction and fallback through an Agree hook, and a memory floor turns misses eager. arena.Arena sub-allocates DeviceBuffer.alloc from a few large chunks (VMM or plain cuMemAlloc) through memory.setHook; borrowed spans and pointers it does not own are freed as before. --- zig/src/cuda/abi.zig | 2 + zig/src/cuda/arena.zig | 331 +++++++++++++++++++++++++++++ zig/src/cuda/graph_cache.zig | 392 +++++++++++++++++++++++++++++++++++ zig/src/cuda/memory.zig | 52 ++++- zig/src/cuda/root.zig | 5 + 5 files changed, 779 insertions(+), 3 deletions(-) create mode 100644 zig/src/cuda/arena.zig create mode 100644 zig/src/cuda/graph_cache.zig diff --git a/zig/src/cuda/abi.zig b/zig/src/cuda/abi.zig index 485edaabde..a099686c67 100644 --- a/zig/src/cuda/abi.zig +++ b/zig/src/cuda/abi.zig @@ -39,6 +39,8 @@ pub const DeviceAttribute = enum(c_int) { compute_capability_minor = 76, max_shared_memory_per_multiprocessor = 81, max_shared_memory_per_block_optin = 97, + virtual_memory_management_supported = 102, + gpu_direct_rdma_with_cuda_vmm_supported = 126, }; pub const FunctionAttribute = enum(c_int) { diff --git a/zig/src/cuda/arena.zig b/zig/src/cuda/arena.zig new file mode 100644 index 0000000000..2c80b7581b --- /dev/null +++ b/zig/src/cuda/arena.zig @@ -0,0 +1,331 @@ +//! A device sub-allocator under DeviceBuffer.alloc: buffers carved from a few large chunks, not one cuMemAlloc each. + +const std = @import("std"); +const abi = @import("abi.zig"); +const Driver = @import("driver.zig").Driver; +const memory = @import("memory.zig"); + +pub const Error = error{ DriverUnavailable, MissingSymbol, Unsupported, CudaFailed, OutOfMemory }; + +/// Where chunks come from: `chunk` maps `len` bytes (a multiple of `granularity`), `release` gives a chunk back. +pub const Backend = struct { + ctx: *anyopaque, + granularity: usize, + chunk: *const fn (ctx: *anyopaque, len: usize) ?u64, + release: *const fn (ctx: *anyopaque, base: u64, len: usize) void, + name: []const u8, +}; + +pub const Options = struct { + small_max: usize = 4 << 20, // requests up to this many bytes share chunks + small_chunk: usize = 64 << 20, // a shared chunk's bytes, rounded up to the granularity + alignment: usize = 512, // every range starts at a multiple of this +}; + +pub const Stats = struct { small: u64 = 0, large: u64 = 0, reused: u64 = 0, fallback: u64 = 0, chunk_bytes: u64 = 0 }; + +/// Small requests bump 512-byte aligned through shared chunks and are reused by size; large ones get their own chunk. +pub const Arena = struct { + gpa: std.mem.Allocator, + be: Backend, + o: Options, + mutex: std.atomic.Mutex = .unlocked, + small: std.ArrayList([2]u64) = .empty, // shared chunks (base, len) + cursor: u64 = 0, // the current shared chunk's next free byte + limit: u64 = 0, + free_small: std.AutoHashMapUnmanaged(u64, std.ArrayList(u64)) = .empty, // freed small ranges by rounded size + live: std.AutoHashMapUnmanaged(u64, u64) = .empty, // live small ranges: address to rounded size + large: std.AutoHashMapUnmanaged(u64, u64) = .empty, // large chunks: address to mapped size + stats: Stats = .{}, + + pub fn init(gpa: std.mem.Allocator, be: Backend, o: Options) Arena { + return .{ .gpa = gpa, .be = be, .o = o }; + } + + pub fn deinit(a: *Arena) void { + var it = a.large.iterator(); + while (it.next()) |e| a.be.release(a.be.ctx, e.key_ptr.*, e.value_ptr.*); + for (a.small.items) |c| a.be.release(a.be.ctx, c[0], c[1]); + var f = a.free_small.valueIterator(); + while (f.next()) |l| l.deinit(a.gpa); + a.free_small.deinit(a.gpa); + a.live.deinit(a.gpa); + a.large.deinit(a.gpa); + a.small.deinit(a.gpa); + a.* = undefined; + } + + fn lock(a: *Arena) void { + while (!a.mutex.tryLock()) std.Thread.yield() catch {}; + } + + fn roundUp(x: usize, m: usize) usize { + return std.mem.alignForward(usize, x, m); + } + + /// An address for `len` bytes, or null (the caller then allocates as before). + pub fn alloc(a: *Arena, len: usize) ?u64 { + if (len == 0) return null; + a.lock(); + defer a.mutex.unlock(); + return a.allocLocked(len) catch null; + } + + fn allocLocked(a: *Arena, len: usize) !?u64 { + if (len > a.o.small_max) { + const n = roundUp(len, a.be.granularity); + const p = a.be.chunk(a.be.ctx, n) orelse { + a.stats.fallback += 1; + return null; + }; + errdefer a.be.release(a.be.ctx, p, n); + try a.large.put(a.gpa, p, n); + a.stats.large += 1; + a.stats.chunk_bytes += n; + return p; + } + const n: u64 = roundUp(len, a.o.alignment); + if (a.free_small.getPtr(n)) |l| if (l.pop()) |p| { + try a.live.put(a.gpa, p, n); + a.stats.reused += 1; + return p; + }; + if (a.limit - a.cursor < n) { + const cl = roundUp(@max(a.o.small_chunk, n), a.be.granularity); + const base = a.be.chunk(a.be.ctx, cl) orelse { + a.stats.fallback += 1; + return null; + }; + errdefer a.be.release(a.be.ctx, base, cl); + try a.small.append(a.gpa, .{ base, cl }); + a.cursor = base; + a.limit = base + cl; + a.stats.chunk_bytes += cl; + } + const p = a.cursor; + try a.live.put(a.gpa, p, n); + a.cursor += n; + a.stats.small += 1; + return p; + } + + /// True when `p` was the arena's (and is now free); false: not ours, the caller frees it as before. + pub fn free(a: *Arena, p: u64) bool { + a.lock(); + defer a.mutex.unlock(); + if (a.large.fetchRemove(p)) |e| { + a.be.release(a.be.ctx, p, e.value); + return true; + } + const e = a.live.fetchRemove(p) orelse return false; + const g = a.free_small.getOrPut(a.gpa, e.value) catch return true; // leaked range, still ours + if (!g.found_existing) g.value_ptr.* = .empty; + g.value_ptr.append(a.gpa, p) catch {}; + return true; + } + + // -- memory.zig's hook ---------------------------------------------------------------------------------------- + + fn hookAlloc(ctx: *anyopaque, len: usize) ?u64 { + const a: *Arena = @ptrCast(@alignCast(ctx)); + return a.alloc(len); + } + + fn hookFree(ctx: *anyopaque, p: u64) bool { + const a: *Arena = @ptrCast(@alignCast(ctx)); + return a.free(p); + } + + /// Every later DeviceBuffer.alloc comes from the arena, which must outlive its buffers; only addresses change. + pub fn install(a: *Arena) void { + memory.setHook(.{ .ctx = a, .alloc = hookAlloc, .free = hookFree }); + } +}; + +// -- backends --------------------------------------------------------------------------------------------------------- + +/// One cuMemAlloc a chunk (the driver's own 2 MiB-aligned allocations). +pub const Plain = struct { + d: *const Driver, + + pub fn backend(p: *Plain) Backend { + return .{ .ctx = p, .granularity = 2 << 20, .chunk = chunk, .release = release, .name = "plain" }; + } + + fn chunk(ctx: *anyopaque, len: usize) ?u64 { + const p: *Plain = @ptrCast(@alignCast(ctx)); + var q: abi.DevicePtr = 0; + if (p.d.api.cuMemAlloc_v2(&q, len) != abi.success) return null; + return q; + } + + fn release(ctx: *anyopaque, base: u64, _: usize) void { + const p: *Plain = @ptrCast(@alignCast(ctx)); + _ = p.d.api.cuMemFree_v2(base); + } +}; + +/// CUDA's virtual memory management: one reserved virtual range, chunks of the recommended granularity mapped in order. +pub const Vmm = struct { + lib: std.DynLib, + device: abi.Device, + prop: Prop, + granularity: usize, + va: u64 = 0, + va_len: usize = 0, + next: u64 = 0, + api: Api, + + const Api = struct { + reserve: *const fn (*u64, usize, usize, u64, c_ulonglong) callconv(.c) abi.Result, + addressFree: *const fn (u64, usize) callconv(.c) abi.Result, + create: *const fn (*u64, usize, *const Prop, c_ulonglong) callconv(.c) abi.Result, + release: *const fn (u64) callconv(.c) abi.Result, + map: *const fn (u64, usize, usize, u64, c_ulonglong) callconv(.c) abi.Result, + unmap: *const fn (u64, usize) callconv(.c) abi.Result, + setAccess: *const fn (u64, usize, *const AccessDesc, usize) callconv(.c) abi.Result, + granularity: *const fn (*usize, *const Prop, c_int) callconv(.c) abi.Result, + }; + + pub const Prop = extern struct { // CUmemAllocationProp + type: c_int = 1, // CU_MEM_ALLOCATION_TYPE_PINNED + handle_types: c_int = 0, + loc_type: c_int = 1, // CU_MEM_LOCATION_TYPE_DEVICE + loc_id: c_int = 0, + win32: ?*anyopaque = null, + compression: u8 = 0, + rdma: u8 = 0, + usage: u16 = 0, + reserved: [4]u8 = .{ 0, 0, 0, 0 }, + }; + const AccessDesc = extern struct { loc_type: c_int = 1, loc_id: c_int, flags: c_int = 3 }; // CUmemAccessDesc, read-write + + comptime { + std.debug.assert(@sizeOf(Prop) == 32 and @offsetOf(Prop, "compression") == 24 and @sizeOf(AccessDesc) == 12); + } + + /// `va_bytes` of virtual addresses reserved (cheap: no memory behind them until a chunk maps). + pub fn open(d: *const Driver, device: abi.Device, va_bytes: usize) Error!Vmm { + var vmm_ok: c_int = 0; + if (d.api.cuDeviceGetAttribute(&vmm_ok, .virtual_memory_management_supported, device) != abi.success or vmm_ok == 0) return error.Unsupported; // VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED + var lib = std.DynLib.open("libcuda.so.1") catch return error.DriverUnavailable; + errdefer lib.close(); + var api: Api = undefined; + const names = .{ .{ "reserve", "cuMemAddressReserve" }, .{ "addressFree", "cuMemAddressFree" }, .{ "create", "cuMemCreate" }, .{ "release", "cuMemRelease" }, .{ "map", "cuMemMap" }, .{ "unmap", "cuMemUnmap" }, .{ "setAccess", "cuMemSetAccess" }, .{ "granularity", "cuMemGetAllocationGranularity" } }; + inline for (names) |n| @field(api, n[0]) = lib.lookup(@FieldType(Api, n[0]), n[1]) orelse return error.MissingSymbol; + var prop: Prop = .{ .loc_id = device }; + var rdma: c_int = 0; // GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED: keep the memory registrable as cuMemAlloc's is + if (d.api.cuDeviceGetAttribute(&rdma, .gpu_direct_rdma_with_cuda_vmm_supported, device) == abi.success and rdma != 0) prop.rdma = 1; + var g: usize = 0; + if (api.granularity(&g, &prop, 1) != abi.success or g == 0) return error.Unsupported; // RECOMMENDED + g = @max(g, 2 << 20); + const len = std.mem.alignForward(usize, va_bytes, g); + var va: u64 = 0; + if (api.reserve(&va, len, g, 0, 0) != abi.success) return error.CudaFailed; + return .{ .lib = lib, .device = device, .prop = prop, .granularity = g, .va = va, .va_len = len, .next = va, .api = api }; + } + + pub fn close(v: *Vmm) void { + _ = v.api.addressFree(v.va, v.va_len); + v.lib.close(); + } + + pub fn backend(v: *Vmm) Backend { + return .{ .ctx = v, .granularity = v.granularity, .chunk = chunk, .release = release, .name = "vmm" }; + } + + fn chunk(ctx: *anyopaque, len: usize) ?u64 { + const v: *Vmm = @ptrCast(@alignCast(ctx)); + if (v.va + v.va_len - v.next < len) return null; + var h: u64 = 0; + if (v.api.create(&h, len, &v.prop, 0) != abi.success) return null; + defer _ = v.api.release(h); // the mapping holds the memory until it is unmapped + const at = v.next; + if (v.api.map(at, len, 0, h, 0) != abi.success) return null; + const desc: AccessDesc = .{ .loc_id = v.device }; + if (v.api.setAccess(at, len, &desc, 1) != abi.success) { + _ = v.api.unmap(at, len); + return null; + } + v.next += len; + return at; + } + + fn release(ctx: *anyopaque, base: u64, len: usize) void { + const v: *Vmm = @ptrCast(@alignCast(ctx)); + _ = v.api.unmap(base, len); // the virtual range is not reused: the reservation is sized for the process + } +}; + +// -- tests (a fake backend: addresses only) --------------------------------------------------------------------------- + +const Fake = struct { + next: u64 = 1 << 40, + chunks: u32 = 0, + released: u32 = 0, + fail: bool = false, + + fn backend(f: *Fake, g: usize) Backend { + return .{ .ctx = f, .granularity = g, .chunk = chunk, .release = release, .name = "fake" }; + } + fn chunk(ctx: *anyopaque, len: usize) ?u64 { + const f: *Fake = @ptrCast(@alignCast(ctx)); + if (f.fail) return null; + const p = f.next; + f.next += len; + f.chunks += 1; + return p; + } + fn release(ctx: *anyopaque, _: u64, _: usize) void { + const f: *Fake = @ptrCast(@alignCast(ctx)); + f.released += 1; + } +}; + +test "arena: small requests pack 512-aligned into shared chunks, are reused by size, large ones get their own" { + const t = std.testing; + var f: Fake = .{}; + var a = Arena.init(t.allocator, f.backend(2 << 20), .{ .small_max = 1 << 20, .small_chunk = 4 << 20 }); + defer a.deinit(); + const p1 = a.alloc(20480).?; // a norm weight + const p2 = a.alloc(983040).?; // an mHC fn16 + const p3 = a.alloc(100).?; + try t.expectEqual(p1 + 20480, p2); + try t.expectEqual(p2 + 983040, p3); + for ([_]u64{ p1, p2, p3 }) |p| try t.expectEqual(@as(u64, 0), p % 512); + try t.expectEqual(@as(u32, 1), f.chunks); + // a freed range comes back for the same rounded size + try t.expect(a.free(p3)); + try t.expectEqual(p3, a.alloc(17).?); + try t.expectEqual(@as(u64, 1), a.stats.reused); + // three more 1 MiB ranges fit the 4 MiB chunk, the fourth opens a new shared chunk + for (0..3) |_| _ = a.alloc(1 << 20).?; + try t.expectEqual(@as(u32, 1), f.chunks); + _ = a.alloc(1 << 20).?; + try t.expectEqual(@as(u32, 2), f.chunks); + // large: its own chunk, rounded to the granularity, released when freed + const big = a.alloc(5 << 20).?; + try t.expectEqual(@as(u32, 3), f.chunks); + try t.expectEqual(@as(u64, 6 << 20), a.large.get(big).?); + try t.expect(a.free(big)); + try t.expectEqual(@as(u32, 1), f.released); + // not ours + try t.expect(!a.free(12345)); + // a backend that cannot map: null, the caller allocates as before + f.fail = true; + try t.expectEqual(@as(?u64, null), a.alloc(64 << 20)); + try t.expectEqual(@as(u64, 1), a.stats.fallback); +} + +test "arena: under DeviceBuffer.alloc through memory.setHook" { + const t = std.testing; + var f: Fake = .{}; + var a = Arena.init(t.allocator, f.backend(2 << 20), .{}); + defer a.deinit(); + a.install(); + defer memory.setHook(null); + const h = memory.currentHook().?; + const p = h.alloc(h.ctx, 4096).?; + try t.expect(h.free(h.ctx, p)); + try t.expect(!h.free(h.ctx, 7)); +} diff --git a/zig/src/cuda/graph_cache.zig b/zig/src/cuda/graph_cache.zig new file mode 100644 index 0000000000..a2fb8e18bb --- /dev/null +++ b/zig/src/cuda/graph_cache.zig @@ -0,0 +1,392 @@ +//! Keyed CUDA graphs for any family: one captured executable per key, launched again when the key comes back. + +const std = @import("std"); +const graph = @import("graph.zig"); +const Driver = @import("driver.zig").Driver; +const Stream = @import("stream.zig").Stream; +const quiet = @import("builtin").is_test; // a test's stderr is the build runner's channel + +/// One step's launches on `stream`; it must not synchronize, allocate device memory or read the device on the host. +pub const Body = struct { + ctx: *anyopaque, + run: *const fn (ctx: *anyopaque, stream: Stream) anyerror!void, +}; + +/// Every rank's verdict AND-ed, so ranks capture, evict and fall back at the same steps; never called in a capture. +pub const Agree = struct { + ctx: *anyopaque, + all: *const fn (ctx: *anyopaque, ok: bool) anyerror!bool, +}; + +/// Capture a body into an executable graph, launch it, free it: a fake engine runs the policy in host tests. +pub const Engine = struct { + ctx: *anyopaque, + capture: *const fn (ctx: *anyopaque, stream: Stream, body: Body) anyerror!*anyopaque, + launch: *const fn (ctx: *anyopaque, exec: *anyopaque, stream: Stream) anyerror!void, + free: *const fn (ctx: *anyopaque, exec: *anyopaque) void, +}; + +pub const Settings = struct { + on: bool = true, // off: every step runs eagerly + max: u32 = 48, // graphs held at most + min_keep: u32 = 8, // under the memory floor the cache evicts down to this many, then stops capturing + hold: bool = false, // under the floor a miss runs eagerly with no eviction: held graphs cost next to no memory +}; + +pub const Outcome = enum { replayed, captured, eager }; + +pub fn Cache(comptime Key: type) type { + return struct { + const Self = @This(); + const Entry = struct { exec: *anyopaque, used: u64 }; + + gpa: std.mem.Allocator, + engine: Engine, + settings: Settings, + agree: ?Agree = null, + room: ?*const fn () bool = null, // false when this rank is below its memory floor; null: always room + entries: std.AutoArrayHashMapUnmanaged(Key, Entry) = .empty, + cap: u32, + tick: u64 = 0, + fingerprint: ?u64 = null, + failed: bool = false, // a capture failed on some rank: eager from here on + stats: struct { replayed: u64 = 0, captured: u64 = 0, eager: u64 = 0, evicted: u64 = 0, dropped: u64 = 0, floor_hits: u64 = 0 } = .{}, + + pub fn init(gpa: std.mem.Allocator, engine: Engine, settings: Settings) Self { + return .{ .gpa = gpa, .engine = engine, .settings = settings, .cap = settings.max }; + } + + pub fn deinit(c: *Self) void { + c.dropAll(); + c.entries.deinit(c.gpa); + } + + pub fn count(c: *const Self) usize { + return c.entries.count(); + } + + pub fn has(c: *const Self, key: Key) bool { + return c.entries.contains(key); + } + + /// Every graph freed (a fingerprint change, a shutdown). + pub fn dropAll(c: *Self) void { + for (c.entries.values()) |e| c.engine.free(c.engine.ctx, e.exec); + c.stats.dropped += c.entries.count(); + c.entries.clearRetainingCapacity(); + } + + fn agreed(c: *Self, ok: bool) !bool { + const a = c.agree orelse return ok; + return a.all(a.ctx, ok); + } + + fn evictOldest(c: *Self) void { + var at: ?usize = null; + for (c.entries.values(), 0..) |e, i| { + if (at == null or e.used < c.entries.values()[at.?].used) at = i; + } + const i = at orelse return; + c.engine.free(c.engine.ctx, c.entries.values()[i].exec); + c.entries.swapRemoveAt(i); + c.stats.evicted += 1; + } + + fn eager(c: *Self, stream: Stream, body: Body) !Outcome { + try body.run(body.ctx, stream); + c.stats.eager += 1; + return .eager; + } + + /// One step: the key's graph launched, captured first on a miss; a new `fingerprint` drops every graph first. + pub fn run(c: *Self, key: Key, stream: Stream, fingerprint: u64, body: Body) !Outcome { + if (!c.settings.on or c.failed) return c.eager(stream, body); + if (c.fingerprint != null and c.fingerprint.? != fingerprint) c.dropAll(); + c.fingerprint = fingerprint; + c.tick += 1; + if (c.entries.getPtr(key)) |e| { + e.used = c.tick; + try c.engine.launch(c.engine.ctx, e.exec, stream); + c.stats.replayed += 1; + return .replayed; + } + // a miss: room on every rank, else shed the least recent quarter and lower the cap (once below it, eager) + const room_here = if (c.room) |f| f() else true; + if (!try c.agreed(room_here and c.cap > 0)) { + if (c.settings.hold) { + if (!quiet and c.stats.floor_hits == 0) std.log.warn("graph cache: under the memory floor on {s} rank: this key eager, {d} graphs held (hold)", .{ if (room_here) "another" else "this", c.entries.count() }); + c.stats.floor_hits += 1; + return c.eager(stream, body); + } + if (c.entries.count() > c.settings.min_keep) { + const drop = @max(1, c.entries.count() / 4); + for (0..drop) |_| if (c.entries.count() > c.settings.min_keep) c.evictOldest(); + } + const was = c.cap; + c.cap = @intCast(@max(c.entries.count(), c.settings.min_keep)); + c.stats.floor_hits += 1; + if (!quiet and c.cap != was) std.log.warn("graph cache: under the memory floor on {s} rank: cap {d} -> {d} ({d} held, {d} evicted, {d} eager so far)", .{ if (room_here) "another" else "this", was, c.cap, c.entries.count(), c.stats.evicted, c.stats.eager + 1 }); + return c.eager(stream, body); + } + if (c.cap < c.settings.max) { + if (!quiet) std.log.info("graph cache: room again on every rank: cap {d} -> {d}", .{ c.cap, c.settings.max }); + c.cap = c.settings.max; + } + while (c.entries.count() >= c.cap) c.evictOldest(); + const exec = c.engine.capture(c.engine.ctx, stream, body) catch |err| blk: { + // the only trace of why graphs went off (the step then runs eagerly, which may fail the same way) + if (!quiet) std.log.warn("graph capture failed: {s} (key {any}, {d} graphs held)", .{ @errorName(err), key, c.entries.count() }); + break :blk null; + }; + if (!try c.agreed(exec != null)) { + if (exec) |x| c.engine.free(c.engine.ctx, x); + c.failed = true; + c.dropAll(); + return c.eager(stream, body); + } + try c.entries.put(c.gpa, key, .{ .exec = exec.?, .used = c.tick }); + try c.engine.launch(c.engine.ctx, exec.?, stream); + c.stats.captured += 1; + return .captured; + } + }; +} + +/// The driver's engine: thread-local stream capture (other threads' CUDA calls stay out), instantiate, upload, launch. +pub const CudaEngine = struct { + d: *const Driver, + + pub fn engine(e: *CudaEngine) Engine { + return .{ .ctx = e, .capture = capture, .launch = launch, .free = free }; + } + + const Held = struct { exec: graph.Exec }; + + fn capture(ctx: *anyopaque, stream: Stream, body: Body) anyerror!*anyopaque { + const e: *CudaEngine = @ptrCast(@alignCast(ctx)); + try graph.beginCapture(stream, .thread_local); + body.run(body.ctx, stream) catch |err| { + // end the capture so the stream is usable again; the partial graph is dropped + if (graph.endCapture(stream)) |g| { + var gg = g; + gg.deinit(); + } else |_| {} + return err; + }; + var g = try graph.endCapture(stream); + defer g.deinit(); + const h = try std.heap.c_allocator.create(Held); + errdefer std.heap.c_allocator.destroy(h); + h.exec = try g.instantiate(); + errdefer h.exec.deinit(); + try h.exec.upload(stream); + _ = e; + return h; + } + + fn launch(_: *anyopaque, exec: *anyopaque, stream: Stream) anyerror!void { + const h: *Held = @ptrCast(@alignCast(exec)); + try h.exec.launchOn(stream); + } + + fn free(_: *anyopaque, exec: *anyopaque) void { + const h: *Held = @ptrCast(@alignCast(exec)); + h.exec.deinit(); + std.heap.c_allocator.destroy(h); + } +}; + +/// FNV-1a over the addresses a body captures (buffer bases, sizes): the fingerprint `run` compares. +pub fn fingerprintOf(words: []const u64) u64 { + var h: u64 = 0xcbf29ce484222325; + for (words) |w| { + var x = w; + for (0..8) |_| { + h = (h ^ (x & 0xff)) *% 0x100000001b3; + x >>= 8; + } + } + return h; +} + +// -- host tests: the policy with a fake engine -------------------------------------------------------------------- + +const Fake = struct { + captures: u32 = 0, + launches: u32 = 0, + live: u32 = 0, + fail_capture: bool = false, + slots: [64]u8 = undefined, + + fn engine(f: *Fake) Engine { + return .{ .ctx = f, .capture = cap, .launch = lau, .free = fre }; + } + fn cap(ctx: *anyopaque, stream: Stream, body: Body) anyerror!*anyopaque { + const f: *Fake = @ptrCast(@alignCast(ctx)); + if (f.fail_capture) return error.CaptureFailed; + _ = stream; + _ = body; // a capture records the body; nothing runs + f.captures += 1; + f.live += 1; + return &f.slots[f.captures % f.slots.len]; + } + fn lau(ctx: *anyopaque, _: *anyopaque, _: Stream) anyerror!void { + const f: *Fake = @ptrCast(@alignCast(ctx)); + f.launches += 1; + } + fn fre(ctx: *anyopaque, _: *anyopaque) void { + const f: *Fake = @ptrCast(@alignCast(ctx)); + f.live -= 1; + } +}; + +const Ran = struct { + n: u32 = 0, + fn body(r: *Ran) Body { + return .{ .ctx = r, .run = go }; + } + fn go(ctx: *anyopaque, _: Stream) anyerror!void { + const r: *Ran = @ptrCast(@alignCast(ctx)); + r.n += 1; + } +}; + +const no_stream: Stream = .{ .d = undefined, .handle = null }; + +test "graph cache: capture on first use, replay after, LRU at the cap" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 2 }); + defer c.deinit(); + try std.testing.expectEqual(Outcome.captured, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.replayed, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.captured, try c.run(2, no_stream, 7, r.body())); + _ = try c.run(1, no_stream, 7, r.body()); // 1 is now the most recent + try std.testing.expectEqual(Outcome.captured, try c.run(3, no_stream, 7, r.body())); // evicts 2 + try std.testing.expect(c.has(1) and c.has(3) and !c.has(2)); + try std.testing.expectEqual(@as(u32, 0), r.n); // the body never ran eagerly + try std.testing.expectEqual(@as(u32, 2), f.live); + try std.testing.expectEqual(@as(u32, 5), f.launches); +} + +test "graph cache: a new fingerprint drops every graph" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{}); + defer c.deinit(); + _ = try c.run(1, no_stream, 7, r.body()); + _ = try c.run(2, no_stream, 7, r.body()); + try std.testing.expectEqual(Outcome.captured, try c.run(1, no_stream, 8, r.body())); + try std.testing.expectEqual(@as(usize, 1), c.count()); + try std.testing.expectEqual(@as(u32, 1), f.live); +} + +const Peer = struct { + theirs: []const bool, // the other rank's verdicts, in call order + at: usize = 0, + fn agree(p: *Peer) Agree { + return .{ .ctx = p, .all = all }; + } + fn all(ctx: *anyopaque, ok: bool) anyerror!bool { + const p: *Peer = @ptrCast(@alignCast(ctx)); + const t = p.theirs[p.at]; + p.at += 1; + return ok and t; + } +}; + +test "graph cache: a capture failing on another rank turns graphs off here too" { + var f: Fake = .{}; + var r: Ran = .{}; + // per miss: room, then the capture's verdict; the peer's capture fails on the second key + var peer: Peer = .{ .theirs = &.{ true, true, true, false } }; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{}); + c.agree = peer.agree(); + defer c.deinit(); + try std.testing.expectEqual(Outcome.captured, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.eager, try c.run(2, no_stream, 7, r.body())); + try std.testing.expect(c.failed); + try std.testing.expectEqual(@as(u32, 0), f.live); // this rank's good graphs are freed as well + try std.testing.expectEqual(Outcome.eager, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 2), r.n); +} + +var room_flag = true; +fn roomFn() bool { + return room_flag; +} + +test "graph cache: below the memory floor the least recent quarter goes, the cap follows, the step runs eagerly" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 48, .min_keep = 2 }); + c.room = roomFn; + defer c.deinit(); + room_flag = true; + for (0..8) |k| _ = try c.run(@intCast(k), no_stream, 7, r.body()); + room_flag = false; + try std.testing.expectEqual(Outcome.eager, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 6), c.count()); + try std.testing.expectEqual(@as(u32, 6), c.cap); + try std.testing.expect(!c.has(0) and !c.has(1) and c.has(7)); + room_flag = true; + // room again: the cap is back to max, the miss captures without evicting + try std.testing.expectEqual(Outcome.captured, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 7), c.count()); + try std.testing.expectEqual(@as(u32, 48), c.cap); + try std.testing.expect(c.has(7) and c.has(2)); + try std.testing.expectEqual(@as(u64, 1), c.stats.floor_hits); +} + +test "graph cache: a dip under the floor does not leave later keys churning once memory is back" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 16, .min_keep = 2 }); + c.room = roomFn; + defer c.deinit(); + room_flag = true; + for (0..12) |k| _ = try c.run(@intCast(k), no_stream, 7, r.body()); + room_flag = false; // a long prompt's workspace held + _ = try c.run(200, no_stream, 7, r.body()); + room_flag = true; // freed: the decode after it meets 12 new keys (a new context bucket), then replays them + const before = c.stats.evicted; + for (0..12) |k| _ = try c.run(@intCast(300 + k), no_stream, 7, r.body()); + for (0..12) |k| try std.testing.expectEqual(Outcome.replayed, try c.run(@intCast(300 + k), no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 16), c.cap); + // only what the cap of 16 forces: 9 left after the dip + 12 new = 21 -> 5 evictions, none during the replays + try std.testing.expectEqual(before + 5, c.stats.evicted); +} + +test "graph cache: hold - below the floor a miss runs eagerly, nothing is evicted, the held keys replay" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .max = 48, .min_keep = 2, .hold = true }); + c.room = roomFn; + defer c.deinit(); + room_flag = true; + for (0..8) |k| _ = try c.run(@intCast(k), no_stream, 7, r.body()); + room_flag = false; + try std.testing.expectEqual(Outcome.eager, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(Outcome.eager, try c.run(101, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 8), c.count()); + try std.testing.expectEqual(@as(u32, 48), c.cap); + for (0..8) |k| try std.testing.expectEqual(Outcome.replayed, try c.run(@intCast(k), no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u64, 0), c.stats.evicted); + try std.testing.expectEqual(@as(u64, 2), c.stats.floor_hits); + room_flag = true; + try std.testing.expectEqual(Outcome.captured, try c.run(100, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(usize, 9), c.count()); +} + +test "graph cache: off runs every step eagerly" { + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), .{ .on = false }); + defer c.deinit(); + try std.testing.expectEqual(Outcome.eager, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 0), f.captures); +} + +test "graph cache: fingerprint" { + try std.testing.expect(fingerprintOf(&.{ 1, 2 }) != fingerprintOf(&.{ 2, 1 })); +} diff --git a/zig/src/cuda/memory.zig b/zig/src/cuda/memory.zig index 2c469bae2c..6d802980c7 100644 --- a/zig/src/cuda/memory.zig +++ b/zig/src/cuda/memory.zig @@ -49,16 +49,37 @@ fn lockCounts() void { while (!counts_mutex.tryLock()) std.Thread.yield() catch {}; } +/// A sub-allocator under DeviceBuffer (arena.zig): `alloc` gives an address or null for cuMemAlloc, `free` owns it. +pub const Hook = struct { + ctx: *anyopaque, + alloc: *const fn (ctx: *anyopaque, len: usize) ?abi.DevicePtr, + free: *const fn (ctx: *anyopaque, p: abi.DevicePtr) bool, // true when the address was the hook's +}; + +var alloc_hook: ?Hook = null; + +/// Installs, or with null removes, the sub-allocator: once at boot, before the buffers it should hold. +pub fn setHook(h: ?Hook) void { + alloc_hook = h; +} + +pub fn currentHook() ?Hook { + return alloc_hook; +} + pub const DeviceBuffer = struct { d: *const Driver, ptr: abi.DevicePtr, len: usize, borrowed: bool = false, // a span of memory another value owns (a carveout's): free leaves it alone - /// `len` bytes, 256-byte aligned by the driver; zero bytes allocate nothing and hold address 0. + /// `len` bytes, 256-byte aligned (512 from an installed arena); zero bytes allocate nothing and hold address 0. pub fn alloc(d: *const Driver, len: usize) Error!DeviceBuffer { var p: abi.DevicePtr = 0; - if (len > 0) try d.check(d.api.cuMemAlloc_v2(&p, len), "cuMemAlloc"); + if (len > 0) { + if (alloc_hook) |h| p = h.alloc(h.ctx, len) orelse 0; + if (p == 0) try d.check(d.api.cuMemAlloc_v2(&p, len), "cuMemAlloc"); + } held(&device_bytes, len); return .{ .d = d, .ptr = p, .len = len }; } @@ -73,7 +94,10 @@ pub const DeviceBuffer = struct { pub fn free(self: *DeviceBuffer) void { if (!self.borrowed) { - if (self.ptr != 0) _ = self.d.api.cuMemFree_v2(self.ptr); + if (self.ptr != 0) { + const mine = if (alloc_hook) |h| h.free(h.ctx, self.ptr) else false; + if (!mine) _ = self.d.api.cuMemFree_v2(self.ptr); + } freed(&device_bytes, self.len); } self.* = undefined; @@ -209,6 +233,28 @@ test "freeing a borrowed span leaves the driver and the counts alone" { try std.testing.expectEqual(before.device, usage(false).device); } +test "a buffer the hook owns is freed by the hook, never by the driver" { + const Owner = struct { + var frees: u32 = 0; + fn alloc(_: *anyopaque, _: usize) ?abi.DevicePtr { + return null; + } + fn free(_: *anyopaque, p: abi.DevicePtr) bool { + frees += 1; + return p == 0x2000; + } + }; + var ctx: u8 = 0; + setHook(.{ .ctx = &ctx, .alloc = Owner.alloc, .free = Owner.free }); + defer setHook(null); + held(&device_bytes, 4096); + var b: DeviceBuffer = .{ .d = undefined, .ptr = 0x2000, .len = 4096 }; + b.free(); + var c: DeviceBuffer = .{ .d = undefined, .ptr = 0x3000, .len = 64, .borrowed = true }; + c.free(); + try std.testing.expectEqual(@as(u32, 1), Owner.frees); +} + test "peak reset preserves an allocation between its read and store" { const before = usage(true); held(&device_bytes, 100); diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 3053c09a6b..22f9454a83 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -6,6 +6,7 @@ pub const Error = @import("driver.zig").Error; pub const Context = @import("context.zig").Context; pub const DeviceBuffer = @import("memory.zig").DeviceBuffer; pub const HostBuffer = @import("memory.zig").HostBuffer; +pub const memory = @import("memory.zig"); pub const usage = @import("memory.zig").usage; pub const Usage = @import("memory.zig").Usage; pub const carveout = @import("carveout.zig"); @@ -36,6 +37,8 @@ pub const embed_rows = @import("embed_rows.zig"); pub const experts = @import("experts.zig"); pub const sample = @import("sample.zig"); pub const modelopt = @import("modelopt.zig"); +pub const graph_cache = @import("graph_cache.zig"); +pub const arena = @import("arena.zig"); test { _ = @import("memory.zig"); @@ -54,4 +57,6 @@ test { _ = carveout; _ = @import("nvfp4.zig"); _ = @import("modelopt.zig"); + _ = graph_cache; + _ = arena; } From 69e8a655adf0a37426f8b8f760c422734f140359 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:24:18 +0700 Subject: [PATCH 02/43] A failed CUDA kernel launch names its kernel and launch configuration; CUresult names fall back to a table when the driver cannot answer. cuFuncGetName is optional (CUDA 12.3+), so an older driver still opens. During the exit's teardown cuGetErrorName fails too, which printed CUDA_ERROR_DEINITIALIZED as "CUDA_ERROR_UNKNOWN_CODE (4)"; code 4 now says the process is exiting. --- zig/src/cuda/driver.zig | 67 ++++++++++++++++++++++++++++++++++++++--- zig/src/cuda/launch.zig | 36 ++++++++++++++++++++-- 2 files changed, 96 insertions(+), 7 deletions(-) diff --git a/zig/src/cuda/driver.zig b/zig/src/cuda/driver.zig index 0a860eac92..88ee9050dc 100644 --- a/zig/src/cuda/driver.zig +++ b/zig/src/cuda/driver.zig @@ -8,6 +8,10 @@ pub const Error = error{ DriverUnavailable, MissingSymbol, CudaFailed, OutOfDevi pub const Driver = struct { lib: std.DynLib, api: abi.Api, + /// cuFuncGetName (CUDA 12.3+; optional: an older driver opens without it): a failed launch names its kernel + func_name: ?FuncGetName = null, + + pub const FuncGetName = *const fn (*?[*:0]const u8, abi.Function) callconv(.c) abi.Result; pub fn open() Error!Driver { return openPath("libcuda.so.1"); @@ -25,7 +29,7 @@ pub const Driver = struct { return error.MissingSymbol; }; } - const d: Driver = .{ .lib = lib, .api = api }; + const d: Driver = .{ .lib = lib, .api = api, .func_name = lib.lookup(FuncGetName, "cuFuncGetName") }; try d.check(api.cuInit(0), "cuInit"); return d; } @@ -38,18 +42,33 @@ pub const Driver = struct { pub fn check(self: *const Driver, res: abi.Result, what: []const u8) Error!void { if (res == abi.success) return; if (res == abi.error_not_ready) return error.NotReady; - std.log.err("{s}: {s} ({d}) {s}", .{ what, self.errorName(res), res, self.errorText(res) }); + std.log.err("{s}: {s} ({d}) {s}{s}", .{ what, self.errorName(res), res, self.errorText(res), shutdownNote(res) }); + return errorOf(res); + } + + /// The error a failed call returns (`check` without its log line, for callers that log more). + pub fn errorOf(res: abi.Result) Error { return switch (res) { 2 => error.OutOfDeviceMemory, abi.error_not_found => error.NotFound, + abi.error_not_ready => error.NotReady, else => error.CudaFailed, }; } + /// The driver's name for the code, else knownName's: during the exit's teardown cuGetErrorName fails too. pub fn errorName(self: *const Driver, res: abi.Result) []const u8 { var s: ?[*:0]const u8 = null; - if (self.api.cuGetErrorName(res, &s) != abi.success) return "CUDA_ERROR_UNKNOWN_CODE"; - return if (s) |p| std.mem.span(p) else "CUDA_ERROR_UNKNOWN_CODE"; + if (self.api.cuGetErrorName(res, &s) == abi.success) if (s) |p| return std.mem.span(p); + return knownName(res) orelse "CUDA_ERROR_UNKNOWN_CODE"; + } + + /// A kernel's symbol name (cuFuncGetName), or "?" when the driver cannot say. + pub fn functionName(self: *const Driver, f: abi.Function) []const u8 { + const get = self.func_name orelse return "?"; + var s: ?[*:0]const u8 = null; + if (get(&s, f) != abi.success) return "?"; + return if (s) |p| std.mem.span(p) else "?"; } pub fn errorText(self: *const Driver, res: abi.Result) []const u8 { @@ -82,3 +101,43 @@ pub const Driver = struct { return @intCast(10 * major + minor); } }; + +/// CUresult names (cuda.h) for when cuGetErrorName cannot answer. +pub fn knownName(res: abi.Result) ?[]const u8 { + return switch (res) { + 1 => "CUDA_ERROR_INVALID_VALUE", + 2 => "CUDA_ERROR_OUT_OF_MEMORY", + 3 => "CUDA_ERROR_NOT_INITIALIZED", + 4 => "CUDA_ERROR_DEINITIALIZED", + 101 => "CUDA_ERROR_INVALID_DEVICE", + 200 => "CUDA_ERROR_INVALID_IMAGE", + 201 => "CUDA_ERROR_INVALID_CONTEXT", + 209 => "CUDA_ERROR_NO_BINARY_FOR_GPU", + 400 => "CUDA_ERROR_INVALID_HANDLE", + 500 => "CUDA_ERROR_NOT_FOUND", + 600 => "CUDA_ERROR_NOT_READY", + 700 => "CUDA_ERROR_ILLEGAL_ADDRESS", + 701 => "CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES", + 702 => "CUDA_ERROR_LAUNCH_TIMEOUT", + 709 => "CUDA_ERROR_CONTEXT_IS_DESTROYED", + 710 => "CUDA_ERROR_ASSERT", + 714 => "CUDA_ERROR_HARDWARE_STACK_ERROR", + 715 => "CUDA_ERROR_ILLEGAL_INSTRUCTION", + 716 => "CUDA_ERROR_MISALIGNED_ADDRESS", + 719 => "CUDA_ERROR_LAUNCH_FAILED", + 999 => "CUDA_ERROR_UNKNOWN", + else => null, + }; +} + +/// CUDA_ERROR_DEINITIALIZED: the process is exiting (a stop signal, then exit) while a thread still calls the driver. +pub fn shutdownNote(res: abi.Result) []const u8 { + return if (res == 4) " (the CUDA driver is shutting down: the process is exiting)" else ""; +} + +test "CUresult names without the driver: DEINITIALIZED is 4, unknown codes stay unknown" { + try std.testing.expectEqualStrings("CUDA_ERROR_DEINITIALIZED", knownName(4).?); + try std.testing.expectEqualStrings("CUDA_ERROR_ILLEGAL_ADDRESS", knownName(700).?); + try std.testing.expect(knownName(12345) == null); + try std.testing.expect(shutdownNote(4).len > 0 and shutdownNote(1).len == 0); +} diff --git a/zig/src/cuda/launch.zig b/zig/src/cuda/launch.zig index 56bba5272a..fe161fabf1 100644 --- a/zig/src/cuda/launch.zig +++ b/zig/src/cuda/launch.zig @@ -4,7 +4,8 @@ const std = @import("std"); const abi = @import("abi.zig"); const Function = @import("module.zig").Function; const Stream = @import("stream.zig").Stream; -const Error = @import("driver.zig").Error; +const driver = @import("driver.zig"); +const Error = driver.Error; pub const Dim3 = abi.Dim3; @@ -88,7 +89,7 @@ pub fn launch(f: Function, cfg: Config, stream: Stream, args: *Args) Error!void const g = cfg.grid; const b = cfg.block; if (cfg.plain()) { - return d.check(d.api.cuLaunchKernel(f.handle, g.x, g.y, g.z, b.x, b.y, b.z, cfg.shared, stream.handle, args.pointers(), null), "cuLaunchKernel"); + return launched(f, cfg, d.api.cuLaunchKernel(f.handle, g.x, g.y, g.z, b.x, b.y, b.z, cfg.shared, stream.handle, args.pointers(), null), "cuLaunchKernel"); } var attrs: [4]abi.LaunchAttribute = undefined; const n = fillAttributes(cfg, &attrs); @@ -104,7 +105,36 @@ pub fn launch(f: Function, cfg: Config, stream: Stream, args: *Args) Error!void .attrs = &attrs, .num_attrs = @intCast(n), }; - try d.check(d.api.cuLaunchKernelEx(&lc, f.handle, args.pointers(), null), "cuLaunchKernelEx"); + try launched(f, cfg, d.api.cuLaunchKernelEx(&lc, f.handle, args.pointers(), null), "cuLaunchKernelEx"); +} + +/// A launch's result; a failure names the kernel and its configuration with the code. +fn launched(f: Function, cfg: Config, res: abi.Result, what: []const u8) Error!void { + if (res == abi.success) return; + const d = f.d; + var buf: [160]u8 = undefined; + std.log.err("{s}: {s} ({d}) {s}; {s}{s}", .{ what, d.errorName(res), res, d.errorText(res), describe(&buf, d.functionName(f.handle), cfg), driver.shutdownNote(res) }); + return driver.Driver.errorOf(res); +} + +/// "kernel , grid XxYxZ, block XxYxZ, shared N B" (+ cluster / PDL / cooperative when set). +pub fn describe(buf: []u8, name: []const u8, cfg: Config) []const u8 { + const g = cfg.grid; + const b = cfg.block; + var w = std.Io.Writer.fixed(buf); + w.print("kernel {s}, grid {d}x{d}x{d}, block {d}x{d}x{d}, shared {d} B", .{ name, g.x, g.y, g.z, b.x, b.y, b.z, cfg.shared }) catch return w.buffered(); + if (cfg.cluster) |c| w.print(", cluster {d}x{d}x{d}", .{ c.x, c.y, c.z }) catch return w.buffered(); + if (cfg.pdl) w.writeAll(", pdl") catch return w.buffered(); + if (cfg.cooperative) w.writeAll(", cooperative") catch return w.buffered(); + return w.buffered(); +} + +test "a failed launch's description: the kernel, its grid, block and shared bytes, bounded" { + var buf: [160]u8 = undefined; + const cfg: Config = .{ .grid = .{ .x = 70000, .y = 2, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 }, .shared = 49152, .pdl = true }; + try std.testing.expectEqualStrings("kernel _stream_pf, grid 70000x2x1, block 128x1x1, shared 49152 B, pdl", describe(&buf, "_stream_pf", cfg)); + var small: [24]u8 = undefined; + try std.testing.expect(describe(&small, "a_very_long_kernel_symbol_name", cfg).len <= small.len); } fn attribute(id: abi.LaunchAttributeId) abi.LaunchAttribute { From d7e89dff4fa4b3f4ac699a6cb2de51fe0e111bed Mon Sep 17 00:00:00 2001 From: Jonathan Regan-Beasley <197533392+jregan-beasley-smc@users.noreply.github.com> Date: Sun, 11 Oct 2026 10:23:24 +1100 Subject: [PATCH 03/43] fix(cuda): correct RDMA VMM capability attribute Use CUDA device attribute 110 for GPUDirect RDMA with VMM support. Attribute 126 reports the memory sync domain count and can incorrectly enable the RDMA allocation flag. Add a regression for the external ABI contract. --- zig/src/cuda/abi.zig | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/zig/src/cuda/abi.zig b/zig/src/cuda/abi.zig index a099686c67..5418ce0109 100644 --- a/zig/src/cuda/abi.zig +++ b/zig/src/cuda/abi.zig @@ -40,7 +40,7 @@ pub const DeviceAttribute = enum(c_int) { max_shared_memory_per_multiprocessor = 81, max_shared_memory_per_block_optin = 97, virtual_memory_management_supported = 102, - gpu_direct_rdma_with_cuda_vmm_supported = 126, + gpu_direct_rdma_with_cuda_vmm_supported = 110, }; pub const FunctionAttribute = enum(c_int) { @@ -243,3 +243,10 @@ comptime { std.debug.assert(@sizeOf(KernelNodeParams) == 72 and @offsetOf(KernelNodeParams, "params") == 40); std.debug.assert(@sizeOf(ExecUpdateResultInfo) == 24); } + +test "CUDA device attribute matches the RDMA VMM capability ABI" { + // CUDA's CUdevice_attribute assigns 110 to RDMA VMM support; 126 is the memory sync domain count. + try @import("std").testing.expectEqual(@as(c_int, 110), @backingInt( + DeviceAttribute.gpu_direct_rdma_with_cuda_vmm_supported, + )); +} From 5826f89bd31ce0f4d875051b61dec78cfd764ae7 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 17:29:31 +0700 Subject: [PATCH 04/43] Graph cache: a capture is freed until the map holds it, the keep floor stays under the max, a max of 0 is eager - A capture whose agreement throws, or whose insert runs out of memory, no longer leaks its executable: the map's slot is reserved before the capture, and the executable is owned here until the insert, which cannot fail. - A host map that cannot grow counts as no room in the existing room vote, so every rank declines that miss together. - `min_keep` is clamped to `max`, so under the memory floor the cap never rises above the configured max. A max of 0 runs every step eagerly, with no collective. - A cap lowered to 0 under the floor goes back to max once every rank has room, instead of staying eager for good. - Hits make no collective: with `Agree` set, keys, settings and fingerprint changes must already match on every rank. The docs now say so: the fingerprint changes at the same step on every rank, and `fingerprintOf` is for one rank. Co-authored-by: Jonathan Regan-Beasley <197533392+jregan-beasley-smc@users.noreply.github.com> --- zig/src/cuda/graph_cache.zig | 31 ++-- zig/src/cuda/graph_cache_regression_test.zig | 181 +++++++++++++++++++ 2 files changed, 201 insertions(+), 11 deletions(-) create mode 100644 zig/src/cuda/graph_cache_regression_test.zig diff --git a/zig/src/cuda/graph_cache.zig b/zig/src/cuda/graph_cache.zig index a2fb8e18bb..693f8ce25a 100644 --- a/zig/src/cuda/graph_cache.zig +++ b/zig/src/cuda/graph_cache.zig @@ -12,7 +12,7 @@ pub const Body = struct { run: *const fn (ctx: *anyopaque, stream: Stream) anyerror!void, }; -/// Every rank's verdict AND-ed, so ranks capture, evict and fall back at the same steps; never called in a capture. +/// Every rank's verdict AND-ed on a miss (hits never call it), so ranks capture, evict and fall back together. pub const Agree = struct { ctx: *anyopaque, all: *const fn (ctx: *anyopaque, ok: bool) anyerror!bool, @@ -53,7 +53,9 @@ pub fn Cache(comptime Key: type) type { stats: struct { replayed: u64 = 0, captured: u64 = 0, eager: u64 = 0, evicted: u64 = 0, dropped: u64 = 0, floor_hits: u64 = 0 } = .{}, pub fn init(gpa: std.mem.Allocator, engine: Engine, settings: Settings) Self { - return .{ .gpa = gpa, .engine = engine, .settings = settings, .cap = settings.max }; + var s = settings; + s.min_keep = @min(s.min_keep, s.max); // the floor never holds more than the configured max + return .{ .gpa = gpa, .engine = engine, .settings = s, .cap = s.max }; } pub fn deinit(c: *Self) void { @@ -98,9 +100,9 @@ pub fn Cache(comptime Key: type) type { return .eager; } - /// One step: the key's graph launched, captured first on a miss; a new `fingerprint` drops every graph first. + /// One step: the key's graph, captured first on a miss; a new `fingerprint` (same step on all ranks) drops all. pub fn run(c: *Self, key: Key, stream: Stream, fingerprint: u64, body: Body) !Outcome { - if (!c.settings.on or c.failed) return c.eager(stream, body); + if (!c.settings.on or c.failed or c.settings.max == 0) return c.eager(stream, body); if (c.fingerprint != null and c.fingerprint.? != fingerprint) c.dropAll(); c.fingerprint = fingerprint; c.tick += 1; @@ -110,9 +112,10 @@ pub fn Cache(comptime Key: type) type { c.stats.replayed += 1; return .replayed; } - // a miss: room on every rank, else shed the least recent quarter and lower the cap (once below it, eager) - const room_here = if (c.room) |f| f() else true; - if (!try c.agreed(room_here and c.cap > 0)) { + // a miss: room on every rank (a host map that cannot grow is none), else shed a quarter and lower the cap + const reserved = if (c.entries.ensureUnusedCapacity(c.gpa, 1)) |_| true else |_| false; + const room_here = reserved and (if (c.room) |f| f() else true); + if (!try c.agreed(room_here)) { if (c.settings.hold) { if (!quiet and c.stats.floor_hits == 0) std.log.warn("graph cache: under the memory floor on {s} rank: this key eager, {d} graphs held (hold)", .{ if (room_here) "another" else "this", c.entries.count() }); c.stats.floor_hits += 1; @@ -123,7 +126,7 @@ pub fn Cache(comptime Key: type) type { for (0..drop) |_| if (c.entries.count() > c.settings.min_keep) c.evictOldest(); } const was = c.cap; - c.cap = @intCast(@max(c.entries.count(), c.settings.min_keep)); + c.cap = @intCast(@max(c.entries.count(), c.settings.min_keep)); // min_keep <= max (init) c.stats.floor_hits += 1; if (!quiet and c.cap != was) std.log.warn("graph cache: under the memory floor on {s} rank: cap {d} -> {d} ({d} held, {d} evicted, {d} eager so far)", .{ if (room_here) "another" else "this", was, c.cap, c.entries.count(), c.stats.evicted, c.stats.eager + 1 }); return c.eager(stream, body); @@ -138,13 +141,15 @@ pub fn Cache(comptime Key: type) type { if (!quiet) std.log.warn("graph capture failed: {s} (key {any}, {d} graphs held)", .{ @errorName(err), key, c.entries.count() }); break :blk null; }; + var owned = exec; // freed here until the map holds it, so a failed agreement cannot leak it + defer if (owned) |x| c.engine.free(c.engine.ctx, x); if (!try c.agreed(exec != null)) { - if (exec) |x| c.engine.free(c.engine.ctx, x); c.failed = true; c.dropAll(); return c.eager(stream, body); } - try c.entries.put(c.gpa, key, .{ .exec = exec.?, .used = c.tick }); + c.entries.putAssumeCapacityNoClobber(key, .{ .exec = exec.?, .used = c.tick }); + owned = null; try c.engine.launch(c.engine.ctx, exec.?, stream); c.stats.captured += 1; return .captured; @@ -196,7 +201,7 @@ pub const CudaEngine = struct { } }; -/// FNV-1a over the addresses a body captures (buffer bases, sizes): the fingerprint `run` compares. +/// FNV-1a over the addresses a body captures (buffer bases, sizes): the fingerprint `run` compares, on one rank only. pub fn fingerprintOf(words: []const u64) u64 { var h: u64 = 0xcbf29ce484222325; for (words) |w| { @@ -390,3 +395,7 @@ test "graph cache: off runs every step eagerly" { test "graph cache: fingerprint" { try std.testing.expect(fingerprintOf(&.{ 1, 2 }) != fingerprintOf(&.{ 2, 1 })); } + +test { + _ = @import("graph_cache_regression_test.zig"); +} diff --git a/zig/src/cuda/graph_cache_regression_test.zig b/zig/src/cuda/graph_cache_regression_test.zig new file mode 100644 index 0000000000..7f2da1063a --- /dev/null +++ b/zig/src/cuda/graph_cache_regression_test.zig @@ -0,0 +1,181 @@ +//! Failure ownership, rank verdict counts and bounded-cap regressions for the graph cache. + +const std = @import("std"); +const cache = @import("graph_cache.zig"); +const Stream = @import("stream.zig").Stream; +const no_stream: Stream = .{ .d = undefined, .handle = null }; +const expect = std.testing.expect; +const equal = std.testing.expectEqual; + +const Engine = struct { + captures: usize = 0, + live: usize = 0, + launches: usize = 0, + bodies: usize = 0, + handles: [32]u8 = undefined, + + fn engine(e: *Engine) cache.Engine { + return .{ .ctx = e, .capture = capture, .launch = launch, .free = free }; + } + fn capture(p: *anyopaque, _: Stream, _: cache.Body) !*anyopaque { + const e: *Engine = @ptrCast(@alignCast(p)); + e.captures += 1; + e.live += 1; + return &e.handles[e.captures % e.handles.len]; + } + fn launch(p: *anyopaque, _: *anyopaque, _: Stream) !void { + const e: *Engine = @ptrCast(@alignCast(p)); + e.launches += 1; + } + fn free(p: *anyopaque, _: *anyopaque) void { + const e: *Engine = @ptrCast(@alignCast(p)); + std.debug.assert(e.live > 0); + e.live -= 1; + } + fn body(e: *Engine) cache.Body { + return .{ .ctx = e, .run = run }; + } + fn run(p: *anyopaque, _: Stream) !void { + const e: *Engine = @ptrCast(@alignCast(p)); + e.bodies += 1; + } +}; + +const Peer = struct { + verdicts: []const bool, + at: usize = 0, + throw_at: ?usize = null, + fn agree(p: *Peer) cache.Agree { + return .{ .ctx = p, .all = all }; + } + fn all(ctx: *anyopaque, ok: bool) !bool { + const p: *Peer = @ptrCast(@alignCast(ctx)); + const at = p.at; + p.at += 1; + if (p.throw_at == at) return error.CollectiveFailed; + if (at >= p.verdicts.len) return error.UnexpectedCollective; + return ok and p.verdicts[at]; + } +}; + +test "graph cache regression: capture agreement throw frees the not-yet-inserted executable" { + var e: Engine = .{}; + // per miss: room, then the capture's verdict + var p: Peer = .{ .verdicts = &.{ true, true }, .throw_at = 1 }; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{}); + defer c.deinit(); + c.agree = p.agree(); + try std.testing.expectError(error.CollectiveFailed, c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 1), e.captures); + try equal(@as(usize, 0), e.live); + try equal(@as(usize, 0), c.count()); + try equal(@as(usize, 2), p.at); +} + +test "graph cache regression: local map OOM declines before any capture and can recover" { + var e: Engine = .{}; + var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = 0 }); + var c = cache.Cache(u32).init(failing.allocator(), e.engine(), .{}); + defer c.deinit(); + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try expect(failing.has_induced_failure); + try equal(@as(usize, 0), e.captures); + try equal(@as(usize, 0), e.live); + try expect(!c.failed); + failing.fail_index = std.math.maxInt(usize); + try equal(cache.Outcome.captured, try c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 1), e.live); +} + +test "graph cache regression: peer map OOM makes a locally prepared rank eager before capture" { + var e: Engine = .{}; + var p: Peer = .{ .verdicts = &.{false} }; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{}); + defer c.deinit(); + c.agree = p.agree(); + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 1), p.at); + try equal(@as(usize, 0), e.captures); + try equal(@as(usize, 1), e.bodies); +} + +test "graph cache regression: local OOM still contributes the room verdict" { + var e: Engine = .{}; + var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = 0 }); + var p: Peer = .{ .verdicts = &.{true} }; + var c = cache.Cache(u32).init(failing.allocator(), e.engine(), .{}); + defer c.deinit(); + c.agree = p.agree(); + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 1), p.at); + try equal(@as(usize, 0), e.captures); +} + +test "graph cache regression: hits make no collective; a fingerprint change shared by every rank recaptures" { + var e: Engine = .{}; + var p: Peer = .{ .verdicts = &.{ true, true, true, true, true, true } }; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{}); + defer c.deinit(); + c.agree = p.agree(); + _ = try c.run(1, no_stream, 7, e.body()); + _ = try c.run(2, no_stream, 7, e.body()); + try equal(@as(usize, 4), p.at); + for (0..4) |_| try equal(cache.Outcome.replayed, try c.run(1, no_stream, 7, e.body())); + try equal(@as(usize, 4), p.at); + // a new layout generation on every rank: both graphs go, the next miss agrees twice as any other + try equal(cache.Outcome.captured, try c.run(1, no_stream, 8, e.body())); + try equal(@as(usize, 6), p.at); + try equal(@as(u64, 2), c.stats.dropped); + try equal(@as(usize, 1), c.count()); + try equal(@as(usize, 1), e.live); +} + +var room = true; +fn available() bool { + return room; +} + +test "graph cache regression: zero maximum remains eager across pressure recovery, with no collective" { + var e: Engine = .{}; + var p: Peer = .{ .verdicts = &.{} }; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{ .max = 0 }); + defer c.deinit(); + c.room = available; + c.agree = p.agree(); + for ([_]bool{ false, true, false, true }) |ok| { + room = ok; + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try equal(@as(u32, 0), c.cap); + try equal(@as(usize, 0), c.count()); + } + try equal(@as(usize, 0), e.captures); + try equal(@as(usize, 0), p.at); +} + +test "graph cache regression: default floor cannot enlarge a two-entry maximum" { + var e: Engine = .{}; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{ .max = 2 }); + defer c.deinit(); + c.room = available; + try equal(@as(u32, 2), c.settings.min_keep); + room = false; + _ = try c.run(1, no_stream, 7, e.body()); + room = true; + for (0..8) |key| { + _ = try c.run(@intCast(key), no_stream, 7, e.body()); + try expect(c.count() <= 2 and c.cap <= 2 and e.live <= 2); + } +} + +test "graph cache regression: zero keep floor recovers after reducing the cap to zero" { + var e: Engine = .{}; + var c = cache.Cache(u32).init(std.testing.allocator, e.engine(), .{ .max = 2, .min_keep = 0 }); + defer c.deinit(); + c.room = available; + room = false; + try equal(cache.Outcome.eager, try c.run(1, no_stream, 7, e.body())); + try equal(@as(u32, 0), c.cap); + room = true; + try equal(cache.Outcome.captured, try c.run(1, no_stream, 7, e.body())); + try equal(@as(u32, 2), c.cap); +} From 5806e551cdcdcee9bff7cb90c596379826681d0e Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:33:36 +0700 Subject: [PATCH 05/43] CUDA arena: the address maps are rehashed before tombstones fill them The arena's `live` (small ranges) and `large` (own chunks) maps drop an entry on every free and gain one on every allocation for the life of the process. std's tombstones never clear, so in time they take every free slot, and each allocation, free and miss then probes the whole map. Each map now counts its removals and is rehashed in place once they reach a tenth of its capacity. Test: 20,000 random allocations and frees (small ranges and large chunks) keep a tenth of each map free at every step; without the rehash it fails. --- zig/src/cuda/arena.zig | 43 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/zig/src/cuda/arena.zig b/zig/src/cuda/arena.zig index 2c80b7581b..3886c81403 100644 --- a/zig/src/cuda/arena.zig +++ b/zig/src/cuda/arena.zig @@ -36,6 +36,7 @@ pub const Arena = struct { free_small: std.AutoHashMapUnmanaged(u64, std.ArrayList(u64)) = .empty, // freed small ranges by rounded size live: std.AutoHashMapUnmanaged(u64, u64) = .empty, // live small ranges: address to rounded size large: std.AutoHashMapUnmanaged(u64, u64) = .empty, // large chunks: address to mapped size + dead: [2]u32 = .{ 0, 0 }, // removals from `live` and `large` since each was last rehashed stats: Stats = .{}, pub fn init(gpa: std.mem.Allocator, be: Backend, o: Options) Arena { @@ -114,16 +115,26 @@ pub const Arena = struct { a.lock(); defer a.mutex.unlock(); if (a.large.fetchRemove(p)) |e| { + removed(&a.large, &a.dead[1]); a.be.release(a.be.ctx, p, e.value); return true; } const e = a.live.fetchRemove(p) orelse return false; + removed(&a.live, &a.dead[0]); const g = a.free_small.getOrPut(a.gpa, e.value) catch return true; // leaked range, still ours if (!g.found_existing) g.value_ptr.* = .empty; g.value_ptr.append(a.gpa, p) catch {}; return true; } + /// std's tombstones never clear and keys fill 80 %: rehash at a tenth, so a tenth of the map always stays free. + fn removed(map: *std.AutoHashMapUnmanaged(u64, u64), dead: *u32) void { + dead.* += 1; + if (@as(u64, dead.*) * 10 < map.capacity()) return; + map.rehash(std.hash_map.AutoContext(u64){}); + dead.* = 0; + } + // -- memory.zig's hook ---------------------------------------------------------------------------------------- fn hookAlloc(ctx: *anyopaque, len: usize) ?u64 { @@ -329,3 +340,35 @@ test "arena: under DeviceBuffer.alloc through memory.setHook" { try t.expect(h.free(h.ctx, p)); try t.expect(!h.free(h.ctx, 7)); } + +/// Slots of `m` that hold neither a key nor a tombstone. +fn freeSlots(m: *const std.AutoHashMapUnmanaged(u64, u64)) usize { + const meta = m.metadata orelse return 0; + var n: usize = 0; + for (meta[0..m.capacity()]) |x| n += @intFromBool(x.isFree()); + return n; +} + +test "arena: a long run of allocations and frees leaves both address maps free slots" { + const t = std.testing; + var f: Fake = .{}; + var a = Arena.init(t.allocator, f.backend(2 << 20), .{ .small_max = 1 << 20, .small_chunk = 4 << 20 }); + defer a.deinit(); + var held: std.ArrayList(u64) = .empty; + defer held.deinit(t.allocator); + var rng = std.Random.DefaultPrng.init(3); + for (0..20000) |_| { + // a serving process's buffers come and go (prompt workspaces, grown buffers): small and large alike + if (held.items.len > 0 and rng.random().boolean()) { + const k = rng.random().uintLessThan(usize, held.items.len); + try t.expect(a.free(held.swapRemove(k))); + } else { + const big = rng.random().uintLessThan(u32, 4) == 0; + try held.append(t.allocator, a.alloc(if (big) 3 << 20 else 512 * (1 + rng.random().uintLessThan(u64, 64))).?); + } + for ([_]*const std.AutoHashMapUnmanaged(u64, u64){ &a.live, &a.large }) |m| { + if (m.capacity() == 0) continue; + try t.expect(freeSlots(m) * 10 >= m.capacity()); + } + } +} From d335520663539b54ee1cada6b32a599c3e2a2bee Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Wed, 14 Oct 2026 21:05:29 +0700 Subject: [PATCH 06/43] Graph cache: TF_GRAPH_CACHE=0 or off runs every step eagerly in any family that installs it, without a rebuild Cache.init applies the operator's value through Settings.operator, so a deployment that measures graphs slower can turn them off at serve time. The value is read once a cache, and is set alike on every rank like the settings themselves: no agreement is added, and hits still make no collective. --- RUNBOOK.md | 2 ++ zig/src/cuda/graph_cache.zig | 25 ++++++++++++++++++++++++- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/RUNBOOK.md b/RUNBOOK.md index 408f2d91e8..525cdf49e5 100644 --- a/RUNBOOK.md +++ b/RUNBOOK.md @@ -157,6 +157,8 @@ budget free for the next states. `TF_HEAT_HIGH` and `TF_HEAT_LOW`, in degrees Celsius and set together, make each prompt chunk wait while the hottest thermal zone (under `/sys/class/thermal`, or `TF_HEAT_ROOT`) is above the high band, until it is at or under the low one. The server logs `heat_wait_s`, a cancel ends the wait, and the reply's tokens do not change. +`TF_GRAPH_CACHE=0` (or `off`) runs every step eagerly in any family that installs the keyed CUDA graph cache, without a +rebuild; set it alike on every tensor-parallel rank. Archives whose names contain `host-only` are CPU verification artifacts and cannot serve CUDA inference. ## Flash Next and paired Metal serving diff --git a/zig/src/cuda/graph_cache.zig b/zig/src/cuda/graph_cache.zig index 693f8ce25a..9c62e4fc7d 100644 --- a/zig/src/cuda/graph_cache.zig +++ b/zig/src/cuda/graph_cache.zig @@ -31,6 +31,14 @@ pub const Settings = struct { max: u32 = 48, // graphs held at most min_keep: u32 = 8, // under the memory floor the cache evicts down to this many, then stops capturing hold: bool = false, // under the floor a miss runs eagerly with no eviction: held graphs cost next to no memory + + /// These settings under the operator's TF_GRAPH_CACHE value: `0` or `off` runs every step eagerly; else unchanged. + pub fn operator(s: Settings, value: ?[]const u8) Settings { + var out = s; + const v = value orelse return out; + if (std.mem.eql(u8, v, "0") or std.ascii.eqlIgnoreCase(v, "off")) out.on = false; + return out; + } }; pub const Outcome = enum { replayed, captured, eager }; @@ -52,8 +60,10 @@ pub fn Cache(comptime Key: type) type { failed: bool = false, // a capture failed on some rank: eager from here on stats: struct { replayed: u64 = 0, captured: u64 = 0, eager: u64 = 0, evicted: u64 = 0, dropped: u64 = 0, floor_hits: u64 = 0 } = .{}, + /// The cache with `settings`, unless TF_GRAPH_CACHE turns it off (set it alike on every rank, as `settings`). pub fn init(gpa: std.mem.Allocator, engine: Engine, settings: Settings) Self { - var s = settings; + var s = settings.operator(if (std.c.getenv("TF_GRAPH_CACHE")) |v| std.mem.span(v) else null); + if (!quiet and settings.on and !s.on) std.log.info("graph cache: off (TF_GRAPH_CACHE), every step eager", .{}); s.min_keep = @min(s.min_keep, s.max); // the floor never holds more than the configured max return .{ .gpa = gpa, .engine = engine, .settings = s, .cap = s.max }; } @@ -392,6 +402,19 @@ test "graph cache: off runs every step eagerly" { try std.testing.expectEqual(@as(u32, 0), f.captures); } +test "graph cache: TF_GRAPH_CACHE=0 or off turns a family's cache off without a rebuild; anything else leaves it" { + const on: Settings = .{}; + for ([_][]const u8{ "0", "off", "OFF" }) |v| try std.testing.expect(!on.operator(v).on); + for ([_]?[]const u8{ null, "1", "on", "" }) |v| try std.testing.expect(on.operator(v).on); + try std.testing.expect(!(Settings{ .on = false }).operator("1").on); // the operator only turns a cache off + var f: Fake = .{}; + var r: Ran = .{}; + var c = Cache(u32).init(std.testing.allocator, f.engine(), on.operator("off")); + defer c.deinit(); + for (0..3) |_| try std.testing.expectEqual(Outcome.eager, try c.run(1, no_stream, 7, r.body())); + try std.testing.expectEqual(@as(u32, 0), f.captures); +} + test "graph cache: fingerprint" { try std.testing.expect(fingerprintOf(&.{ 1, 2 }) != fingerprintOf(&.{ 2, 1 })); } From 7ff83d8057b36b87293f8e843c6994828e1fab36 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:38:24 +0700 Subject: [PATCH 07/43] Tensor parallelism across CUDA processes: bootstrap, NCCL, mailbox and RoCE collectives, plan link, fail-fast. zig/src/fabric/tp brings a rank up from TF_TP_WORLD, TF_TP_RANK, TF_TP_MASTER and TF_TP_PORT, agrees the settings digest over every rank, and gives a family all-gather, all-reduce and broadcast over NCCL (TF_COMM_BACKEND=nccl, the default), a pinned host mailbox for ranks on one host, or RoCE verbs between two hosts (TF_TP_ROCE_*), with an NCCL fallback when RoCE setup fails. Rank 0 sends each round's plan to the others over its own link (TF_TP_PLAN_*). A failure or a vanished rank ends every rank shortly after (TF_TP_FAILFAST, TF_TP_FAILFAST_GRACE_S). No family uses it yet. zig build test-tp runs the host tests; zig build tf-tp-test builds the two-GPU runner (-Dtp-nvcc= or -Dtp-fatbin= for the mailbox kernel). The RoCE all-gather adapts b12x's RoCEnante protocol. --- THIRD_PARTY_NOTICES.md | 6 + build.zig | 2 + zig/build/tp.zig | 59 ++++ zig/src/cuda/nccl.zig | 7 + zig/src/fabric/tp/bootstrap.zig | 262 +++++++++++++++ zig/src/fabric/tp/collective.zig | 136 ++++++++ zig/src/fabric/tp/config.zig | 245 ++++++++++++++ zig/src/fabric/tp/fate.zig | 261 +++++++++++++++ zig/src/fabric/tp/fate_test.zig | 146 ++++++++ zig/src/fabric/tp/gpu_bench.zig | 82 +++++ zig/src/fabric/tp/gpu_checks.zig | 188 +++++++++++ zig/src/fabric/tp/gpu_main.zig | 92 +++++ zig/src/fabric/tp/host.zig | 293 ++++++++++++++++ zig/src/fabric/tp/host_test.zig | 154 +++++++++ zig/src/fabric/tp/hybrid.zig | 167 ++++++++++ zig/src/fabric/tp/log.zig | 15 + zig/src/fabric/tp/mailbox.cu | 556 +++++++++++++++++++++++++++++++ zig/src/fabric/tp/mailbox.zig | 215 ++++++++++++ zig/src/fabric/tp/nccl.zig | 163 +++++++++ zig/src/fabric/tp/planlink.zig | 467 ++++++++++++++++++++++++++ zig/src/fabric/tp/reference.zig | 121 +++++++ zig/src/fabric/tp/roce.zig | 458 +++++++++++++++++++++++++ zig/src/fabric/tp/roce_host.zig | 58 ++++ zig/src/fabric/tp/roce_test.zig | 116 +++++++ zig/src/fabric/tp/roce_verbs.zig | 297 +++++++++++++++++ zig/src/fabric/tp/run_hosts.sh | 24 ++ zig/src/fabric/tp/run_pair.sh | 52 +++ zig/src/fabric/tp/sock.zig | 169 ++++++++++ zig/src/fabric/tp/tp.zig | 212 ++++++++++++ zig/src/fabric/verbs.zig | 2 +- 30 files changed, 5024 insertions(+), 1 deletion(-) create mode 100644 zig/build/tp.zig create mode 100644 zig/src/fabric/tp/bootstrap.zig create mode 100644 zig/src/fabric/tp/collective.zig create mode 100644 zig/src/fabric/tp/config.zig create mode 100644 zig/src/fabric/tp/fate.zig create mode 100644 zig/src/fabric/tp/fate_test.zig create mode 100644 zig/src/fabric/tp/gpu_bench.zig create mode 100644 zig/src/fabric/tp/gpu_checks.zig create mode 100644 zig/src/fabric/tp/gpu_main.zig create mode 100644 zig/src/fabric/tp/host.zig create mode 100644 zig/src/fabric/tp/host_test.zig create mode 100644 zig/src/fabric/tp/hybrid.zig create mode 100644 zig/src/fabric/tp/log.zig create mode 100644 zig/src/fabric/tp/mailbox.cu create mode 100644 zig/src/fabric/tp/mailbox.zig create mode 100644 zig/src/fabric/tp/nccl.zig create mode 100644 zig/src/fabric/tp/planlink.zig create mode 100644 zig/src/fabric/tp/reference.zig create mode 100644 zig/src/fabric/tp/roce.zig create mode 100644 zig/src/fabric/tp/roce_host.zig create mode 100644 zig/src/fabric/tp/roce_test.zig create mode 100644 zig/src/fabric/tp/roce_verbs.zig create mode 100644 zig/src/fabric/tp/run_hosts.sh create mode 100644 zig/src/fabric/tp/run_pair.sh create mode 100644 zig/src/fabric/tp/sock.zig create mode 100644 zig/src/fabric/tp/tp.zig diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index b482c601a3..6be05bb01a 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -125,6 +125,12 @@ values agree within fp16/bf16 rounding (under 0.01 on random groups), not bit fo dequant are written for TensorFold and checked against an independent reference of that arithmetic. The native engine has no quantized KV cache and refuses `--kv-dtype`. +The Zig tensor-parallel fabric's RoCE path (`zig/src/fabric/tp/mailbox.cu`, `roce.zig`, `roce_verbs.zig`) adapts +the one-shot all-gather protocol of [b12x](https://github.com/local-inference-lab/b12x)'s RoCEnante +(`b12x/comm/roce/`, PRs #295 and #315): a device-resident epoch, sequence flags after each payload, two slots by +sequence parity and a timeout that fails the transport. Copyright 2026 Luke Alonso and the b12x contributors, +Apache License 2.0 (see [the license text](LICENSES/Apache-2.0.txt)); b12x ships no NOTICE file. + ## Vendored code and weights In the Python engine, frozen at 0.6.6, `src/tensorfold/drafters/vendor/z_lab_dflash/model_mlx.py` is the unmodified diff --git a/build.zig b/build.zig index 08765bd1df..d9ec51a4e0 100644 --- a/build.zig +++ b/build.zig @@ -12,6 +12,7 @@ const counters_build = @import("zig/build/metal_counters.zig"); const exl3_build = @import("zig/build/exl3.zig"); const exlr8_build = @import("zig/build/exlr8.zig"); const golden_build = @import("zig/build/golden.zig"); +const tp_build = @import("zig/build/tp.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -189,6 +190,7 @@ pub fn build(b: *std.Build) void { test_step.dependOn(&warning_run.step); b.step("test-learn-warning", "CPU production-log refusal, backoff and recovery test").dependOn(&warning_run.step); grammar_build.steps(b, target, xgrammar); + tp_build.add(b, target, optimize, test_step); } /// `zig build native -Dcpu=apple_m1`: tensorfold-native with the Metal engines for the Python package's bundle (a native M5 build traps on M1-M4). diff --git a/zig/build/tp.zig b/zig/build/tp.zig new file mode 100644 index 0000000000..f5dacc46e3 --- /dev/null +++ b/zig/build/tp.zig @@ -0,0 +1,59 @@ +//! Tensor parallelism's build: host tests on any OS, and on Linux the `tf-tp-test` GPU runner with the mailbox kernel. + +const std = @import("std"); + +/// The CUDA runtime without the engine's kernels (tp only needs the driver, memory, streams and NCCL). +fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + const options = b.addOptions(); + options.addOption(bool, "with_kernels", false); + const cuda = b.createModule(.{ .root_source_file = b.path("zig/src/cuda/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); + cuda.addOptions("kernel_options", options); + cuda.addImport("stagger", b.createModule(.{ .root_source_file = b.path("zig/src/core/stagger.zig"), .target = target, .optimize = optimize })); + return cuda; +} + +/// The `tp` module over `cuda` (an engine passes its own runtime module); verbs (zig/src/fabric/verbs.zig) for RoCE. +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/tp/tp.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addImport("cuda", cuda); + m.addImport("verbs", b.createModule(.{ .root_source_file = b.path("zig/src/fabric/verbs.zig"), .target = target, .optimize = optimize, .link_libc = true })); + return m; +} + +/// mailbox.cu as a fatbin: sm_120 (RTX PRO 6000), sm_121 (GB10), and compute_90 PTX for anything else. +fn mailboxFatbin(b: *std.Build, nvcc: []const u8) std.Build.LazyPath { + const run = b.addSystemCommand(&.{ nvcc, "-fatbin", "-O3", "-std=c++17", "-gencode=arch=compute_120,code=sm_120", "-gencode=arch=compute_121,code=sm_121", "-gencode=arch=compute_90,code=compute_90", "-o" }); + const out = run.addOutputFileArg("tp_mailbox.fatbin"); + run.addFileArg(b.path("zig/src/fabric/tp/mailbox.cu")); + return out; +} + +/// `pub const mailbox: []const u8`, the fatbin's bytes or empty. +fn kernelModule(b: *std.Build, image: ?std.Build.LazyPath) *std.Build.Module { + const wf = b.addWriteFiles(); + if (image) |file| { + _ = wf.addCopyFile(file, "tp_mailbox.fatbin"); + // the driver takes 8-byte aligned images (Module.load refuses others); @embedFile alone promises nothing + return b.createModule(.{ .root_source_file = wf.add("tp_kernel.zig", "const bytes align(16) = @embedFile(\"tp_mailbox.fatbin\").*;\npub const mailbox: []const u8 = &bytes;\n") }); + } + return b.createModule(.{ .root_source_file = wf.add("tp_kernel.zig", "pub const mailbox: []const u8 = \"\";\n") }); +} + +pub fn add(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, test_step: *std.Build.Step) void { + const host = b.graph.host; + const tests = b.addRunArtifact(b.addTest(.{ .root_module = module(b, host, .debug, runtime(b, host, .debug)) })); + b.step("test-tp", "Tensor parallelism's host tests: bootstrap, fate channel, in-process ranks (no GPU)").dependOn(&tests.step); + test_step.dependOn(&tests.step); + if (target.result.os.tag != .linux) return; + const nvcc = b.option([]const u8, "tp-nvcc", "nvcc that builds the TP mailbox kernel"); + const prebuilt = b.option([]const u8, "tp-fatbin", "absolute path of a prebuilt tp_mailbox.fatbin to embed"); + const image: ?std.Build.LazyPath = if (prebuilt) |p| b.graph.cwdRelativePath(p) else if (nvcc) |tool| mailboxFatbin(b, tool) else null; + const cuda = runtime(b, target, optimize); + const runner = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/tp/gpu_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + runner.addImport("cuda", cuda); + runner.addImport("tp", module(b, target, optimize, cuda)); + runner.addImport("tp_kernel", kernelModule(b, image)); + const exe = b.addExecutable(.{ .name = "tf-tp-test", .root_module = runner }); + b.installArtifact(exe); + b.step("tf-tp-test", "TP=2 on two GPUs: bootstrap, every collective bit-exact, graph capture, fail-fast, latency").dependOn(&b.addInstallArtifact(exe, .{}).step); +} diff --git a/zig/src/cuda/nccl.zig b/zig/src/cuda/nccl.zig index f742d62549..85f88bace8 100644 --- a/zig/src/cuda/nccl.zig +++ b/zig/src/cuda/nccl.zig @@ -12,6 +12,8 @@ pub const RedOp = enum(c_int) { sum = 0, prod = 1, max = 2, min = 3, avg = 4 }; pub const Error = error{ LibraryUnavailable, MissingSymbol, NcclFailed }; +pub const in_progress: Result = 7; // ncclInProgress: a non-blocking communicator's call is still running + const R = Result; const D = abi.DevicePtr; @@ -25,6 +27,11 @@ pub const Api = struct { ncclAllReduce: *const fn (D, D, usize, DataType, RedOp, Comm, abi.Stream) callconv(.c) R, ncclAllGather: *const fn (D, D, usize, DataType, Comm, abi.Stream) callconv(.c) R, ncclBroadcast: *const fn (D, D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, + ncclSend: *const fn (D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, + ncclRecv: *const fn (D, usize, DataType, c_int, Comm, abi.Stream) callconv(.c) R, + /// Frees the communicator after ending its pending work; safe from another thread while a kernel waits. + ncclCommAbort: *const fn (Comm) callconv(.c) R, + ncclCommGetAsyncError: *const fn (Comm, *Result) callconv(.c) R, ncclGroupStart: *const fn () callconv(.c) R, ncclGroupEnd: *const fn () callconv(.c) R, }; diff --git a/zig/src/fabric/tp/bootstrap.zig b/zig/src/fabric/tp/bootstrap.zig new file mode 100644 index 0000000000..1486e3a8f7 --- /dev/null +++ b/zig/src/fabric/tp/bootstrap.zig @@ -0,0 +1,262 @@ +//! Rank discovery over TCP (a star on rank 0): small host blobs all-gathered, settings agreed, sockets kept for fate. +const std = @import("std"); +const log = @import("log.zig"); +const sock = @import("sock.zig"); +const Config = @import("config.zig").Config; + +pub const max_ranks = 64; +/// Largest blob a rank contributes to one `allGather`. +pub const max_blob = 64 * 1024; + +const magic: u32 = 0x31425054; // "TPB1" +const version: u32 = 2; + +pub const Error = error{ Protocol, Mismatch, TooLarge } || sock.Error; + +/// `role`: 0 the control connection (bootstrap, then the fate channel), 1 the plan link. +const Hello = extern struct { magic: u32, version: u32, world: u32, rank: u32, role: u32 }; + +pub const Bootstrap = struct { + rank: u32, + world: u32, + /// Rank 0: the socket of rank i at index i (index 0 unused); other ranks: rank 0's socket at index 0. + fds: [max_ranks]sock.Fd = @splat(-1), + /// The plan link's sockets, laid out as `fds`. + plan_fds: [max_ranks]sock.Fd = @splat(-1), + /// Each call's limit for the peers' bytes (the connect timeout). + timeout_ns: u64, + + fn deadline(b: *const Bootstrap) u64 { + return sock.nowNs() + b.timeout_ns; + } + + /// Rank 0 listens on `link_port` and accepts every other rank; the others connect to `master` (connect timeout). + pub fn open(cfg: Config) Error!Bootstrap { + if (cfg.world > max_ranks or cfg.rank >= cfg.world) return error.Protocol; + var b: Bootstrap = .{ .rank = cfg.rank, .world = cfg.world, .timeout_ns = cfg.connect_timeout_ns }; + errdefer b.close(); + if (cfg.world == 1) return b; + if (cfg.rank == 0) { + const l = try sock.listen(cfg.link_port); + defer sock.close(l); + try b.acceptAll(l); + } else { + const until = b.deadline(); + for ([_]u32{ 0, 1 }) |role| { + const fd = try sock.connect(cfg.master, cfg.link_port, until); + if (role == 0) b.fds[0] = fd else b.plan_fds[0] = fd; + const h: Hello = .{ .magic = magic, .version = version, .world = cfg.world, .rank = cfg.rank, .role = role }; + try sock.sendAll(fd, std.mem.asBytes(&h)); + } + } + return b; + } + + /// Rank 0 over an already listening socket (tests: an ephemeral port). + pub fn acceptOn(l: sock.Fd, world: u32, timeout_ns: u64) Error!Bootstrap { + var b: Bootstrap = .{ .rank = 0, .world = world, .timeout_ns = timeout_ns }; + errdefer b.close(); + try b.acceptAll(l); + return b; + } + + fn acceptAll(b: *Bootstrap, l: sock.Fd) Error!void { + // every other rank's control and plan connections, in any order + var joined: u32 = 0; + const until = b.deadline(); + while (joined < 2 * (b.world - 1)) { + const fd = try sock.accept(l, until); + var h: Hello = undefined; + sock.recvExact(fd, std.mem.asBytes(&h), until) catch { + sock.close(fd); + continue; + }; + if (h.magic != magic or h.version != version) { + log.warn("tp bootstrap: a connection that is not a rank (magic {x}); ignored", .{h.magic}); + sock.close(fd); + continue; + } + const table = if (h.role == 0) &b.fds else &b.plan_fds; + if (h.world != b.world or h.rank == 0 or h.rank >= b.world or h.role > 1 or table[h.rank] >= 0) { + log.warn("tp bootstrap: rank {d} of {d} joined a world of {d}", .{ h.rank, h.world, b.world }); + sock.close(fd); + return error.Mismatch; + } + table[h.rank] = fd; + joined += 1; + } + } + + /// Takes the control sockets (the fate channel keeps them); `close` then leaves them open. + pub fn release(b: *Bootstrap) [max_ranks]sock.Fd { + const out = b.fds; + b.fds = @splat(-1); + return out; + } + + /// Takes the plan link's sockets. + pub fn releasePlan(b: *Bootstrap) [max_ranks]sock.Fd { + const out = b.plan_fds; + b.plan_fds = @splat(-1); + return out; + } + + pub fn close(b: *Bootstrap) void { + for ([_]*[max_ranks]sock.Fd{ &b.fds, &b.plan_fds }) |table| for (table) |*fd| if (fd.* >= 0) { + sock.close(fd.*); + fd.* = -1; + }; + } + + /// Every rank's `mine` in rank order into `out`, returned as slices of `out`; every rank calls it equally often. + pub fn allGather(b: *Bootstrap, mine: []const u8, out: []u8, blobs: [][]const u8) Error![]const []const u8 { + if (mine.len > max_blob) return error.TooLarge; + std.debug.assert(blobs.len >= b.world); + if (b.world == 1) { + if (out.len < mine.len) return error.TooLarge; + @memcpy(out[0..mine.len], mine); + blobs[0] = out[0..mine.len]; + return blobs[0..1]; + } + if (b.rank != 0) { + try sendFrame(b.fds[0], mine); + return b.readAll(b.fds[0], out, blobs); + } + // rank 0: everyone's blob, then the whole table to everyone + var used: usize = 0; + for (0..b.world) |i| { + const len: u32 = if (i == 0) @intCast(mine.len) else blk: { + var l: u32 = undefined; + try sock.recvExact(b.fds[i], std.mem.asBytes(&l), b.deadline()); + break :blk l; + }; + if (len > max_blob or used + len > out.len) return error.TooLarge; + if (i == 0) @memcpy(out[used..][0..len], mine) else try sock.recvExact(b.fds[i], out[used..][0..len], b.deadline()); + blobs[i] = out[used..][0..len]; + used += len; + } + for (1..b.world) |i| { + for (blobs[0..b.world]) |blob| try sendFrame(b.fds[i], blob); + } + return blobs[0..b.world]; + } + + fn readAll(b: *Bootstrap, fd: sock.Fd, out: []u8, blobs: [][]const u8) Error![]const []const u8 { + var used: usize = 0; + for (0..b.world) |i| { + var l: u32 = undefined; + try sock.recvExact(fd, std.mem.asBytes(&l), b.deadline()); + if (l > max_blob or used + l > out.len) return error.TooLarge; + try sock.recvExact(fd, out[used..][0..l], b.deadline()); + blobs[i] = out[used..][0..l]; + used += l; + } + return blobs[0..b.world]; + } + + /// Rank 0's `bytes` on every rank (into `out` on the others; `bytes` is ignored there). + pub fn broadcast(b: *Bootstrap, bytes: []const u8, out: []u8) Error!void { + var blobs: [max_ranks][]const u8 = undefined; + var scratch: [4096]u8 = undefined; + if (out.len * b.world > scratch.len) return error.TooLarge; + const all = try b.allGather(if (b.rank == 0) bytes else &.{}, &scratch, &blobs); + if (all[0].len != out.len) return error.Protocol; + @memcpy(out, all[0]); + } + + /// Every rank's `bytes` must be equal (settings that change bytes on the wire); `what` names them in the log. + pub fn agree(b: *Bootstrap, bytes: []const u8, what: []const u8) Error!void { + var blobs: [max_ranks][]const u8 = undefined; + var scratch: [8192]u8 = undefined; + const all = try b.allGather(bytes, &scratch, &blobs); + for (all, 0..) |blob, i| if (!std.mem.eql(u8, blob, all[0])) { + log.warn("tp bootstrap: rank {d}'s {s} differs from rank 0's; both ranks refuse to start", .{ i, what }); + return error.Mismatch; + }; + } +}; + +fn sendFrame(fd: sock.Fd, bytes: []const u8) sock.Error!void { + const l: u32 = @intCast(bytes.len); + try sock.sendAll(fd, std.mem.asBytes(&l)); + try sock.sendAll(fd, bytes); +} + +/// An ephemeral listening port for tests. +pub fn testListener() !struct { fd: sock.Fd, port: u16 } { + const l = try sock.listen(0); + var addr: std.posix.sockaddr.in = undefined; + var len: std.posix.socklen_t = @sizeOf(std.posix.sockaddr.in); + _ = std.c.getsockname(l, @ptrCast(&addr), &len); + return .{ .fd = l, .port = std.mem.bigToNative(u16, addr.port) }; +} + +fn testRank(cfg: Config, out: *[3][32]u8, err: *?anyerror) void { + var b = Bootstrap.open(cfg) catch |e| { + err.* = e; + return; + }; + defer b.close(); + var id: [32]u8 = undefined; + b.broadcast(&@as([32]u8, @splat(0xAB)), &id) catch |e| { + err.* = e; + return; + }; + var blobs: [max_ranks][]const u8 = undefined; + var buf: [256]u8 = undefined; + const mine: [3]u8 = @splat(@intCast(cfg.rank)); + const all = b.allGather(mine[0 .. cfg.rank + 1], &buf, &blobs) catch |e| { + err.* = e; + return; + }; + for (all, 0..) |blob, i| @memcpy(out[i][0..blob.len], blob); + b.agree("same", "settings") catch |e| { + err.* = e; + }; +} + +test "three ranks gather in rank order and agree" { + const l = try testListener(); + defer sock.close(l.fd); + var outs: [3][3][32]u8 = undefined; + var errs: [3]?anyerror = @splat(null); + var threads: [2]std.Thread = undefined; + for (0..2) |i| threads[i] = try std.Thread.spawn(.{}, testRank, .{ Config{ .world = 3, .rank = @intCast(i + 1), .link_port = l.port, .connect_timeout_ns = 5 * std.time.ns_per_s }, &outs[i + 1], &errs[i + 1] }); + var b = try Bootstrap.acceptOn(l.fd, 3, 5 * std.time.ns_per_s); + defer b.close(); + var id: [32]u8 = undefined; + try b.broadcast(&@as([32]u8, @splat(0xAB)), &id); + var blobs: [max_ranks][]const u8 = undefined; + var buf: [256]u8 = undefined; + const all = try b.allGather(&.{0}, &buf, &blobs); + try b.agree("same", "settings"); + for (threads) |t| t.join(); + for (errs) |e| try std.testing.expect(e == null); + try std.testing.expectEqual(@as(usize, 3), all.len); + try std.testing.expectEqualSlices(u8, &.{ 2, 2, 2 }, all[2]); + try std.testing.expectEqualSlices(u8, &.{ 1, 1 }, outs[2][1][0..2]); +} + +test "settings that differ stop every rank" { + const l = try testListener(); + defer sock.close(l.fd); + const T = struct { + fn go(port: u16, err: *?anyerror) void { + var b = Bootstrap.open(.{ .world = 2, .rank = 1, .link_port = port, .connect_timeout_ns = 5 * std.time.ns_per_s }) catch |e| { + err.* = e; + return; + }; + defer b.close(); + b.agree("knobs-b", "settings") catch |e| { + err.* = e; + }; + } + }; + var err: ?anyerror = null; + const t = try std.Thread.spawn(.{}, T.go, .{ l.port, &err }); + var b = try Bootstrap.acceptOn(l.fd, 2, 5 * std.time.ns_per_s); + defer b.close(); + try std.testing.expectError(error.Mismatch, b.agree("knobs-a", "settings")); + t.join(); + try std.testing.expectEqual(@as(?anyerror, error.Mismatch), err); +} diff --git a/zig/src/fabric/tp/collective.zig b/zig/src/fabric/tp/collective.zig new file mode 100644 index 0000000000..d2334dc74a --- /dev/null +++ b/zig/src/fabric/tp/collective.zig @@ -0,0 +1,136 @@ +//! The tensor-parallel communicator every backend implements: stream-ordered, graph-capturable, rank-order sums. +const std = @import("std"); +const cuda = @import("cuda"); + +pub const DevicePtr = cuda.abi.DevicePtr; +pub const Stream = cuda.abi.Stream; + +pub const Error = error{ Aborted, Unsupported, Invalid, PeerFailed, Timeout, BackendFailed }; + +/// Element types the collectives move (`reduce` takes the float ones and the 32/64-bit integers). +pub const DType = enum(u8) { + u8, + i32, + i64, + f16, + bf16, + f32, + f64, + + pub fn size(t: DType) usize { + return switch (t) { + .u8 => 1, + .f16, .bf16 => 2, + .i32, .f32 => 4, + .i64, .f64 => 8, + }; + } +}; + +pub const Op = enum(u8) { sum, max, min }; + +/// Which implementation is behind a `Collective` (logs, bench labels). +pub const Kind = enum { nccl, mailbox, hybrid, roce, host }; + +pub const Collective = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + kind: Kind, + rank: *const fn (ptr: *anyopaque) u32, + world: *const fn (ptr: *anyopaque) u32, + /// recv[count] <- the element-wise reduction of every rank's send[count], added in rank order. + all_reduce: *const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, count: usize, t: DType, op: Op, stream: Stream) Error!void, + /// recv[world * count] <- every rank's send[count], rank 0's first (NCCL's layout). + all_gather: *const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, count: usize, t: DType, stream: Stream) Error!void, + /// Two ranks trade: this rank's send[count] lands in `peer`'s recv[count] and the peer's in ours (one group). + exchange: *const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void, + send: *const fn (ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void, + recv: *const fn (ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void, + /// Host: returns once every rank has called it and this rank's earlier collectives have finished. + barrier: *const fn (ptr: *anyopaque) Error!void, + /// Host: a failure the transport recorded since the last check (a peer that timed out, an NCCL async error). + check: *const fn (ptr: *anyopaque) Error!void, + /// Any thread: end pending and later work at once; safe while another thread is inside a call. + abort: *const fn (ptr: *anyopaque) void, + /// recv[world][stride]: rank r's first lens[r] bytes (device int32 [world]) land at recv + r * stride. + all_gather_v: ?*const fn (ptr: *anyopaque, send: DevicePtr, recv: DevicePtr, stride: usize, lens: DevicePtr, t: DType, stream: Stream) Error!void = null, + /// Whether `all_gather_v` moves only the lengths for shards of `bytes` here (rank-local). Null: never. + var_fits: ?*const fn (ptr: *anyopaque, bytes: usize) bool = null, + /// Host: whether every rank's transport takes device lengths, agreed once per communicator. Null: false. + var_agreed: ?*const fn (ptr: *anyopaque) Error!bool = null, + }; + + pub fn kind(c: Collective) Kind { + return c.vtable.kind; + } + + pub fn rank(c: Collective) u32 { + return c.vtable.rank(c.ptr); + } + + pub fn world(c: Collective) u32 { + return c.vtable.world(c.ptr); + } + + pub fn allReduce(c: Collective, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, op: Op, stream: Stream) Error!void { + return c.vtable.all_reduce(c.ptr, src, dst, count, t, op, stream); + } + + pub fn allGather(c: Collective, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, stream: Stream) Error!void { + return c.vtable.all_gather(c.ptr, src, dst, count, t, stream); + } + + /// `allGather` of `stride` elements a rank where only rank r's first lens[r] bytes matter (`VTable.all_gather_v`). + pub fn allGatherV(c: Collective, src: DevicePtr, dst: DevicePtr, stride: usize, lens: DevicePtr, t: DType, stream: Stream) Error!void { + if (c.vtable.all_gather_v) |f| return f(c.ptr, src, dst, stride, lens, t, stream); + return c.vtable.all_gather(c.ptr, src, dst, stride, t, stream); + } + + /// Whether `allGatherV` of `bytes`-byte strides moves only the lengths here (callers agree first: `varAgreed`). + pub fn varFits(c: Collective, bytes: usize) bool { + const f = c.vtable.var_fits orelse return false; + return f(c.ptr, bytes); + } + + /// Every rank's transport takes device lengths: one agreement a communicator, called alike outside any graph. + pub fn varAgreed(c: Collective) Error!bool { + const f = c.vtable.var_agreed orelse return false; + return f(c.ptr); + } + + pub fn exchange(c: Collective, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void { + return c.vtable.exchange(c.ptr, src, dst, count, t, peer, stream); + } + + pub fn send(c: Collective, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void { + return c.vtable.send(c.ptr, buf, count, t, peer, stream); + } + + pub fn recv(c: Collective, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: Stream) Error!void { + return c.vtable.recv(c.ptr, buf, count, t, peer, stream); + } + + pub fn barrier(c: Collective) Error!void { + return c.vtable.barrier(c.ptr); + } + + pub fn check(c: Collective) Error!void { + return c.vtable.check(c.ptr); + } + + pub fn abort(c: Collective) void { + c.vtable.abort(c.ptr); + } +}; + +/// Refuses a peer that is this rank or out of range. +pub fn checkPeer(rank: u32, world: u32, peer: u32) Error!void { + if (peer == rank or peer >= world) return error.Invalid; +} + +test "dtype sizes" { + try std.testing.expectEqual(@as(usize, 2), DType.bf16.size()); + try std.testing.expectEqual(@as(usize, 8), DType.i64.size()); +} diff --git a/zig/src/fabric/tp/config.zig b/zig/src/fabric/tp/config.zig new file mode 100644 index 0000000000..13647f5921 --- /dev/null +++ b/zig/src/fabric/tp/config.zig @@ -0,0 +1,245 @@ +//! A rank's TP settings from the environment: rank, master, port, backend, fail-fast knobs. +const std = @import("std"); +const log = @import("log.zig"); + +pub const Backend = enum { + /// NCCL for everything (the default). + nccl, + /// Exchanges of at most `mailbox_max_bytes` a rank through the one-shot mailbox kernel, NCCL above that. + mailbox, + /// The same split between two hosts: the small exchanges over RoCE (doorbell + RDMA proxy), NCCL above. + roce, +}; + +/// What `roce` does when its setup or load-time probe fails on any rank (TF_TP_ROCE_FALLBACK). +pub const Fallback = enum { nccl, @"error" }; + +pub const Error = error{BadSetting}; + +/// TF_TP_PLAN_PIN: the CPU the plan threads pin to at their first plan (auto: the highest cpu_capacity). +pub const PlanPin = union(enum) { auto, cpu: u32 }; + +pub const Config = struct { + world: u32 = 2, + rank: u32 = 0, + /// Rank 0's address as the other ranks reach it. + master: []const u8 = "localhost", + /// The rendezvous port (TF_TP_PORT); the link listens on `link_port`. + port: u16 = 29500, + /// Rank 0 listens here: TF_TP_PLAN_PORT, else `port` + 7. + link_port: u16 = 29507, + /// How long rank 0 waits for the others and they keep retrying the connect (TF_TP_PLAN_TIMEOUT). + connect_timeout_ns: u64 = 300 * std.time.ns_per_s, + /// CUDA device ordinal this rank uses. + device: u32 = 0, + backend: Backend = .nccl, + mailbox_max_bytes: usize = 256 * 1024, + /// A mailbox wait for the peer gives up after this long and poisons the transport (fail-stop). + mailbox_timeout_ns: u64 = 120 * std.time.ns_per_s, + failfast: bool = true, + /// After a failure: time for the HTTP threads to answer before exit 70. + grace_ns: u64 = std.time.ns_per_s, + /// Rank 1's memory report period on the fate channel (also how fast it notices rank 0 hanging up). + report_ns: u64 = std.time.ns_per_s / 2, + nccl_lib: []const u8 = "libnccl.so.2", + /// `name[:gid],...` RDMA devices for `roce` (null: every active Ethernet port with an IPv4 RoCE v2 GID). + roce_hca: ?[]const u8 = null, + /// Stripe over at most this many HCAs (auto detection only). + roce_hcas: u8 = 2, + roce_tc: u8 = 0, + /// TF_TP_ROCE_FAST as the RoCE kernel's `opts` word: bit 0 on, bits 8-15 pollers, bits 16-31 stagger ns (local). + roce_fast: u32 = 0, + /// TF_TP_ROCE_BLOCKS: the largest RoCE grid, a power of two 1..16; launch shape only, the bytes are the same. + roce_blocks: u32 = 16, + /// Pin the proxy thread (it busy-spins one core between exchanges). + roce_cpu: ?u32 = null, + /// `shm` runs the RoCE protocol on one host without RDMA (tests); `verbs` is the real wire. + roce_wire: enum { verbs, shm } = .verbs, + roce_fallback: Fallback = .nccl, + /// TF_TP_PLAN_LINK=rdma: the follower polls the plan socket this long before blocking (TF_TP_PLAN_SPIN_US). + plan_spin_ns: u64 = 0, + plan_pin: ?PlanPin = null, + + /// Reads every setting from `get` (the environment in `fromEnv`); unset names keep the defaults. + pub fn parse(get: *const fn ([]const u8) ?[]const u8) Error!Config { + var c: Config = .{}; + c.world = try int(u32, get, "TF_TP_WORLD", c.world, 1, 64); + c.rank = try int(u32, get, "TF_TP_RANK", c.rank, 0, c.world - 1); + if (get("TF_TP_MASTER")) |m| c.master = m; + c.port = try int(u16, get, "TF_TP_PORT", c.port, 1, 65535 - 8); + c.link_port = try int(u16, get, "TF_TP_PLAN_PORT", c.port + 7, 1, 65535); + c.connect_timeout_ns = try seconds(get, "TF_TP_PLAN_TIMEOUT", 300, 0.1, 86400); + c.device = try int(u32, get, "TF_TP_DEVICE", c.device, 0, 63); + const backend = "TF_COMM_BACKEND"; + if (get(backend)) |b| c.backend = std.meta.stringToEnum(Backend, std.mem.trim(u8, b, " ")) orelse { + log.warn("{s}={s}: expected nccl, mailbox or roce", .{ backend, b }); + return error.BadSetting; + }; + c.mailbox_max_bytes = 1024 * try int(usize, get, "TF_TP_MAILBOX_MAX_KB", 256, 4, 16 * 1024); + c.mailbox_timeout_ns = try seconds(get, "TF_TP_MAILBOX_TIMEOUT_S", 120, 0.001, 86400); + if (get("TF_TP_FAILFAST")) |v| c.failfast = !(std.mem.eql(u8, v, "0") or std.mem.eql(u8, v, "off")); + c.grace_ns = try seconds(get, "TF_TP_FAILFAST_GRACE_S", 1.0, 0, 60); + c.report_ns = try seconds(get, "TF_TP_PEER_MEM_S", 0.5, 0.01, 60); + if (get("TF_NCCL_LIB")) |p| c.nccl_lib = p; + c.roce_hca = get("TF_TP_ROCE_HCA"); + c.roce_hcas = try int(u8, get, "TF_TP_ROCE_HCAS", 2, 1, 2); + c.roce_tc = try int(u8, get, "TF_TP_ROCE_TC", try int(u8, get, "NCCL_IB_TC", 0, 0, 255), 0, 255); + if (try int(u8, get, "TF_TP_ROCE_FAST", 0, 0, 1) == 1) { + const pollers = try int(u32, get, "TF_TP_ROCE_POLLERS", 1, 1, 8); + const stagger = try int(u32, get, "TF_TP_ROCE_STAGGER_NS", 150, 0, 65535); + c.roce_fast = 1 | (pollers << 8) | (stagger << 16); + } + c.roce_blocks = try int(u32, get, "TF_TP_ROCE_BLOCKS", 16, 1, 16); + if (c.roce_blocks & (c.roce_blocks - 1) != 0) { + log.warn("TF_TP_ROCE_BLOCKS={d}: expected a power of two", .{c.roce_blocks}); + return error.BadSetting; + } + if (get("TF_TP_ROCE_CPU")) |_| c.roce_cpu = try int(u32, get, "TF_TP_ROCE_CPU", 0, 0, 4095); + if (get("TF_TP_ROCE_WIRE")) |w| c.roce_wire = std.meta.stringToEnum(@TypeOf(c.roce_wire), w) orelse { + log.warn("TF_TP_ROCE_WIRE={s}: expected verbs or shm", .{w}); + return error.BadSetting; + }; + const fallback = "TF_TP_ROCE_FALLBACK"; + if (get(fallback)) |f| c.roce_fallback = std.meta.stringToEnum(Fallback, f) orelse { + log.warn("{s}={s}: expected nccl or error", .{ fallback, f }); + return error.BadSetting; + }; + if (get("TF_TP_PLAN_LINK")) |l| { + const mode = std.mem.trim(u8, l, " "); + if (std.mem.eql(u8, mode, "rdma")) { + c.plan_spin_ns = 1000 * try int(u64, get, "TF_TP_PLAN_SPIN_US", 5000, 0, 10_000_000); + } else if (!std.mem.eql(u8, mode, "tcp") and !std.mem.eql(u8, mode, "nccl")) { + log.warn("TF_TP_PLAN_LINK={s}: expected nccl, tcp or rdma", .{l}); + return error.BadSetting; + } + } + if (get("TF_TP_PLAN_PIN")) |raw| { + const v = std.mem.trim(u8, raw, " "); + if (std.mem.eql(u8, v, "auto")) { + c.plan_pin = .auto; + } else if (!std.mem.eql(u8, v, "-1") and !std.mem.eql(u8, v, "off")) { + c.plan_pin = .{ .cpu = try int(u32, get, "TF_TP_PLAN_PIN", 0, 0, 4095) }; + } + } + return c; + } + + pub fn fromEnv() Error!Config { + return parse(&envGet); + } + + /// The settings both ranks must share (everything but rank, device and timeouts), as bytes to compare. + pub fn digest(c: Config) [32]u8 { + var h = std.crypto.hash.sha2.Sha256.init(.{}); + h.update(std.mem.asBytes(&c.world)); + h.update(std.mem.asBytes(&c.backend)); + h.update(std.mem.asBytes(&c.mailbox_max_bytes)); + h.update(std.mem.asBytes(&c.failfast)); + h.update(std.mem.asBytes(&c.roce_wire)); + h.update(std.mem.asBytes(&c.roce_fallback)); + var out: [32]u8 = undefined; + h.final(&out); + return out; + } +}; + +fn envGet(name: []const u8) ?[]const u8 { + var buf: [64]u8 = undefined; + if (name.len >= buf.len) return null; + @memcpy(buf[0..name.len], name); + buf[name.len] = 0; + const v = std.c.getenv(buf[0..name.len :0]) orelse return null; + const s = std.mem.span(v); + return if (s.len == 0) null else s; +} + +fn int(comptime T: type, get: *const fn ([]const u8) ?[]const u8, name: []const u8, default: T, lo: T, hi: T) Error!T { + const raw = get(name) orelse return default; + const v = std.fmt.parseInt(T, std.mem.trim(u8, raw, " "), 10) catch { + log.warn("{s}={s}: expected an integer", .{ name, raw }); + return error.BadSetting; + }; + if (v < lo or v > hi) { + log.warn("{s}={d}: expected {d}..{d}", .{ name, v, lo, hi }); + return error.BadSetting; + } + return v; +} + +fn seconds(get: *const fn ([]const u8) ?[]const u8, name: []const u8, default: f64, lo: f64, hi: f64) Error!u64 { + const v = if (get(name)) |raw| std.fmt.parseFloat(f64, std.mem.trim(u8, raw, " ")) catch { + log.warn("{s}={s}: expected seconds", .{ name, raw }); + return error.BadSetting; + } else default; + if (!(v >= lo and v <= hi)) { + log.warn("{s}={d}: expected {d}..{d} s", .{ name, v, lo, hi }); + return error.BadSetting; + } + return @intFromFloat(v * std.time.ns_per_s); +} + +const Fake = struct { + var pairs: []const [2][]const u8 = &.{}; + fn get(name: []const u8) ?[]const u8 { + for (pairs) |p| if (std.mem.eql(u8, p[0], name)) return p[1]; + return null; + } +}; + +test "defaults: the plan link 7 ports past the rendezvous, NCCL, fail-fast on" { + Fake.pairs = &.{ .{ "TF_TP_RANK", "1" }, .{ "TF_TP_PORT", "29600" } }; + const c = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(u32, 1), c.rank); + try std.testing.expectEqual(@as(u16, 29607), c.link_port); + try std.testing.expectEqual(Backend.nccl, c.backend); + try std.testing.expectEqual(std.time.ns_per_s, c.grace_ns); + try std.testing.expectEqual(std.time.ns_per_s / 2, c.report_ns); +} + +test "bad settings are refused and the digest ignores the rank" { + Fake.pairs = &.{.{ "TF_TP_RANK", "2" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{.{ "TF_COMM_BACKEND", "rdma" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_ROCE_HCA", "rocep1s0f1,roceP2p1s0f1" }, .{ "NCCL_IB_TC", "106" } }; + const r = try Config.parse(&Fake.get); + try std.testing.expectEqual(Backend.roce, r.backend); + try std.testing.expectEqual(@as(u8, 106), r.roce_tc); + try std.testing.expectEqualStrings("rocep1s0f1,roceP2p1s0f1", r.roce_hca.?); + Fake.pairs = &.{ .{ "TF_COMM_BACKEND", "mailbox" }, .{ "TF_TP_MAILBOX_MAX_KB", "64" }, .{ "TF_TP_FAILFAST_GRACE_S", "0.25" } }; + const a = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(usize, 64 * 1024), a.mailbox_max_bytes); + try std.testing.expectEqual(std.time.ns_per_s / 4, a.grace_ns); + var b = a; + b.rank = 1; + try std.testing.expectEqualSlices(u8, &a.digest(), &b.digest()); + b.mailbox_max_bytes = 1024; + try std.testing.expect(!std.mem.eql(u8, &a.digest(), &b.digest())); +} + +test "the plan link's spin and pin, the RoCE grid cap" { + Fake.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_MAILBOX_MAX_KB", "1024" }, .{ "TF_TP_ROCE_FALLBACK", "error" }, .{ "TF_TP_PLAN_LINK", "rdma" }, .{ "TF_TP_PLAN_PIN", "auto" } }; + const p = try Config.parse(&Fake.get); + try std.testing.expectEqual(Backend.roce, p.backend); + try std.testing.expectEqual(@as(usize, 1024 * 1024), p.mailbox_max_bytes); + try std.testing.expectEqual(Fallback.@"error", p.roce_fallback); + try std.testing.expectEqual(@as(u64, 5000 * 1000), p.plan_spin_ns); + try std.testing.expectEqual(PlanPin.auto, p.plan_pin.?); + Fake.pairs = &.{ .{ "TF_TP_PLAN_LINK", "tcp" }, .{ "TF_TP_PLAN_PIN", "7" } }; // tcp blocks at once; an explicit CPU + const o = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(u64, 0), o.plan_spin_ns); + try std.testing.expectEqual(PlanPin{ .cpu = 7 }, o.plan_pin.?); + Fake.pairs = &.{ .{ "TF_TP_PLAN_LINK", "rdma" }, .{ "TF_TP_PLAN_SPIN_US", "200" }, .{ "TF_TP_PLAN_PIN", "-1" } }; + const s = try Config.parse(&Fake.get); + try std.testing.expectEqual(@as(u64, 200 * 1000), s.plan_spin_ns); + try std.testing.expect(s.plan_pin == null); + Fake.pairs = &.{.{ "TF_TP_PLAN_LINK", "udp" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{}; // the RoCE grid cap: 16 unset, powers of two only + try std.testing.expectEqual(@as(u32, 16), (try Config.parse(&Fake.get)).roce_blocks); + Fake.pairs = &.{.{ "TF_TP_ROCE_BLOCKS", "4" }}; + try std.testing.expectEqual(@as(u32, 4), (try Config.parse(&Fake.get)).roce_blocks); + Fake.pairs = &.{.{ "TF_TP_ROCE_BLOCKS", "12" }}; + try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); + Fake.pairs = &.{}; +} diff --git a/zig/src/fabric/tp/fate.zig b/zig/src/fabric/tp/fate.zig new file mode 100644 index 0000000000..5084c4648b --- /dev/null +++ b/zig/src/fabric/tp/fate.zig @@ -0,0 +1,261 @@ +//! Fail fast on every rank: a failure or a vanished rank ends all ranks with exit 70 about a second later. +const std = @import("std"); +const log = @import("log.zig"); +const sock = @import("sock.zig"); +const Config = @import("config.zig").Config; +const posix = std.posix; +/// Linux POLLRDHUP: the peer shut its side down (std has no name for it). +const pollrdhup: i16 = 0x2000; + +pub const exit_code: u8 = 70; +const magic: u32 = 0x46463431; +const kind_mem: u32 = 1; +const kind_fail: u32 = 2; +/// rank 0's clean stop: the hang-up that follows is no failure (sent before rank 0 closes its sockets) +const kind_stop: u32 = 3; +const Frame = extern struct { magic: u32, kind: u32, value: i64 }; + +pub const max_ranks = 64; +pub const max_hooks = 8; + +/// A callback with its context (an abort of a collective, rank 0's request failing); `reason` lives until exit. +pub const Hook = struct { ctx: *anyopaque, call: *const fn (ctx: *anyopaque, reason: []const u8) void }; + +pub const Fate = struct { + rank: u32, + world: u32, + fds: [max_ranks]sock.Fd, + grace_ns: u64, + report_ns: u64, + /// This rank's usable memory in bytes for the MEM report (default: MemAvailable). + usable: *const fn () u64 = memAvailable, + /// The process exit; tests replace it (it may return there). + exit_fn: *const fn (u8) void = exitNow, + hooks: [max_hooks]Hook = undefined, + n_hooks: std.atomic.Value(u32) = .init(0), + dead: std.atomic.Value(bool) = .init(false), + closing: std.atomic.Value(bool) = .init(false), + /// Set once the exit was asked for (tests read it; the process is gone in serving). + exited: std.atomic.Value(bool) = .init(false), + reason_buf: [512]u8 = undefined, + reason_len: usize = 0, + send_lock: std.atomic.Value(bool) = .init(false), + /// Rank 0: each rank's last report (bytes) and its monotonic time. + peer_mem: [max_ranks]std.atomic.Value(u64) = @splat(.init(0)), + peer_mem_at: [max_ranks]std.atomic.Value(u64) = @splat(.init(0)), + wake: [2]sock.Fd = .{ -1, -1 }, + thread: ?std.Thread = null, + + /// The channel over the bootstrap's sockets (`Bootstrap.release`); `start` runs its thread. + pub fn init(cfg: Config, fds: [max_ranks]sock.Fd) Fate { + return .{ .rank = cfg.rank, .world = cfg.world, .fds = fds, .grace_ns = cfg.grace_ns, .report_ns = cfg.report_ns }; + } + + pub fn start(f: *Fate) !void { + var p: [2]c_int = undefined; + if (std.c.pipe2(&p, .{ .CLOEXEC = true }) != 0) return error.SocketFailed; + f.wake = p; + f.thread = if (f.rank == 0) try std.Thread.spawn(.{}, watch0, .{f}) else try std.Thread.spawn(.{}, watchK, .{f}); + } + + /// Runs `hook` when the ranks go down; collectives register their abort here. + pub fn onFail(f: *Fate, hook: Hook) void { + const i = f.n_hooks.load(.acquire); + std.debug.assert(i < max_hooks); + f.hooks[i] = hook; + f.n_hooks.store(i + 1, .release); + } + + /// This rank failed (`where`: what it was doing): every rank goes down. Returns only in tests. + pub fn fatal(f: *Fate, where: []const u8, err: anyerror) void { + var buf: [256]u8 = undefined; + const msg = std.fmt.bufPrint(&buf, "rank {d} failed in {s}: {t}", .{ f.rank, where, err }) catch "rank failed"; + f.die(msg, true); + } + + /// This rank failed while booting: every rank goes down, and the thread is joined before the caller frees this. + pub fn failBoot(f: *Fate, where: []const u8, err: anyerror) void { + f.fatal(where, err); + f.close(); + } + + /// A normal stop is coming (rank 0 before it tells the others; the others when told): later hang-ups are fine. + pub fn expectStop(f: *Fate) void { + f.closing.store(true, .release); + } + + /// Rank 0's clean stop of every rank: a STOP frame to each ahead of its close, so followers end quietly. + pub fn announceStop(f: *Fate) void { + if (f.closing.swap(true, .acq_rel)) return; + if (f.rank != 0) return; + const h: Frame = .{ .magic = magic, .kind = kind_stop, .value = 0 }; + f.lock(); + defer f.unlock(); + for (f.fds[1..f.world]) |fd| if (fd >= 0) { + _ = sock.trySend(fd, std.mem.asBytes(&h)); + }; + } + + /// A normal stop: the thread ends and the sockets close. + pub fn close(f: *Fate) void { + f.closing.store(true, .release); + if (f.wake[1] >= 0) _ = std.c.write(f.wake[1], "x", 1); + if (f.thread) |t| t.join(); + f.thread = null; + for (&f.fds) |*fd| if (fd.* >= 0) { + sock.close(fd.*); + fd.* = -1; + }; + for (f.wake) |fd| if (fd >= 0) sock.close(fd); + f.wake = .{ -1, -1 }; + } + + /// Rank 0: rank `r`'s last reported usable bytes and the age of the report, or null before its first. + pub fn peerMem(f: *const Fate, r: u32) ?struct { bytes: u64, age_ns: u64 } { + const at = f.peer_mem_at[r].load(.acquire); + if (at == 0) return null; + return .{ .bytes = f.peer_mem[r].load(.acquire), .age_ns = sock.nowNs() -| at }; + } + + pub fn reason(f: *const Fate) []const u8 { + return f.reason_buf[0..f.reason_len]; + } + + fn lock(f: *Fate) void { + while (f.send_lock.swap(true, .acquire)) std.atomic.spinLoopHint(); + } + + fn unlock(f: *Fate) void { + f.send_lock.store(false, .release); + } + + fn die(f: *Fate, why: []const u8, own: bool) void { + if (f.dead.swap(true, .acq_rel)) { + // the first caller exits the process; a second (the fate thread or a round thread) waits for it + while (!f.exited.load(.acquire)) sock.sleepNs(10 * std.time.ns_per_ms); + return; + } + const n = @min(why.len, f.reason_buf.len); + @memcpy(f.reason_buf[0..n], why[0..n]); + f.reason_len = n; + const r = f.reason(); + log.warn("[tensorfold] tp fail-fast: rank {d}: {s}; exiting ({d}) in {d} ms so every rank restarts", .{ f.rank, r, exit_code, f.grace_ns / std.time.ns_per_ms }); + // aborts first: kernels waiting on the dead peer end now; a hook that blocks cannot hold the exit + if (std.Thread.spawn(.{}, runHooks, .{ f, r })) |t| t.detach() else |_| runHooks(f, r); + if (!f.closing.load(.acquire)) { + const head: Frame = .{ .magic = magic, .kind = kind_fail, .value = @intCast(r.len) }; + f.lock(); + for (f.fds[0..f.world]) |fd| if (fd >= 0) { + if (own or f.rank == 0) { + _ = sock.trySend(fd, std.mem.asBytes(&head)); + _ = sock.trySend(fd, r); + } + sock.hangup(fd); + }; + f.unlock(); + } + sock.sleepNs(f.grace_ns); + f.exited.store(true, .release); + f.exit_fn(exit_code); + } + + fn runHooks(f: *Fate, r: []const u8) void { + const n = f.n_hooks.load(.acquire); + for (f.hooks[0..n]) |h| h.call(h.ctx, r); + } + + /// Reads one frame from `fd`: a MEM report, a FAIL (the ranks go down), or the end (the peer is gone). + fn readFrame(f: *Fate, peer: u32, fd: sock.Fd) bool { + var h: Frame = undefined; + const deadline = sock.nowNs() + std.time.ns_per_s; + sock.recvExact(fd, std.mem.asBytes(&h), deadline) catch { + if (f.closing.load(.acquire)) return false; + var buf: [128]u8 = undefined; + f.die(std.fmt.bufPrint(&buf, "rank {d} closed the fate channel (it failed, exited or was killed)", .{peer}) catch "peer gone", false); + return false; + }; + if (h.magic != magic) { + f.die("a frame that is not the fate channel's", false); + return false; + } + if (h.kind == kind_stop) { + // rank 0 stops every rank: its hang-up is next, and no failure + f.closing.store(true, .release); + return false; + } + if (h.kind == kind_mem) { + f.peer_mem[peer].store(@intCast(@max(h.value, 0)), .release); + f.peer_mem_at[peer].store(sock.nowNs(), .release); + return true; + } + var text: [512]u8 = undefined; + const len: usize = @intCast(std.math.clamp(h.value, 0, text.len)); + sock.recvExact(fd, text[0..len], deadline) catch {}; + f.die(if (len > 0) text[0..len] else "a peer failed", false); + return false; + } + + fn watch0(f: *Fate) void { + var p: [max_ranks + 1]std.c.pollfd = undefined; + var who: [max_ranks + 1]u32 = undefined; + var n: usize = 0; + for (f.fds[1..f.world], 1..) |fd, r| if (fd >= 0) { + p[n] = .{ .fd = fd, .events = posix.POLL.IN | pollrdhup, .revents = 0 }; + who[n] = @intCast(r); + n += 1; + }; + p[n] = .{ .fd = f.wake[0], .events = posix.POLL.IN, .revents = 0 }; + while (!f.dead.load(.acquire) and !f.closing.load(.acquire)) { + if (std.c.poll(&p, @intCast(n + 1), -1) <= 0) continue; + for (p[0..n], who[0..n]) |*q, r| if (q.revents != 0) { + if (!f.readFrame(r, q.fd)) return; + }; + } + } + + fn watchK(f: *Fate) void { + const fd = f.fds[0]; + var p = [_]std.c.pollfd{ + .{ .fd = fd, .events = posix.POLL.IN | pollrdhup, .revents = 0 }, + .{ .fd = f.wake[0], .events = posix.POLL.IN, .revents = 0 }, + }; + const ms: c_int = @intCast(@max(f.report_ns / std.time.ns_per_ms, 1)); + f.report(fd); + while (!f.dead.load(.acquire) and !f.closing.load(.acquire)) { + const ready = std.c.poll(&p, 2, ms); + if (ready == 0) { + f.report(fd); + continue; + } + if (ready > 0 and p[0].revents != 0 and !f.readFrame(0, fd)) return; + } + } + + /// One MEM frame, never blocking (a full buffer drops it: the next comes one period later). + fn report(f: *Fate, fd: sock.Fd) void { + const h: Frame = .{ .magic = magic, .kind = kind_mem, .value = @intCast(f.usable()) }; + f.lock(); + defer f.unlock(); + _ = sock.trySend(fd, std.mem.asBytes(&h)); + } +}; + +fn exitNow(code: u8) void { + std.c._exit(code); +} + +/// MemAvailable from /proc/meminfo in bytes (0 when unreadable). +pub fn memAvailable() u64 { + const fd = std.c.open("/proc/meminfo", .{ .ACCMODE = .RDONLY }); + if (fd < 0) return 0; + defer _ = std.c.close(fd); + var buf: [4096]u8 = undefined; + const n = std.c.read(fd, &buf, buf.len); + if (n <= 0) return 0; + const text = buf[0..@intCast(n)]; + const key = "MemAvailable:"; + const at = std.mem.indexOf(u8, text, key) orelse return 0; + var it = std.mem.tokenizeAny(u8, text[at + key.len ..], " \n"); + const kb = std.fmt.parseInt(u64, it.next() orelse return 0, 10) catch return 0; + return kb * 1024; +} diff --git a/zig/src/fabric/tp/fate_test.zig b/zig/src/fabric/tp/fate_test.zig new file mode 100644 index 0000000000..6019055f50 --- /dev/null +++ b/zig/src/fabric/tp/fate_test.zig @@ -0,0 +1,146 @@ +//! The fate channel between two ranks over real TCP: a failure or a vanished rank ends both, a stop ends neither. +const std = @import("std"); +const sock = @import("sock.zig"); +const fate = @import("fate.zig"); +const bootstrap = @import("bootstrap.zig"); +const Config = @import("config.zig").Config; +const Fate = fate.Fate; + +var exits: std.atomic.Value(u32) = .init(0); +var hooks: std.atomic.Value(u32) = .init(0); + +fn recordExit(code: u8) void { + std.debug.assert(code == fate.exit_code); + _ = exits.fetchAdd(1, .acq_rel); +} + +fn hook(_: *anyopaque, reason: []const u8) void { + std.debug.assert(reason.len > 0); + _ = hooks.fetchAdd(1, .acq_rel); +} + +const grace = 50 * std.time.ns_per_ms; + +/// Two connected ranks' fate channels (rank 0 first), started, with recorded exits and hooks. +fn pair(fates: *[2]Fate) !void { + return pairAt(.{ &fates[0], &fates[1] }); +} + +fn pairAt(fates: [2]*Fate) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var fds0: [fate.max_ranks]sock.Fd = @splat(-1); + var fds1: [fate.max_ranks]sock.Fd = @splat(-1); + fds0[1] = b; + fds1[0] = a; + for (fates, [2][fate.max_ranks]sock.Fd{ fds0, fds1 }, 0..) |f, fds, r| { + f.* = Fate.init(.{ .world = 2, .rank = @intCast(r), .grace_ns = grace, .report_ns = 20 * std.time.ns_per_ms }, fds); + f.exit_fn = recordExit; + f.usable = fakeUsable; + f.onFail(.{ .ctx = f, .call = hook }); + } + for (fates) |f| try f.start(); +} + +fn fakeUsable() u64 { + return 123 << 20; +} + +fn waitExits(want: u32) !u64 { + const t0 = sock.nowNs(); + while (exits.load(.acquire) < want) { + if (sock.nowNs() - t0 > 5 * std.time.ns_per_s) return error.Timeout; + sock.sleepNs(std.time.ns_per_ms); + } + return sock.nowNs() - t0; +} + +fn reset() void { + exits.store(0, .release); + hooks.store(0, .release); +} + +test "rank 1's failure ends both ranks with its reason" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer for (&fates) |*f| f.close(); + sock.sleepNs(60 * std.time.ns_per_ms); + const mem = fates[0].peerMem(1) orelse return error.NoReport; + try std.testing.expectEqual(@as(u64, 123 << 20), mem.bytes); + fates[1].fatal("prefill", error.OutOfMemory); + const took = try waitExits(2); + try std.testing.expect(took < grace + 500 * std.time.ns_per_ms); + try std.testing.expectEqual(@as(u32, 2), hooks.load(.acquire)); + try std.testing.expect(std.mem.indexOf(u8, fates[0].reason(), "rank 1 failed in prefill: OutOfMemory") != null); +} + +test "rank 0's failure ends rank 1" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer for (&fates) |*f| f.close(); + fates[0].fatal("window", error.Unexpected); + _ = try waitExits(2); + try std.testing.expect(std.mem.indexOf(u8, fates[1].reason(), "rank 0 failed in window") != null); +} + +test "a rank that vanishes without a word ends the other" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer fates[0].close(); + // rank 1 killed: its socket closes with no FAIL frame (its own fate thread ends quietly first) + fates[1].expectStop(); + fates[1].close(); + _ = try waitExits(1); + try std.testing.expect(std.mem.indexOf(u8, fates[0].reason(), "rank 1 closed the fate channel") != null); +} + +test "a normal stop ends nobody" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + fates[0].expectStop(); + fates[1].expectStop(); + fates[1].close(); + fates[0].close(); + sock.sleepNs(2 * grace); + try std.testing.expectEqual(@as(u32, 0), exits.load(.acquire)); +} + +test "a rank failing at boot ends both with its reason, and its channel is down before its owner frees it" { + // a follower's boot failed after its session came up; its Fate is freed while the fate thread could still read it + reset(); + // rank 1's Fate where a Model keeps it: on the heap, its pages unmapped when freed (a stale read faults) + const pa = std.heap.page_allocator; + var f0: Fate = undefined; + const f1 = try pa.create(Fate); + try pairAt(.{ &f0, f1 }); + defer f0.close(); + sock.sleepNs(60 * std.time.ns_per_ms); + f1.failBoot("boot", error.RopeTooShort); + try std.testing.expect(f1.thread == null); + pa.destroy(f1); + _ = try waitExits(2); + // several report periods with rank 1's Fate gone: its thread would have reported from freed memory + sock.sleepNs(5 * 20 * std.time.ns_per_ms); + try std.testing.expect(std.mem.indexOf(u8, f0.reason(), "rank 1 failed in boot: RopeTooShort") != null); +} + +test "rank 0's clean stop ends every rank with no fail-fast exit, whatever order the ranks close in" { + // rank 0 ended its run and closed; the follower, still tearing down, must not see the hang-up as a failure + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + sock.sleepNs(60 * std.time.ns_per_ms); + fates[0].announceStop(); + fates[0].close(); // rank 0 gone before rank 1 even started its own stop + sock.sleepNs(2 * grace + 100 * std.time.ns_per_ms); + fates[1].close(); + try std.testing.expectEqual(@as(u32, 0), exits.load(.acquire)); + try std.testing.expectEqual(@as(u32, 0), hooks.load(.acquire)); +} diff --git a/zig/src/fabric/tp/gpu_bench.zig b/zig/src/fabric/tp/gpu_bench.zig new file mode 100644 index 0000000000..99831e6d77 --- /dev/null +++ b/zig/src/fabric/tp/gpu_bench.zig @@ -0,0 +1,82 @@ +//! Latency and bandwidth of each backend's collectives, eager and graphed; RESULT lines are JSON, busbw as nccl-tests. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const Ctx = @import("gpu_checks.zig").Ctx; +const Collective = tp.Collective; + +pub const Kind = enum { all_reduce, all_gather }; + +fn issue(c: *Ctx, comm: Collective, kind: Kind, bytes: usize) !void { + const count = bytes / 2; + switch (kind) { + .all_reduce => try comm.allReduce(c.in.ptr, c.out.ptr, count, .bf16, .sum, c.stream.handle), + .all_gather => try comm.allGather(c.in.ptr, c.out.ptr, count, .bf16, c.stream.handle), + } +} + +/// GPU microseconds a call, `iters` calls issued back to back on the stream. +fn eager(c: *Ctx, comm: Collective, kind: Kind, bytes: usize, iters: u32) !f64 { + var e0 = try cuda.Event.init(c.d, true); + defer e0.deinit(); + var e1 = try cuda.Event.init(c.d, true); + defer e1.deinit(); + for (0..10) |_| try issue(c, comm, kind, bytes); + try c.stream.synchronize(); + try comm.barrier(); + try e0.record(c.stream); + for (0..iters) |_| try issue(c, comm, kind, bytes); + try e1.record(c.stream); + try e1.synchronize(); + return @as(f64, try cuda.Event.elapsedMs(e0, e1)) * 1000.0 / @as(f64, @floatFromInt(iters)); +} + +/// GPU microseconds a call inside one graph of `per_graph` calls, replayed `replays` times. +fn graphed(c: *Ctx, comm: Collective, kind: Kind, bytes: usize, per_graph: u32, replays: u32) !f64 { + try cuda.graph.beginCapture(c.stream, .thread_local); + const err: ?tp.collective.Error = blk: { + for (0..per_graph) |_| issue(c, comm, kind, bytes) catch |e| break :blk e; + break :blk null; + }; + var g = try cuda.graph.endCapture(c.stream); + defer g.deinit(); + if (err) |e| return e; + var exec = try g.instantiate(); + defer exec.deinit(); + try exec.upload(c.stream); + try exec.launchOn(c.stream); + try c.stream.synchronize(); + try comm.barrier(); + var e0 = try cuda.Event.init(c.d, true); + defer e0.deinit(); + var e1 = try cuda.Event.init(c.d, true); + defer e1.deinit(); + try e0.record(c.stream); + for (0..replays) |_| try exec.launchOn(c.stream); + try e1.record(c.stream); + try e1.synchronize(); + try comm.check(); + return @as(f64, try cuda.Event.elapsedMs(e0, e1)) * 1000.0 / @as(f64, @floatFromInt(per_graph * replays)); +} + +/// Every power of two from `lo` to `hi` bytes a rank, both kinds, eager and graphed. +pub fn sweep(c: *Ctx, comm: Collective, label: []const u8, lo: usize, hi: usize) !void { + for ([_]Kind{ .all_reduce, .all_gather }) |kind| { + var bytes = lo; + while (bytes <= hi) : (bytes *= 2) { + const per_graph: u32 = if (bytes >= 1 << 20) 20 else 100; + const e = try eager(c, comm, kind, bytes, if (bytes >= 1 << 20) 50 else 200); + const g = try graphed(c, comm, kind, bytes, per_graph, 5); + const moved: f64 = @floatFromInt(if (kind == .all_gather) 2 * bytes else bytes); + const factor: f64 = if (kind == .all_gather) 0.5 else 1.0; + const busbw = moved / (g * 1e-6) * factor / 1e9; + if (c.rank == 0) std.debug.print("RESULT {{\"backend\":\"{s}\",\"op\":\"{t}\",\"bytes\":{d},\"eager_us\":{d:.2},\"graph_us\":{d:.2},\"busbw_GBps\":{d:.2}}}\n", .{ label, kind, bytes, e, g, busbw }); + } + } +} + +/// One decode window's exchanges: `n` all-gathers of one row's bf16 partials (`row_bytes` a rank) in one graph. +pub fn decodeWindow(c: *Ctx, comm: Collective, label: []const u8, n: u32, row_bytes: usize) !void { + const per = try graphed(c, comm, .all_gather, row_bytes, n, 20); + if (c.rank == 0) std.debug.print("RESULT {{\"backend\":\"{s}\",\"op\":\"decode_window\",\"exchanges\":{d},\"bytes\":{d},\"per_exchange_us\":{d:.2},\"window_us\":{d:.1}}}\n", .{ label, n, row_bytes, per, per * @as(f64, @floatFromInt(n)) }); +} diff --git a/zig/src/fabric/tp/gpu_checks.zig b/zig/src/fabric/tp/gpu_checks.zig new file mode 100644 index 0000000000..5dc4ed6d4c --- /dev/null +++ b/zig/src/fabric/tp/gpu_checks.zig @@ -0,0 +1,188 @@ +//! Every collective on real GPUs against the host reference, eager and replayed from a graph. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const Collective = tp.Collective; +const DType = tp.DType; +const Op = tp.Op; +const reference = tp.reference; + +pub const Ctx = struct { + d: *const cuda.Driver, + gpa: std.mem.Allocator, + stream: cuda.Stream, + rank: u32, + /// Device buffers sized for the largest case: input, output (world x), a spare. + in: cuda.DeviceBuffer, + out: cuda.DeviceBuffer, + /// Host scratch: this rank's input, the peer's, the expected and the downloaded output. + mine: []u8, + peer: []u8, + want: []u8, + got: []u8, + failures: u32 = 0, + passes: u32 = 0, + + pub fn init(d: *const cuda.Driver, gpa: std.mem.Allocator, rank: u32, max_bytes: usize) !Ctx { + return .{ + .d = d, + .gpa = gpa, + .stream = try cuda.Stream.init(d, true), + .rank = rank, + .in = try cuda.DeviceBuffer.alloc(d, 2 * max_bytes + 256), + .out = try cuda.DeviceBuffer.alloc(d, 2 * max_bytes + 256), + .mine = try gpa.alloc(u8, max_bytes), + .peer = try gpa.alloc(u8, max_bytes), + .want = try gpa.alloc(u8, 2 * max_bytes), + .got = try gpa.alloc(u8, 2 * max_bytes), + }; + } + + pub fn deinit(c: *Ctx) void { + c.stream.deinit(); + c.in.free(); + c.out.free(); + for ([_][]u8{ c.mine, c.peer, c.want, c.got }) |b| c.gpa.free(b); + } + + fn verdict(c: *Ctx, ok: bool, comptime fmt: []const u8, args: anytype) void { + if (ok) c.passes += 1 else c.failures += 1; + std.debug.print("{s} " ++ fmt ++ "\n", .{if (ok) "PASS" else "FAIL"} ++ args); + } + + /// Both ranks' inputs for case `seed`: this rank's into `mine`, the peer's into `peer` (n bytes each). + fn inputs(c: *Ctx, n: usize, t: DType, seed: u64) void { + reference.fill(c.mine[0..n], t, c.rank, seed); + reference.fill(c.peer[0..n], t, 1 - c.rank, seed); + } + + fn rankInput(c: *Ctx, r: u32, n: usize) []const u8 { + return if (r == c.rank) c.mine[0..n] else c.peer[0..n]; + } + + fn compare(c: *Ctx, comm: Collective, label: []const u8, what: []const u8, t: DType, n: usize, out_off: usize, want: []const u8) !void { + try c.stream.synchronize(); + comm.check() catch |e| { + c.verdict(false, "{s} {s} {t} {d} B: transport error {t}", .{ label, what, t, n, e }); + return; + }; + try c.out.download(out_off, c.got[0..want.len]); + const ok = std.mem.eql(u8, c.got[0..want.len], want); + var first: usize = 0; + if (!ok) first = std.mem.indexOfDiff(u8, c.got[0..want.len], want) orelse 0; + c.verdict(ok, "{s} {s} {t} {d} B a rank{s}", .{ label, what, t, n, if (ok) "" else " (first difference at a byte shown below)" }); + if (!ok) std.debug.print(" byte {d}: got {x} want {x}\n", .{ first, c.got[first], want[first] }); + } + + /// all-reduce (out of place, in place, misaligned by `skew` bytes), all-gather (out of place, in place), exchange. + pub fn collectives(c: *Ctx, comm: Collective, label: []const u8, n_bytes: usize, seed: u64) !void { + const reduce_types = [_]DType{ .f32, .bf16, .f16, .i32 }; + for (reduce_types, 0..) |t, ti| { + const n = n_bytes - n_bytes % t.size(); + const count = n / t.size(); + c.inputs(n, t, seed + ti); + reference.reduce(c.want[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, t, .sum); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.in.ptr, c.out.ptr, count, t, .sum, c.stream.handle); + try c.compare(comm, label, "all-reduce sum", t, n, 0, c.want[0..n]); + // in place, at an offset of one element (not 16-byte aligned for 2-byte types) + const skew = t.size(); + try c.out.uploadAsync(skew, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.out.ptr + skew, c.out.ptr + skew, count, t, .sum, c.stream.handle); + try c.compare(comm, label, "all-reduce sum in place, skewed", t, n, skew, c.want[0..n]); + } + for ([_]Op{ .max, .min }) |op| { + const n = n_bytes - n_bytes % 4; + c.inputs(n, .f32, seed + 11); + reference.reduce(c.want[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, .f32, op); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.in.ptr, c.out.ptr, n / 4, .f32, op, c.stream.handle); + try c.compare(comm, label, if (op == .max) "all-reduce max" else "all-reduce min", .f32, n, 0, c.want[0..n]); + } + for ([_]DType{ .u8, .bf16, .f32 }, 0..) |t, ti| { + const n = n_bytes - n_bytes % t.size(); + c.inputs(n, t, seed + 20 + ti); + @memcpy(c.want[0..n], c.rankInput(0, n)); + @memcpy(c.want[n..][0..n], c.rankInput(1, n)); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.allGather(c.in.ptr, c.out.ptr, n / t.size(), t, c.stream.handle); + try c.compare(comm, label, "all-gather", t, n, 0, c.want[0 .. 2 * n]); + try c.out.fill8(0, c.stream.handle); + try c.out.uploadAsync(c.rank * n, c.mine[0..n], c.stream.handle); + try comm.allGather(c.out.ptr + c.rank * n, c.out.ptr, n / t.size(), t, c.stream.handle); + try c.compare(comm, label, "all-gather in place", t, n, 0, c.want[0 .. 2 * n]); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + try comm.exchange(c.in.ptr, c.out.ptr, n / t.size(), t, 1 - c.rank, c.stream.handle); + try c.compare(comm, label, "exchange", t, n, 0, c.peer[0..n]); + } + } + + /// send then recv one way, then the other way (NCCL behind every backend). + pub fn pointToPoint(c: *Ctx, comm: Collective, label: []const u8, n_bytes: usize, seed: u64) !void { + const n = n_bytes - n_bytes % 4; + c.inputs(n, .f32, seed); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + for (0..2) |dir| { + const sender: u32 = @intCast(dir); + if (c.rank == sender) { + try comm.send(c.in.ptr, n / 4, .f32, 1 - c.rank, c.stream.handle); + try c.stream.synchronize(); + } else { + try comm.recv(c.out.ptr, n / 4, .f32, sender, c.stream.handle); + try c.compare(comm, label, "send/recv", .f32, n, 0, c.peer[0..n]); + } + } + } + + /// A captured chain replayed with fresh inputs, eager calls between replays (the mailbox epoch spans both). + pub fn graphs(c: *Ctx, comm: Collective, label: []const u8, n_bytes: usize, replays: u32) !void { + const n = n_bytes - n_bytes % 16; + const half = n / 2; // bf16 elements + try cuda.graph.beginCapture(c.stream, .thread_local); + // x1 = sum(x0) [bf16], x2 = gather(x1), x3 = sum(x2) [f32 view], x4 = exchange(x3) + const x0 = c.in.ptr; + const x1 = c.out.ptr; + const x2 = c.in.ptr + n; + const x3 = c.out.ptr + n; + const x4 = c.out.ptr + 3 * n; + const ops: ?tp.collective.Error = blk: { + comm.allReduce(x0, x1, half, .bf16, .sum, c.stream.handle) catch |e| break :blk e; + comm.allGather(x1, x2, half, .bf16, c.stream.handle) catch |e| break :blk e; + comm.allReduce(x2, x3, n / 2, .f32, .sum, c.stream.handle) catch |e| break :blk e; + comm.exchange(x3, x4, n / 2, .f32, 1 - c.rank, c.stream.handle) catch |e| break :blk e; + break :blk null; + }; + var g = try cuda.graph.endCapture(c.stream); + defer g.deinit(); + if (ops) |e| return e; + var exec = try g.instantiate(); + defer exec.deinit(); + const nodes = try g.nodeCount(); + const want1 = try c.gpa.alloc(u8, 2 * n); + defer c.gpa.free(want1); + const want3 = try c.gpa.alloc(u8, 2 * n); + defer c.gpa.free(want3); + for (0..replays) |k| { + c.inputs(n, .bf16, 1000 + k); + try c.in.uploadAsync(0, c.mine[0..n], c.stream.handle); + reference.reduce(want1[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, .bf16, .sum); + @memcpy(want1[n..][0..n], want1[0..n]); // the gather of equal sums: twice the same + reference.reduce(want3[0 .. 2 * n], &.{ want1[0 .. 2 * n], want1[0 .. 2 * n] }, .f32, .sum); + try exec.launchOn(c.stream); + try c.compare(comm, label, "graph replay: sum, gather, sum, exchange", .bf16, n, n, want3[0 .. 2 * n]); + try c.compare(comm, label, "graph replay: exchange of the sum", .f32, 2 * n, 3 * n, want3[0 .. 2 * n]); + // an eager call between replays + try c.collectivesQuick(comm, label, k); + } + std.debug.print(" ({s}: {d} graph nodes, {d} replays)\n", .{ label, nodes, replays }); + } + + fn collectivesQuick(c: *Ctx, comm: Collective, label: []const u8, k: usize) !void { + const n = 4096; + c.inputs(n, .bf16, 5000 + k); + reference.reduce(c.want[0..n], &.{ c.rankInput(0, n), c.rankInput(1, n) }, .bf16, .sum); + try c.in.uploadAsync(5 * n, c.mine[0..n], c.stream.handle); + try comm.allReduce(c.in.ptr + 5 * n, c.out.ptr + 5 * n, n / 2, .bf16, .sum, c.stream.handle); + try c.compare(comm, label, "eager between replays", .bf16, n, 5 * n, c.want[0..n]); + } +}; diff --git a/zig/src/fabric/tp/gpu_main.zig b/zig/src/fabric/tp/gpu_main.zig new file mode 100644 index 0000000000..c5668a7b06 --- /dev/null +++ b/zig/src/fabric/tp/gpu_main.zig @@ -0,0 +1,92 @@ +//! `tf-tp-test suite | failfast RANK fatal|kill | roce-host`: one rank of a two-process TP=2 test on real GPUs. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const tp_kernel = @import("tp_kernel"); +const checks = @import("gpu_checks.zig"); +const bench = @import("gpu_bench.zig"); + +const max_bytes = (4 << 20) + 4096; + +pub fn main(init: std.process.Init) !u8 { + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 2) { + std.debug.print("usage: tf-tp-test suite | failfast RANK fatal|kill | roce-host\n", .{}); + return 2; + } + const cfg = try tp.Config.fromEnv(); + if (std.mem.eql(u8, args[1], "roce-host")) return @import("roce_host.zig").run(cfg, init.gpa); + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, @intCast(cfg.device)); + defer ctx.deinit(); + var name_buf: [128]u8 = undefined; + const t0 = tp.sock.nowNs(); + var s: tp.Session = undefined; + try s.open(&driver, cfg, tp_kernel.mailbox); + const up_ms = @as(f64, @floatFromInt(tp.sock.nowNs() - t0)) / 1e6; + std.debug.print("INFO rank {d}/{d} on {s} (sm_{d}), backend {t}, mailbox kernel {d} B, bootstrap + NCCL init {d:.0} ms\n", .{ cfg.rank, cfg.world, try ctx.name(&name_buf), try ctx.capability(), s.comm().kind(), tp_kernel.mailbox.len, up_ms }); + if (std.mem.eql(u8, args[1], "suite")) return suite(&s, &driver, init.gpa, cfg); + if (std.mem.eql(u8, args[1], "failfast") and args.len == 4) return failfast(&s, &driver, init.gpa, cfg, try std.fmt.parseInt(u32, args[2], 10), args[3]); + std.debug.print("unknown command {s}\n", .{args[1]}); + return 2; +} + +fn suite(s: *tp.Session, d: *const cuda.Driver, gpa: std.mem.Allocator, cfg: tp.Config) !u8 { + var c = try checks.Ctx.init(d, gpa, cfg.rank, max_bytes); + defer c.deinit(); + var backends: [2]struct { comm: tp.Collective, label: []const u8 } = undefined; + var n_backends: usize = 1; + backends[0] = .{ .comm = s.nccl.iface(), .label = "nccl" }; + if (s.comm().kind() != .nccl) { + backends[1] = .{ .comm = s.comm(), .label = if (s.roce != null) "roce" else "mailbox" }; + n_backends = 2; + } + const sizes = [_]usize{ 6, 1000, 4096, 65536, 256 * 1024, 256 * 1024 + 16, 1 << 20, 4 << 20 }; + for (backends[0..n_backends]) |b| { + for (sizes, 0..) |n, i| try c.collectives(b.comm, b.label, n, 100 * i); + try c.pointToPoint(b.comm, b.label, 64 * 1024, 7); + for ([_]usize{ 4096, 65536, 1 << 20 }) |n| try c.graphs(b.comm, b.label, n, 4); + try b.comm.barrier(); + } + for (backends[0..n_backends]) |b| { + try bench.sweep(&c, b.comm, b.label, 4096, 4 << 20); + try bench.decodeWindow(&c, b.comm, b.label, 86, 8192); + } + // the stop: rank 0 marks it expected first, then both meet and close + if (s.fate) |*f| f.expectStop(); + try s.comm().barrier(); + std.debug.print("SUMMARY rank {d}: {d} passed, {d} failed\n", .{ cfg.rank, c.passes, c.failures }); + s.close(); + return if (c.failures == 0) 0 else 1; +} + +fn realNow() f64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.REALTIME, &ts); + return @as(f64, @floatFromInt(ts.sec)) + @as(f64, @floatFromInt(ts.nsec)) / 1e9; +} + +/// Exchanges until the victim fails; the survivor is left waiting inside one, which only the fate channel ends. +fn failfast(s: *tp.Session, d: *const cuda.Driver, gpa: std.mem.Allocator, cfg: tp.Config, victim: u32, mode: []const u8) !u8 { + var c = try checks.Ctx.init(d, gpa, cfg.rank, 1 << 20); + defer c.deinit(); + const comm = s.comm(); + var k: u32 = 0; + while (true) : (k += 1) { + if (k == 50 and cfg.rank == victim) { + std.debug.print("INJECT {s} rank {d} at {d:.6}\n", .{ mode, cfg.rank, realNow() }); + if (std.mem.eql(u8, mode, "kill")) _ = std.c.raise(.KILL); + s.fatal("window", error.Injected); + return 3; // not reached: the process exits 70 + } + try comm.allReduce(c.in.ptr, c.out.ptr, 2048, .bf16, .sum, c.stream.handle); + c.stream.synchronize() catch |e| { + std.debug.print("rank {d}: stream error {t} after the peer failed\n", .{ cfg.rank, e }); + }; + comm.check() catch |e| { + std.debug.print("rank {d}: transport reports {t} at exchange {d}; waiting for the fate channel\n", .{ cfg.rank, e, k }); + while (true) tp.sock.sleepNs(std.time.ns_per_s); + }; + } +} diff --git a/zig/src/fabric/tp/host.zig b/zig/src/fabric/tp/host.zig new file mode 100644 index 0000000000..051daea537 --- /dev/null +++ b/zig/src/fabric/tp/host.zig @@ -0,0 +1,293 @@ +//! In-process ranks over host memory: the `Collective` contract without a GPU (desyncs refused, abort wakes waiters). +const std = @import("std"); +const collective = @import("collective.zig"); +const reference = @import("reference.zig"); +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Op = collective.Op; +const Error = collective.Error; + +pub const max_ranks = 8; + +const Call = enum(u8) { none, all_reduce, all_gather, all_gather_v, exchange, barrier }; + +const Post = struct { call: Call = .none, send: DevicePtr = 0, recv: DevicePtr = 0, count: usize = 0, t: DType = .u8, op: Op = .sum, peer: u32 = 0, lens: DevicePtr = 0 }; + +/// One-way message slot from rank `src` to rank `dst` (send / recv). +const Slot = struct { buf: std.atomic.Value(DevicePtr) = .init(0), bytes: usize = 0, taken: std.atomic.Value(bool) = .init(false) }; + +pub const Group = struct { + world: u32, + /// A wait longer than this is `error.Timeout` (a peer that never calls). + timeout_ns: u64 = 10 * std.time.ns_per_s, + posts: [max_ranks]Post = @splat(.{}), + arrived: std.atomic.Value(u32) = .init(0), + generation: std.atomic.Value(u32) = .init(0), + aborted: std.atomic.Value(bool) = .init(false), + /// Ranks (a bit each) whose transport refuses device lengths: their `all_gather_v` moves whole strides (tests). + var_refuse: u32 = 0, + slots: [max_ranks][max_ranks]Slot = @splat(@splat(.{})), + + pub fn init(world: u32) Group { + std.debug.assert(world >= 1 and world <= max_ranks); + return .{ .world = world }; + } + + pub fn rank(g: *Group, r: u32) Rank { + std.debug.assert(r < g.world); + return .{ .group = g, .me = r }; + } + + /// Sense-reversing barrier over every rank; `error.Aborted` once any rank aborts. + fn meet(g: *Group) Error!void { + const gen = g.generation.load(.acquire); + if (g.arrived.fetchAdd(1, .acq_rel) + 1 == g.world) { + g.arrived.store(0, .release); + _ = g.generation.fetchAdd(1, .acq_rel); + return if (g.aborted.load(.acquire)) error.Aborted else {}; + } + const t0 = now(); + while (g.generation.load(.acquire) == gen) { + if (g.aborted.load(.acquire)) return error.Aborted; + if (now() - t0 > g.timeout_ns) return error.Timeout; + std.Thread.yield() catch {}; + } + if (g.aborted.load(.acquire)) return error.Aborted; + } +}; + +fn now() u64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); +} + +fn bytesAt(p: DevicePtr, n: usize) []u8 { + if (n == 0) return &.{}; + return @as([*]u8, @ptrFromInt(p))[0..n]; +} + +pub const Rank = struct { + group: *Group, + me: u32, + /// `varAgreed`'s answer, once this communicator agreed. + var_agreed: ?bool = null, + + pub fn iface(self: *Rank) collective.Collective { + return .{ .ptr = self, .vtable = &vtable }; + } + + const vtable: collective.Collective.VTable = .{ + .kind = .host, + .rank = rankOf, + .world = worldOf, + .all_reduce = allReduce, + .all_gather = allGather, + .exchange = exchange, + .send = send, + .recv = recv, + .barrier = barrier, + .check = check, + .abort = abort, + .all_gather_v = allGatherV, + .var_fits = varFits, + .var_agreed = varAgreed, + }; + + fn self_(ptr: *anyopaque) *Rank { + return @ptrCast(@alignCast(ptr)); + } + + fn rankOf(ptr: *anyopaque) u32 { + return self_(ptr).me; + } + + fn worldOf(ptr: *anyopaque) u32 { + return self_(ptr).group.world; + } + + /// Posts this rank's call, waits for all, and refuses calls that do not match rank 0's. + fn enter(self: *Rank, p: Post) Error!void { + const g = self.group; + if (g.aborted.load(.acquire)) return error.Aborted; + g.posts[self.me] = p; + try g.meet(); + const r0 = g.posts[0]; + if (r0.call != p.call or r0.count != p.count or r0.t != p.t or r0.op != p.op) { + abort(self); + return error.Invalid; + } + } + + fn allReduce(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, op: Op, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try self.enter(.{ .call = .all_reduce, .send = s, .recv = r, .count = count, .t = t, .op = op }); + const n = count * t.size(); + var ins: [max_ranks][]const u8 = undefined; + for (0..g.world) |i| ins[i] = bytesAt(g.posts[i].send, n); + // the output may alias this rank's input: reduce into scratch first, then meet before writing + var scratch: [1 << 16]u8 = undefined; + if (n > scratch.len) return error.Unsupported; + if (n > 0) reference.reduce(scratch[0..n], ins[0..g.world], t, op); + try g.meet(); + @memcpy(bytesAt(r, n), scratch[0..n]); + try g.meet(); + } + + fn allGather(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try self.enter(.{ .call = .all_gather, .send = s, .recv = r, .count = count, .t = t }); + const n = count * t.size(); + const out = bytesAt(r, n * g.world); + for (0..g.world) |i| @memcpy(out[i * n ..][0..n], bytesAt(g.posts[i].send, n)); + try g.meet(); + } + + fn lensAt(p: DevicePtr, world: u32) []const i32 { + return @as([*]const i32, @ptrFromInt(p))[0..world]; + } + + /// Copies only each rank's lens[r] bytes (`var_refuse` ranks take whole strides); unequal lengths are refused. + fn allGatherV(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, stride: usize, lens: DevicePtr, t: DType, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try self.enter(.{ .call = .all_gather_v, .send = s, .recv = r, .count = stride, .t = t, .lens = lens }); + const n = stride * t.size(); + const mine = lensAt(lens, g.world); + for (0..g.world) |i| if (!std.mem.eql(i32, mine, lensAt(g.posts[i].lens, g.world))) { + abort(self); + return error.Invalid; + }; + const whole = g.var_refuse & (@as(u32, 1) << @intCast(self.me)) != 0; + const out = bytesAt(r, n * g.world); + for (0..g.world) |i| { + if (mine[i] < 0 or mine[i] > n) { + abort(self); + return error.Invalid; + } + const len: usize = if (whole) n else @intCast(mine[i]); + @memcpy(out[i * n ..][0..len], bytesAt(g.posts[i].send, len)); + } + try g.meet(); + } + + fn varFits(ptr: *anyopaque, bytes: usize) bool { + const self = self_(ptr); + return bytes > 0 and self.group.var_refuse & (@as(u32, 1) << @intCast(self.me)) == 0; + } + + /// One int all-gather of `varFits(16)` the first time, then the kept answer. + fn varAgreed(ptr: *anyopaque) Error!bool { + const self = self_(ptr); + if (self.var_agreed) |v| return v; + const mine: i32 = @intFromBool(varFits(ptr, 16)); + var all: [max_ranks]i32 = undefined; + try allGather(ptr, @intFromPtr(&mine), @intFromPtr(&all), 1, .i32, null); + self.var_agreed = std.mem.min(i32, all[0..self.group.world]) == 1; + return self.var_agreed.?; + } + + fn exchange(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, peer: u32, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try collective.checkPeer(self.me, g.world, peer); + try self.enter(.{ .call = .exchange, .send = s, .recv = r, .count = count, .t = t, .peer = peer }); + if (g.posts[peer].peer != self.me) { + abort(self); + return error.Invalid; + } + const n = count * t.size(); + @memcpy(bytesAt(r, n), bytesAt(g.posts[peer].send, n)); + try g.meet(); + } + + fn send(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try collective.checkPeer(self.me, g.world, peer); + const slot = &g.slots[self.me][peer]; + slot.bytes = count * t.size(); + slot.taken.store(false, .release); + slot.buf.store(if (buf == 0) 1 else buf, .release); + try waitFor(g, &slot.taken, true); + slot.buf.store(0, .release); + } + + fn recv(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, _: collective.Stream) Error!void { + const self = self_(ptr); + const g = self.group; + try collective.checkPeer(self.me, g.world, peer); + const slot = &g.slots[peer][self.me]; + const t0 = now(); + var src = slot.buf.load(.acquire); + while (src == 0) : (src = slot.buf.load(.acquire)) { + if (g.aborted.load(.acquire)) return error.Aborted; + if (now() - t0 > g.timeout_ns) return error.Timeout; + std.Thread.yield() catch {}; + } + const n = count * t.size(); + if (slot.bytes != n) { + abort(self); + return error.Invalid; + } + @memcpy(bytesAt(buf, n), bytesAt(src, n)); + slot.taken.store(true, .release); + } + + fn waitFor(g: *Group, flag: *std.atomic.Value(bool), want: bool) Error!void { + const t0 = now(); + while (flag.load(.acquire) != want) { + if (g.aborted.load(.acquire)) return error.Aborted; + if (now() - t0 > g.timeout_ns) return error.Timeout; + std.Thread.yield() catch {}; + } + } + + fn barrier(ptr: *anyopaque) Error!void { + const self = self_(ptr); + try self.enter(.{ .call = .barrier }); + try self.group.meet(); + } + + fn check(ptr: *anyopaque) Error!void { + if (self_(ptr).group.aborted.load(.acquire)) return error.Aborted; + } + + fn abort(ptr: *anyopaque) void { + self_(ptr).group.aborted.store(true, .release); + } +}; + +/// Runs `body(rank_collective, rank, ctx)` on `world` threads at once and returns the first error any rank hit. +pub fn run(world: u32, ctx: anytype, comptime body: fn (collective.Collective, u32, @TypeOf(ctx)) anyerror!void) !void { + return runWith(world, .{}, ctx, body); +} + +pub const Options = struct { + /// `Group.var_refuse`: ranks whose transport takes no device lengths. + var_refuse: u32 = 0, +}; + +pub fn runWith(world: u32, opts: Options, ctx: anytype, comptime body: fn (collective.Collective, u32, @TypeOf(ctx)) anyerror!void) !void { + var group = Group.init(world); + group.var_refuse = opts.var_refuse; + var ranks: [max_ranks]Rank = undefined; + var errs: [max_ranks]?anyerror = @splat(null); + var threads: [max_ranks]std.Thread = undefined; + const Wrap = struct { + fn go(r: *Rank, c: @TypeOf(ctx), e: *?anyerror) void { + body(r.iface(), r.me, c) catch |err| { + e.* = err; + r.group.aborted.store(true, .release); + }; + } + }; + for (0..world) |i| { + ranks[i] = group.rank(@intCast(i)); + threads[i] = try std.Thread.spawn(.{}, Wrap.go, .{ &ranks[i], ctx, &errs[i] }); + } + for (threads[0..world]) |t| t.join(); + for (errs[0..world]) |e| if (e) |err| return err; +} diff --git a/zig/src/fabric/tp/host_test.zig b/zig/src/fabric/tp/host_test.zig new file mode 100644 index 0000000000..7de32008af --- /dev/null +++ b/zig/src/fabric/tp/host_test.zig @@ -0,0 +1,154 @@ +//! The `Collective` contract on the in-process ranks: every call's bits against the reference, desyncs and aborts. +const std = @import("std"); +const collective = @import("collective.zig"); +const reference = @import("reference.zig"); +const host = @import("host.zig"); +const Collective = collective.Collective; +const DType = collective.DType; + +fn addr(b: []u8) collective.DevicePtr { + return @intFromPtr(b.ptr); +} + +const n = 96; + +const Calls = struct { + fn go(c: Collective, rank: u32, _: void) !void { + const world = c.world(); + inline for (.{ DType.f32, DType.bf16, DType.i32, DType.f16 }) |t| { + const bytes = n * t.size(); + var mine: [n * 8]u8 = undefined; + var out: [n * 8 * host.max_ranks]u8 = undefined; + reference.fill(mine[0..bytes], t, rank, 7); + // all-reduce against the reference sum of every rank's input, in rank order + var ins: [host.max_ranks][n * 8]u8 = undefined; + var views: [host.max_ranks][]const u8 = undefined; + for (0..world) |r| { + reference.fill(ins[r][0..bytes], t, @intCast(r), 7); + views[r] = ins[r][0..bytes]; + } + var want: [n * 8]u8 = undefined; + reference.reduce(want[0..bytes], views[0..world], t, .sum); + try c.allReduce(addr(&mine), addr(&out), n, t, .sum, null); + try std.testing.expectEqualSlices(u8, want[0..bytes], out[0..bytes]); + // in place + try c.allReduce(addr(&mine), addr(&mine), n, t, .sum, null); + try std.testing.expectEqualSlices(u8, want[0..bytes], mine[0..bytes]); + // all-gather: rank order + reference.fill(mine[0..bytes], t, rank, 7); + try c.allGather(addr(&mine), addr(&out), n, t, null); + for (0..world) |r| try std.testing.expectEqualSlices(u8, views[r], out[r * bytes ..][0..bytes]); + } + try c.barrier(); + } +}; + +test "every call on two and three in-process ranks equals the reference" { + try host.run(2, {}, Calls.go); + try host.run(3, {}, Calls.go); +} + +const Pairwise = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var mine: [64]u8 = @splat(@intCast(rank + 1)); + var got: [64]u8 = undefined; + const peer = 1 - rank; + try c.exchange(addr(&mine), addr(&got), 64, .u8, peer, null); + try std.testing.expectEqualSlices(u8, &@as([64]u8, @splat(@intCast(peer + 1))), &got); + // send / recv, ordered so neither waits forever + if (rank == 0) { + try c.send(addr(&mine), 16, .f32, 1, null); + } else { + var buf: [64]u8 = undefined; + try c.recv(addr(&buf), 16, .f32, 0, null); + try std.testing.expectEqualSlices(u8, &@as([64]u8, @splat(1)), &buf); + } + try std.testing.expectError(error.Invalid, c.exchange(addr(&mine), addr(&got), 64, .u8, rank, null)); + } +}; + +test "exchange, send and recv between two ranks" { + try host.run(2, {}, Pairwise.go); +} + +const Desync = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var a: [64]u8 = @splat(0); + var b: [128]u8 = undefined; + // the ranks disagree on the count: both see it instead of hanging + const r = c.allGather(addr(&a), addr(&b), if (rank == 0) 16 else 32, .u8, null); + if (r) |_| return error.ShouldHaveFailed else |e| if (e != error.Invalid and e != error.Aborted) return e; + } +}; + +test "a desynchronised call is refused on every rank" { + try host.run(2, {}, Desync.go); +} + +const Aborts = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var a: [16]u8 = @splat(0); + var b: [32]u8 = undefined; + if (rank == 1) { + // rank 1 never calls: it aborts after a while, which must free rank 0's waiting call + @import("sock.zig").sleepNs(50 * std.time.ns_per_ms); + c.abort(); + return; + } + try std.testing.expectError(error.Aborted, c.allGather(addr(&a), addr(&b), 16, .u8, null)); + try std.testing.expectError(error.Aborted, c.check()); + try std.testing.expectError(error.Aborted, c.barrier()); + } +}; + +test "abort frees a rank waiting in a collective" { + try host.run(2, {}, Aborts.go); +} + +const VarGather = struct { + fn go(c: Collective, rank: u32, refuse: u32) !void { + const world = c.world(); + const stride = 200; + const whole = refuse & (@as(u32, 1) << @intCast(rank)) != 0; + try std.testing.expectEqual(!whole, c.varFits(stride)); + try std.testing.expectEqual(refuse == 0, try c.varAgreed()); + try std.testing.expectEqual(refuse == 0, try c.varAgreed()); // kept: no second collective + var ins: [host.max_ranks][stride]u8 = undefined; + for (0..world) |r| reference.fill(&ins[r], .u8, @intCast(r), 11); + // lengths of 0, odd, 16 and the whole stride, in bytes; .f32 strides count elements + const cases = [_][3]i32{ .{ 0, 17, 200 }, .{ 200, 0, 3 }, .{ 16, 16, 16 }, .{ 199, 200, 1 } }; + inline for (.{ DType.u8, DType.f32 }) |t| for (cases) |lens| { + var out: [stride * host.max_ranks]u8 = @splat(0xEE); + try c.allGatherV(addr(&ins[rank]), addr(&out), stride / t.size(), @intFromPtr(&lens), t, null); + for (0..world) |r| { + const len: usize = @intCast(lens[r]); + try std.testing.expectEqualSlices(u8, ins[r][0..len], out[r * stride ..][0..len]); + // past a length: the transport's choice (untouched here, or the whole stride where lengths are refused) + const past = out[r * stride + len .. (r + 1) * stride]; + if (whole) try std.testing.expectEqualSlices(u8, ins[r][len..], past) else for (past) |x| try std.testing.expectEqual(@as(u8, 0xEE), x); + } + }; + try c.barrier(); + } +}; + +test "a variable-length all-gather moves each rank's length, and the ranks agree on device lengths once" { + for ([_]u32{ 2, 3 }) |w| { + try host.runWith(w, .{}, @as(u32, 0), VarGather.go); + try host.runWith(w, .{ .var_refuse = 2 }, @as(u32, 2), VarGather.go); + } +} + +const VarDesync = struct { + fn go(c: Collective, rank: u32, _: void) !void { + var a: [64]u8 = @splat(0); + var b: [128]u8 = undefined; + const lens = [2]i32{ 8, if (rank == 0) 8 else 16 }; + const r = c.allGatherV(addr(&a), addr(&b), 64, @intFromPtr(&lens), .u8, null); + if (r) |_| return error.ShouldHaveFailed else |e| if (e != error.Invalid and e != error.Aborted) return e; + } +}; + +test "lengths that differ between ranks are refused on every rank" { + try host.run(2, {}, VarDesync.go); +} diff --git a/zig/src/fabric/tp/hybrid.zig b/zig/src/fabric/tp/hybrid.zig new file mode 100644 index 0000000000..2cdeb0e4a4 --- /dev/null +++ b/zig/src/fabric/tp/hybrid.zig @@ -0,0 +1,167 @@ +//! Small exchanges through the mailbox (one host) or RoCE (two hosts), the rest through NCCL; ranks always agree. +const std = @import("std"); +const collective = @import("collective.zig"); +const Mailbox = @import("mailbox.zig").Mailbox; +const Roce = @import("roce.zig").Roce; +const Nccl = @import("nccl.zig").Nccl; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Op = collective.Op; +const Error = collective.Error; + +/// The one-kernel path for small shards: both expose the same calls. +pub const Small = union(enum) { + mailbox: *Mailbox, + roce: *Roce, + + fn fits(s: Small, n: usize) bool { + return switch (s) { + inline else => |p| p.fits(n), + }; + } + + fn sum(s: Small, src: DevicePtr, dst: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + return switch (s) { + inline else => |p| p.sum(src, dst, count, t, stream), + }; + } + + fn gather(s: Small, src: DevicePtr, dst: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return switch (s) { + inline else => |p| p.allGatherBytes(src, dst, n, stream), + }; + } + + fn exchange(s: Small, src: DevicePtr, dst: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return switch (s) { + inline else => |p| p.exchangeBytes(src, dst, n, stream), + }; + } + + fn check(s: Small) Error!void { + return switch (s) { + inline else => |p| p.check(), + }; + } + + fn abort(s: Small) void { + switch (s) { + inline else => |p| p.abort(), + } + } +}; + +pub const Hybrid = struct { + small: Small, + nccl: *Nccl, + /// Largest shard a rank sends through the small path. + max_bytes: usize, + /// Every rank's `varFits(16)`, agreed at `Session.open` (the bootstrap's all-gather): `varAgreed`'s answer. + var_agreed: bool = false, + + pub fn iface(h: *Hybrid) collective.Collective { + return .{ .ptr = h, .vtable = if (h.small == .roce) &roce_vtable else &mailbox_vtable }; + } + + fn vtable(comptime kind: collective.Kind) collective.Collective.VTable { + return .{ + .kind = kind, + .rank = rankOf, + .world = worldOf, + .all_reduce = allReduce, + .all_gather = allGather, + .exchange = exchange, + .send = send, + .recv = recv, + .barrier = barrier, + .check = check, + .abort = abort, + .all_gather_v = allGatherV, + .var_fits = varFits, + .var_agreed = varAgreed, + }; + } + + const mailbox_vtable = vtable(.hybrid); + const roce_vtable = vtable(.roce); + + fn self_(ptr: *anyopaque) *Hybrid { + return @ptrCast(@alignCast(ptr)); + } + + fn routed(h: *const Hybrid, n: usize) bool { + return n <= h.max_bytes and h.small.fits(n); + } + + fn rankOf(ptr: *anyopaque) u32 { + return self_(ptr).nccl.me; + } + + fn worldOf(ptr: *anyopaque) u32 { + return self_(ptr).nccl.size; + } + + fn allReduce(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, op: Op, stream: collective.Stream) Error!void { + const h = self_(ptr); + if (op == .sum and (t == .f32 or t == .bf16) and h.routed(count * t.size())) return h.small.sum(s, r, count, t, stream); + return Nccl.allReduce(h.nccl, s, r, count, t, op, stream); + } + + fn allGather(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const h = self_(ptr); + if (h.routed(count * t.size())) return h.small.gather(s, r, count * t.size(), stream); + return Nccl.allGather(h.nccl, s, r, count, t, stream); + } + + /// RoCE moves each rank's device length; the mailbox and NCCL move whole strides (the same bytes where lens says). + fn allGatherV(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, stride: usize, lens: DevicePtr, t: DType, stream: collective.Stream) Error!void { + const h = self_(ptr); + const n = stride * t.size(); + if (varFits(ptr, n)) return h.small.roce.allGatherV(s, r, n, lens, stream); + return allGather(ptr, s, r, stride, t, stream); + } + + /// RoCE with its device side, for a stride the small path takes (the route is the size's alone: every rank alike). + pub fn varFits(ptr: *anyopaque, n: usize) bool { + const h = self_(ptr); + return switch (h.small) { + .roce => |p| n > 0 and h.routed(n) and p.d != null, + .mailbox => false, + }; + } + + fn varAgreed(ptr: *anyopaque) Error!bool { + return self_(ptr).var_agreed; + } + + fn exchange(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const h = self_(ptr); + try collective.checkPeer(h.nccl.me, h.nccl.size, peer); + if (h.routed(count * t.size())) return h.small.exchange(s, r, count * t.size(), stream); + return Nccl.exchange(h.nccl, s, r, count, t, peer, stream); + } + + fn send(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + return Nccl.iface(self_(ptr).nccl).send(buf, count, t, peer, stream); + } + + fn recv(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + return Nccl.iface(self_(ptr).nccl).recv(buf, count, t, peer, stream); + } + + fn barrier(ptr: *anyopaque) Error!void { + return Nccl.iface(self_(ptr).nccl).barrier(); + } + + fn check(ptr: *anyopaque) Error!void { + const h = self_(ptr); + try h.small.check(); + try Nccl.iface(h.nccl).check(); + } + + fn abort(ptr: *anyopaque) void { + const h = self_(ptr); + h.small.abort(); + Nccl.iface(h.nccl).abort(); + } +}; diff --git a/zig/src/fabric/tp/log.zig b/zig/src/fabric/tp/log.zig new file mode 100644 index 0000000000..4e3234ab49 --- /dev/null +++ b/zig/src/fabric/tp/log.zig @@ -0,0 +1,15 @@ +//! TP's log lines, silent under test: a test's stderr is the build runner's channel. +const std = @import("std"); +const quiet = @import("builtin").is_test; + +pub fn err(comptime format: []const u8, args: anytype) void { + if (!quiet) std.log.err(format, args); +} + +pub fn warn(comptime format: []const u8, args: anytype) void { + if (!quiet) std.log.warn(format, args); +} + +pub fn info(comptime format: []const u8, args: anytype) void { + if (!quiet) std.log.info(format, args); +} diff --git a/zig/src/fabric/tp/mailbox.cu b/zig/src/fabric/tp/mailbox.cu new file mode 100644 index 0000000000..63d53b2484 --- /dev/null +++ b/zig/src/fabric/tp/mailbox.cu @@ -0,0 +1,556 @@ +// One-shot two-rank exchange through pinned host memory both GPUs map (the TP mailbox; mailbox.zig is its host side). + +#include +#include + +// Protocol adapted from b12x's RoCEnante (Apache-2.0): a device epoch, slots by seq parity, a seq flag after the data. +namespace tftp { + +constexpr int kMaxBlocks = 16; +constexpr int kFlagStride = 32; // words: one 128-byte line per flag +// ctrl words: [0] abort (any rank, any host thread); per rank r at 16 + 8 r: failed, seq, seen, block, why, waited us +constexpr int kCtrlAbort = 0; +constexpr int kCtrlRank = 16; + +enum Mode : uint32_t { kGather = 0, kSumF32 = 1, kSumBf16 = 2, kExchange = 3 }; + +struct MailArgs { + const uint8_t *in; // this rank's shard (device) + uint8_t *out; // gather [2 n], exchange / sums [n] + uint8_t *peer_slots; // the peer's inbox data, as this GPU maps it + uint32_t *peer_flags; // the peer's inbox flags + const uint8_t *my_slots; // this rank's inbox data + const uint32_t *my_flags; // this rank's inbox flags + uint32_t *state; // device: [0] epoch, [1] tail counter, [2] poison + uint32_t *ctrl; // shared control words + uint64_t nbytes; // bytes a rank + uint64_t slot_bytes; + uint64_t timeout_ns; + uint32_t rank; + uint32_t mode; +}; + +__device__ __forceinline__ uint32_t ld_relaxed_gpu(const uint32_t *p) { + uint32_t v; + asm volatile("ld.relaxed.gpu.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint32_t ld_relaxed_sys(const uint32_t *p) { + uint32_t v; + asm volatile("ld.relaxed.sys.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint32_t ld_acquire_sys(const uint32_t *p) { + uint32_t v; + asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint4 ld_relaxed_sys_v4(const void *p) { + uint4 v; + asm volatile("ld.relaxed.sys.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint32_t ld_relaxed_sys_u8(const uint8_t *p) { + uint32_t v; + asm volatile("ld.relaxed.sys.global.u8 %0, [%1];" : "=r"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ void st_relaxed_sys(uint32_t *p, uint32_t v) { + asm volatile("st.relaxed.sys.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ void st_release_sys(uint32_t *p, uint32_t v) { + asm volatile("st.release.sys.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ void st_release_gpu(uint32_t *p, uint32_t v) { + asm volatile("st.release.gpu.global.u32 [%0], %1;" ::"l"(p), "r"(v) : "memory"); +} + +__device__ __forceinline__ uint32_t atom_add_acq_rel_gpu(uint32_t *p, uint32_t v) { + uint32_t old; + asm volatile("atom.acq_rel.gpu.global.add.u32 %0, [%1], %2;" : "=r"(old) : "l"(p), "r"(v) : "memory"); + return old; +} + +__device__ __forceinline__ uint32_t atom_add_acq_rel_sys(uint32_t *p, uint32_t v) { + uint32_t old; + asm volatile("atom.acq_rel.sys.global.add.u32 %0, [%1], %2;" : "=r"(old) : "l"(p), "r"(v) : "memory"); + return old; +} + +__device__ __forceinline__ uint32_t atom_cas_gpu(uint32_t *p, uint32_t cmp, uint32_t v) { + uint32_t old; + asm volatile("atom.relaxed.gpu.global.cas.b32 %0, [%1], %2, %3;" : "=r"(old) : "l"(p), "r"(cmp), "r"(v) : "memory"); + return old; +} + +__device__ __forceinline__ void fence_sc_sys() { asm volatile("fence.sc.sys;" ::: "memory"); } + +__device__ __forceinline__ uint64_t globaltimer() { + uint64_t t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + return t; +} + +// dst[i] <- src[i] for i in [0, n), grid-stride, 16-byte vectors when aligned; SysSrc: system-scope loads. +template +__device__ __forceinline__ void copy_bytes(uint8_t *dst, const uint8_t *src, uint64_t n, uint64_t index, + uint64_t stride) { + const uint64_t a = reinterpret_cast(dst) | reinterpret_cast(src) | n; + if ((a & 15u) == 0u) { + const uint64_t packs = n >> 4; + for (uint64_t i = index; i < packs; i += stride) { + const uint4 v = SysSrc ? ld_relaxed_sys_v4(src + (i << 4)) : reinterpret_cast(src)[i]; + reinterpret_cast(dst)[i] = v; + } + } else { + for (uint64_t i = index; i < n; i += stride) { + dst[i] = static_cast(SysSrc ? ld_relaxed_sys_u8(src + i) : src[i]); + } + } +} + +// TF_TP_ROCE_FAST: copy_bytes with two grid strides' loads issued before their stores (same bytes, one round trip). +template +__device__ __forceinline__ void copy_bytes2(uint8_t *dst, const uint8_t *src, uint64_t n, uint64_t index, + uint64_t stride) { + const uint64_t a = reinterpret_cast(dst) | reinterpret_cast(src) | n; + if ((a & 15u) != 0u) { + copy_bytes(dst, src, n, index, stride); + return; + } + const uint64_t packs = n >> 4; + for (uint64_t i = index; i < packs; i += 2u * stride) { + const uint64_t j = i + stride; + uint4 v0, v1 = make_uint4(0u, 0u, 0u, 0u); + if (SysSrc) { + v0 = ld_relaxed_sys_v4(src + (i << 4)); + if (j < packs) v1 = ld_relaxed_sys_v4(src + (j << 4)); + } else { + v0 = reinterpret_cast(src)[i]; + if (j < packs) v1 = reinterpret_cast(src)[j]; + } + reinterpret_cast(dst)[i] = v0; + if (j < packs) reinterpret_cast(dst)[j] = v1; + } +} + +template +__device__ __forceinline__ void copy_any(bool fast, uint8_t *dst, const uint8_t *src, uint64_t n, uint64_t index, + uint64_t stride) { + if (fast) { + copy_bytes2(dst, src, n, index, stride); + } else { + copy_bytes(dst, src, n, index, stride); + } +} + +__device__ __forceinline__ uint32_t add_bf16x2(uint32_t x, uint32_t y) { + const float2 a = __bfloat1622float2(*reinterpret_cast(&x)); + const float2 b = __bfloat1622float2(*reinterpret_cast(&y)); + const __nv_bfloat162 r = __floats2bfloat162_rn(a.x + b.x, a.y + b.y); + return *reinterpret_cast(&r); +} + +__device__ __forceinline__ uint32_t add_f32(uint32_t x, uint32_t y) { + return __float_as_uint(__uint_as_float(x) + __uint_as_float(y)); +} + +__device__ __forceinline__ uint32_t ld_relaxed_sys_u16(const uint8_t *p) { + uint16_t v; + asm volatile("ld.relaxed.sys.global.u16 %0, [%1];" : "=h"(v) : "l"(p) : "memory"); + return v; +} + +__device__ __forceinline__ uint16_t add_bf16(uint16_t x, uint16_t y) { + const __nv_bfloat16 r = __float2bfloat16_rn(__bfloat162float(__ushort_as_bfloat16(x)) + + __bfloat162float(__ushort_as_bfloat16(y))); + return __bfloat16_as_ushort(r); +} + +// out <- mine + peer's, rank 0's first; IEEE addition commutes, so both ranks get the same bits. +template +__device__ __forceinline__ void sum_bytes(uint8_t *out, const uint8_t *mine, const uint8_t *peer, uint64_t n, + uint32_t rank, uint64_t index, uint64_t stride) { + const uint64_t al = reinterpret_cast(out) | reinterpret_cast(mine) | + reinterpret_cast(peer) | n; + if ((al & 15u) == 0u) { + const uint64_t packs = n >> 4; + for (uint64_t i = index; i < packs; i += stride) { + const uint4 m = reinterpret_cast(mine)[i]; + const uint4 p = ld_relaxed_sys_v4(peer + (i << 4)); + const uint4 a = rank == 0 ? m : p; + const uint4 b = rank == 0 ? p : m; + uint4 r; + r.x = Bf16 ? add_bf16x2(a.x, b.x) : add_f32(a.x, b.x); + r.y = Bf16 ? add_bf16x2(a.y, b.y) : add_f32(a.y, b.y); + r.z = Bf16 ? add_bf16x2(a.z, b.z) : add_f32(a.z, b.z); + r.w = Bf16 ? add_bf16x2(a.w, b.w) : add_f32(a.w, b.w); + reinterpret_cast(out)[i] = r; + } + } else if (Bf16) { + for (uint64_t i = index; i < (n >> 1); i += stride) { + const uint16_t m = reinterpret_cast(mine)[i]; + const uint16_t p = static_cast(ld_relaxed_sys_u16(peer + (i << 1))); + reinterpret_cast(out)[i] = rank == 0 ? add_bf16(m, p) : add_bf16(p, m); + } + } else { + for (uint64_t i = index; i < (n >> 2); i += stride) { + const uint32_t m = reinterpret_cast(mine)[i]; + const uint32_t p = ld_relaxed_sys(reinterpret_cast(peer) + i); + reinterpret_cast(out)[i] = rank == 0 ? add_f32(m, p) : add_f32(p, m); + } + } +} + +// Thread 0 of a block: spin until the peer block's flag holds seq; false on abort (why 1) or timeout (why 2). +__device__ __noinline__ bool wait_flag(const uint32_t *flag, uint32_t seq, const uint32_t *ctrl, uint64_t timeout_ns, + uint32_t *seen, uint32_t *why, uint32_t *waited_us) { + const uint64_t t0 = globaltimer(); + uint32_t polls = 0; + while (true) { + const uint32_t v = ld_acquire_sys(flag); + if (v == seq) { + *seen = v; + return true; + } + if ((++polls & 63u) == 0u) { + const uint64_t now = globaltimer(); + const bool aborted = ld_relaxed_sys(ctrl + kCtrlAbort) != 0u; + const bool late = now > t0 && now - t0 > timeout_ns; + if (aborted || late) { + *seen = ld_acquire_sys(flag); + *why = aborted ? 1u : 2u; + *waited_us = now > t0 ? static_cast((now - t0) / 1000u) : 0u; + return *seen == seq; + } + } + } +} + +// TF_TP_ROCE_FAST: wait_flag for one of several pollers of a flag; `done` (shared) ends the others once one saw it. +__device__ __noinline__ bool wait_flag_shared(const uint32_t *flag, uint32_t seq, const uint32_t *ctrl, + uint64_t timeout_ns, uint32_t *seen, uint32_t *why, uint32_t *waited_us, + volatile uint32_t *done) { + const uint64_t t0 = globaltimer(); + uint32_t polls = 0; + while (true) { + const uint32_t v = ld_acquire_sys(flag); + if (v == seq) { + *done = 1u; + *seen = v; + return true; + } + if (*done != 0u) { + *seen = seq; + return true; + } + if ((++polls & 63u) == 0u) { + const uint64_t now = globaltimer(); + const bool aborted = ld_relaxed_sys(ctrl + kCtrlAbort) != 0u; + const bool late = now > t0 && now - t0 > timeout_ns; + if (aborted || late) { + *seen = ld_acquire_sys(flag); + *why = aborted ? 1u : 2u; + *waited_us = now > t0 ? static_cast((now - t0) / 1000u) : 0u; + if (*seen == seq) { + *done = 1u; + return true; + } + return *done != 0u; + } + } + } +} + +} // namespace tftp + +using namespace tftp; + +extern "C" __global__ void __launch_bounds__(512) tp_mailbox(MailArgs a) { + __shared__ uint32_t s_head[2]; // poison, seq + __shared__ uint32_t s_bad; + const uint32_t tid = threadIdx.x; + const uint32_t gdim = gridDim.x; + const uint64_t index = static_cast(blockIdx.x) * blockDim.x + tid; + const uint64_t stride = static_cast(gdim) * blockDim.x; + + // a poisoned transport does nothing (the host raises on its next check); poison and epoch in one round trip + if (tid == 0) { + s_head[0] = ld_relaxed_gpu(a.state + 2); + s_head[1] = ld_relaxed_gpu(a.state) + 1u; + s_bad = 0u; + } + __syncthreads(); + if (s_head[0] != 0u) { + return; + } + const uint32_t seq = s_head[1]; + const uint64_t slot = seq & 1u; + const uint64_t flag_at = (slot * kMaxBlocks + blockIdx.x) * kFlagStride; + + // 1. this block's share of the shard into the peer's inbox, then its flag after it (system-scope release) + copy_bytes(a.peer_slots + slot * a.slot_bytes, a.in, a.nbytes, index, stride); + __syncthreads(); + if (tid == 0) { + fence_sc_sys(); + st_release_sys(a.peer_flags + flag_at, seq); + } + // the local shard while the peer's writes travel + if (a.mode == kGather) { + copy_bytes(a.out + static_cast(a.rank) * a.nbytes, a.in, a.nbytes, index, stride); + } + + // 2. the peer's block of the same index (it wrote exactly the bytes this block reads) + if (tid == 0) { + uint32_t seen = 0u, why = 0u, waited = 0u; + if (!wait_flag(a.my_flags + flag_at, seq, a.ctrl, a.timeout_ns, &seen, &why, &waited)) { + s_bad = 1u; + if (atom_cas_gpu(a.state + 2, 0u, 1u) == 0u) { // one failure record a rank + uint32_t *rec = a.ctrl + kCtrlRank + 8u * a.rank; + st_relaxed_sys(rec + 1, seq); + st_relaxed_sys(rec + 2, seen); + st_relaxed_sys(rec + 3, blockIdx.x); + st_relaxed_sys(rec + 4, why); + st_relaxed_sys(rec + 5, waited); + fence_sc_sys(); + st_relaxed_sys(rec, 1u); + fence_sc_sys(); + } + } + } + __syncthreads(); + + // 3. copy out, gated by this block's own wait + if (s_bad == 0u) { + const uint8_t *inbox = a.my_slots + slot * a.slot_bytes; + switch (a.mode) { + case kGather: + copy_bytes(a.out + static_cast(1u - a.rank) * a.nbytes, inbox, a.nbytes, index, + stride); + break; + case kExchange: + copy_bytes(a.out, inbox, a.nbytes, index, stride); + break; + case kSumF32: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + default: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + } + } + + // 4. the last block publishes the epoch, unless a wait of this launch failed (later launches stay no-ops) + __syncthreads(); + if (tid == 0) { + // the counter is back at 0 for the next launch (stream order), whatever its grid + const bool last = gdim == 1u || atom_add_acq_rel_gpu(a.state + 1, 1u) == gdim - 1u; + const bool ok = gdim == 1u ? s_bad == 0u : ld_relaxed_gpu(a.state + 2) == 0u; + if (last && gdim > 1u) { + st_release_gpu(a.state + 1, 0u); + } + if (last && ok) { + st_release_gpu(a.state, seq); + } + } +} + +// The RoCE variant (roce.zig): blocks stage into the send slot, the last rings the doorbell, a host proxy RDMA-writes. +namespace tftp { + +constexpr int kRoceHcas = 2; +// ctrl words shared with roce.zig's proxy: [0] doorbell seq, [4 + slot] padded bytes, [8] completed, [16..21] failure +constexpr int kRoceSeq = 0; +constexpr int kRoceSlotBytes = 4; +constexpr int kRoceCompleted = 8; +constexpr int kRoceFailed = 16; +constexpr int kRoceAbort = 12; // wait_flag reads ctrl[kCtrlAbort]: the RoCE abort word is passed as that ctrl + +struct RoceArgs { + const uint8_t *in; + uint8_t *out; + uint8_t *send_slots; // this rank's send slots (pinned host, read by the NIC) + const uint8_t *recv_slots; // this rank's receive slots (written by the peer's NIC) + const uint32_t *flags; // this rank's flags [2][kRoceHcas], written by the peer's NIC + uint32_t *ctrl; + uint32_t *state; // device: [0] epoch, [1] tail counter, [2] poison, [3] stage counter + uint64_t nbytes; + uint64_t slot_bytes; + uint64_t timeout_ns; + uint32_t rank; + uint32_t mode; + uint32_t n_hca; + uint32_t opts; // TF_TP_ROCE_FAST: bit 0 fast, bits 8-15 pollers a flag, bits 16-31 start stagger ns; 0: classic + // gather only (nullptr: every shard is nbytes): lens[r] = rank r's bytes, padded to 16 on the wire + const int32_t *lens; +}; + +// Rank r's bytes this op: ``nbytes``, or lens[r] capped at nbytes. +__device__ __forceinline__ uint64_t roce_shard(const RoceArgs &a, uint32_t r) { + if (a.lens == nullptr) { + return a.nbytes; + } + const int32_t v = __ldg(a.lens + r); + const uint64_t n = v > 0 ? static_cast(v) : 0u; + return n < a.nbytes ? n : a.nbytes; +} + +// The bytes a copy of an n-byte shard moves: n rounded up to 16 when that stays inside a 16-aligned stride, else n. +__device__ __forceinline__ uint64_t roce_span(const RoceArgs &a, uint64_t n) { + if (a.lens == nullptr || (a.nbytes & 15u) != 0u) { + return n; + } + return (n + 15u) & ~uint64_t(15u); +} + +} // namespace tftp + +extern "C" __global__ void __launch_bounds__(512) tp_roce(RoceArgs a) { + __shared__ uint32_t s_head[2]; // poison, seq + __shared__ uint32_t s_bad; + __shared__ uint32_t s_done[kRoceHcas]; // fast: some poller of HCA h saw its flag + const bool fast = (a.opts & 1u) != 0u; + const uint32_t tid = threadIdx.x; + const uint32_t gdim = gridDim.x; + const uint64_t index = static_cast(blockIdx.x) * blockDim.x + tid; + const uint64_t stride = static_cast(gdim) * blockDim.x; + + if (tid == 0) { + s_head[0] = ld_relaxed_gpu(a.state + 2); + s_head[1] = ld_relaxed_gpu(a.state) + 1u; + s_bad = 0u; + } + if (tid < static_cast(kRoceHcas)) { + s_done[tid] = 0u; + } + __syncthreads(); + if (s_head[0] != 0u) { + return; + } + const uint32_t seq = s_head[1]; + const uint64_t slot = seq & 1u; + const bool var = a.mode == kGather && a.lens != nullptr; + const uint64_t mine = var ? roce_shard(a, a.rank) : a.nbytes; + const uint64_t theirs = var ? roce_shard(a, 1u - a.rank) : a.nbytes; + const uint64_t padded = var ? (mine < 16u ? 16u : (mine + 15u) & ~uint64_t(15u)) : (a.nbytes + 15u) & ~uint64_t(15u); + + // 1. stage this block's share into the send slot; the last block to finish rings the doorbell + copy_any(fast, a.send_slots + slot * a.slot_bytes, a.in, var ? roce_span(a, mine) : a.nbytes, index, stride); + __syncthreads(); + if (tid == 0 && fast) { + // lean doorbell: release RMW per block arrival, the last acquires, then a release store (no SC system fences) + const bool last = gdim == 1u || atom_add_acq_rel_sys(a.state + 3, 1u) == gdim - 1u; + if (last) { + if (gdim > 1u) { + st_release_gpu(a.state + 3, 0u); + } + st_relaxed_sys(a.ctrl + kRoceSlotBytes + slot, static_cast(padded)); + st_release_sys(a.ctrl + kRoceSeq, seq); + } + } else if (tid == 0) { + fence_sc_sys(); + const bool last = gdim == 1u || atom_add_acq_rel_gpu(a.state + 3, 1u) == gdim - 1u; + if (last) { + if (gdim > 1u) { + st_release_gpu(a.state + 3, 0u); + fence_sc_sys(); // every block's staging (seen through the counter) before the doorbell + } + st_relaxed_sys(a.ctrl + kRoceSlotBytes + slot, static_cast(padded)); + st_release_sys(a.ctrl + kRoceSeq, seq); + } + } + if (a.mode == kGather) { + copy_any(fast, a.out + static_cast(a.rank) * a.nbytes, a.in, + var ? roce_span(a, mine) : a.nbytes, index, stride); + } + + // 2. every HCA's flag for this slot; fast: poller k of HCA h is thread 32 k + h, started k x stagger ns late + const uint32_t pollers = (a.opts >> 8) & 0xffu; + const uint32_t stagger_ns = a.opts >> 16; + if (fast) { + const uint32_t lane = tid & 31u; + const uint32_t k = tid >> 5; + if (lane < a.n_hca && k < (pollers == 0u ? 1u : pollers)) { + if (k > 0u && stagger_ns != 0u) { + __nanosleep(k * stagger_ns); + } + uint32_t seen = 0u, why = 0u, waited = 0u; + if (!wait_flag_shared(a.flags + (slot * kRoceHcas + lane) * kFlagStride, seq, a.ctrl + kRoceAbort, + a.timeout_ns, &seen, &why, &waited, &s_done[lane])) { + s_bad = 1u; + if (atom_cas_gpu(a.state + 2, 0u, 1u) == 0u) { + st_relaxed_sys(a.ctrl + kRoceFailed + 1, seq); + st_relaxed_sys(a.ctrl + kRoceFailed + 2, seen); + st_relaxed_sys(a.ctrl + kRoceFailed + 3, lane); + st_relaxed_sys(a.ctrl + kRoceFailed + 4, why); + st_relaxed_sys(a.ctrl + kRoceFailed + 5, waited); + fence_sc_sys(); + st_relaxed_sys(a.ctrl + kRoceFailed, 1u); + fence_sc_sys(); + } + } + } + } else if (tid == 0) { + for (uint32_t h = 0; h < a.n_hca; h++) { + uint32_t seen = 0u, why = 0u, waited = 0u; + if (!wait_flag(a.flags + (slot * kRoceHcas + h) * kFlagStride, seq, a.ctrl + kRoceAbort, a.timeout_ns, &seen, + &why, &waited)) { + s_bad = 1u; + if (atom_cas_gpu(a.state + 2, 0u, 1u) == 0u) { + st_relaxed_sys(a.ctrl + kRoceFailed + 1, seq); + st_relaxed_sys(a.ctrl + kRoceFailed + 2, seen); + st_relaxed_sys(a.ctrl + kRoceFailed + 3, h); + st_relaxed_sys(a.ctrl + kRoceFailed + 4, why); + st_relaxed_sys(a.ctrl + kRoceFailed + 5, waited); + fence_sc_sys(); + st_relaxed_sys(a.ctrl + kRoceFailed, 1u); + fence_sc_sys(); + } + break; + } + } + } + __syncthreads(); + + // 3. copy out of the receive slot the peer's NIC wrote + if (s_bad == 0u) { + const uint8_t *inbox = a.recv_slots + slot * a.slot_bytes; + switch (a.mode) { + case kGather: + copy_any(fast, a.out + static_cast(1u - a.rank) * a.nbytes, inbox, + var ? roce_span(a, theirs) : a.nbytes, index, stride); + break; + case kExchange: + copy_any(fast, a.out, inbox, a.nbytes, index, stride); + break; + case kSumF32: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + default: + sum_bytes(a.out, a.in, inbox, a.nbytes, a.rank, index, stride); + break; + } + } + + // 4. the epoch, as tp_mailbox; the host also sees the last completed seq + __syncthreads(); + if (tid == 0) { + const bool last = gdim == 1u || atom_add_acq_rel_gpu(a.state + 1, 1u) == gdim - 1u; + const bool ok = gdim == 1u ? s_bad == 0u : ld_relaxed_gpu(a.state + 2) == 0u; + if (last && gdim > 1u) { + st_release_gpu(a.state + 1, 0u); + } + if (last && ok) { + st_release_gpu(a.state, seq); + st_relaxed_sys(a.ctrl + kRoceCompleted, seq); + } + } +} diff --git a/zig/src/fabric/tp/mailbox.zig b/zig/src/fabric/tp/mailbox.zig new file mode 100644 index 0000000000..902bb79393 --- /dev/null +++ b/zig/src/fabric/tp/mailbox.zig @@ -0,0 +1,215 @@ +//! The mailbox backend: two-rank exchanges in one kernel through pinned host memory both GPUs map (host side). +const std = @import("std"); +const log = @import("log.zig"); +const cuda = @import("cuda"); +const collective = @import("collective.zig"); +const Bootstrap = @import("bootstrap.zig").Bootstrap; +const sock = @import("sock.zig"); +const posix = std.posix; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Error = collective.Error; + +pub const max_blocks = 16; +const threads = 512; +const flag_bytes = 2 * max_blocks * 128; +const ctrl_bytes = 4096; +const ctrl_abort = 0; +const ctrl_rank = 16; + +const Mode = enum(u32) { gather = 0, sum_f32 = 1, sum_bf16 = 2, exchange = 3 }; + +/// mailbox.cu's `MailArgs`, field for field. +const MailArgs = extern struct { + in: DevicePtr, + out: DevicePtr, + peer_slots: DevicePtr, + peer_flags: DevicePtr, + my_slots: DevicePtr, + my_flags: DevicePtr, + state: DevicePtr, + ctrl: DevicePtr, + nbytes: u64, + slot_bytes: u64, + timeout_ns: u64, + rank: u32, + mode: u32, +}; + +/// Driver calls the runtime's table does not carry (page-locking memory this process mapped itself). +pub const HostRegister = struct { + lib: std.DynLib, + register: *const fn (?*anyopaque, usize, c_uint) callconv(.c) cuda.abi.Result, + unregister: *const fn (?*anyopaque) callconv(.c) cuda.abi.Result, + + pub fn open() !HostRegister { + var lib = std.DynLib.open("libcuda.so.1") catch return error.DriverUnavailable; + errdefer lib.close(); + return .{ + .register = lib.lookup(@FieldType(HostRegister, "register"), "cuMemHostRegister_v2") orelse return error.MissingSymbol, + .unregister = lib.lookup(@FieldType(HostRegister, "unregister"), "cuMemHostUnregister") orelse return error.MissingSymbol, + .lib = lib, + }; + } +}; + +/// Grid for `n` bytes a rank: one block per 16 KiB, a power of two, at most `max_blocks` (both ranks pick the same). +pub fn blocksFor(n: usize) u32 { + const want = std.math.divCeil(usize, @max(n, 1), 16 * 1024) catch 1; + return @intCast(@min(std.math.ceilPowerOfTwoAssert(usize, want), max_blocks)); +} + +pub const Mailbox = struct { + d: *const cuda.Driver, + me: u32, + slot_bytes: usize, + timeout_ns: u64, + region: []align(std.heap.page_size_min) u8, + dev: DevicePtr, + hr: HostRegister, + state: cuda.DeviceBuffer, + module: cuda.Module, + func: cuda.Function, + aborted: std.atomic.Value(bool) = .init(false), + + /// Both ranks call it after the bootstrap (two ranks, one host); `image` is mailbox.cu's fatbin. + pub fn init(d: *const cuda.Driver, b: *Bootstrap, image: []const u8, max_bytes: usize, timeout_ns: u64) !Mailbox { + if (b.world != 2) return error.Unsupported; + if (image.len == 0) return error.NoKernel; + const slot = std.mem.alignForward(usize, max_bytes, 4096); + const len = ctrl_bytes + 2 * (flag_bytes + 2 * slot); + // rank 0 creates and sizes the file before anyone learns its name, the others open it, then rank 0 unlinks it + var name: [64]u8 = undefined; + var name_len: u8 = 0; + var fd: std.c.fd_t = -1; + defer if (fd >= 0) { + _ = std.c.close(fd); + }; + if (b.rank == 0) blk: { + const s = std.fmt.bufPrint(&name, "/dev/shm/tf-tp-{d}-{x}", .{ std.c.getpid(), sock.nowNs() }) catch break :blk; + name[s.len] = 0; + fd = std.c.open(name[0..s.len :0], .{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true, .CLOEXEC = true }, @as(std.c.mode_t, 0o600)); + if (fd >= 0 and std.c.ftruncate(fd, @intCast(len)) == 0) name_len = @intCast(s.len); + } + // an empty name tells the other rank that rank 0 failed (it must not wait for a name that never comes) + try b.broadcast(std.mem.asBytes(&name_len), std.mem.asBytes(&name_len)); + if (name_len == 0) return error.SharedMemoryUnavailable; + try b.broadcast(name[0..name_len], name[0..name_len]); + name[name_len] = 0; + const path = name[0..name_len :0]; + if (b.rank != 0) { + fd = std.c.open(path, .{ .ACCMODE = .RDWR, .CLOEXEC = true }, @as(std.c.mode_t, 0)); + if (fd < 0) return error.SharedMemoryUnavailable; + } + try b.agree("mapped", "mailbox file"); + const region = try posix.mmap(null, len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); + errdefer posix.munmap(region); + try b.agree("open", "mailbox mapping"); + if (b.rank == 0) _ = std.c.unlink(path); + var hr = try HostRegister.open(); + errdefer hr.lib.close(); + try d.check(hr.register(region.ptr, len, 1 | 2), "cuMemHostRegister"); // PORTABLE | DEVICEMAP + errdefer _ = hr.unregister(region.ptr); + var dev: DevicePtr = 0; + try d.check(d.api.cuMemHostGetDevicePointer_v2(&dev, region.ptr, 0), "cuMemHostGetDevicePointer"); + var state = try cuda.DeviceBuffer.alloc(d, 16); + errdefer state.free(); + try state.fill8(0, null); + var module = try cuda.Module.load(d, image); + errdefer module.unload(); + const func = try module.function("tp_mailbox"); + try b.agree("ready", "mailbox setup"); + return .{ .d = d, .me = b.rank, .slot_bytes = slot, .timeout_ns = timeout_ns, .region = region, .dev = dev, .hr = hr, .state = state, .module = module, .func = func }; + } + + pub fn deinit(m: *Mailbox) void { + m.module.unload(); + m.state.free(); + _ = m.hr.unregister(m.region.ptr); + m.hr.lib.close(); + posix.munmap(m.region); + } + + fn inbox(m: *const Mailbox, r: u32) usize { + return ctrl_bytes + r * (flag_bytes + 2 * m.slot_bytes); + } + + fn ctrlWord(m: *const Mailbox, i: usize) *std.atomic.Value(u32) { + return @ptrCast(@alignCast(m.region[4 * i ..].ptr)); + } + + /// Whether `n` bytes a rank fit one slot. + pub fn fits(m: *const Mailbox, n: usize) bool { + return n <= m.slot_bytes; + } + + fn launch(m: *Mailbox, mode: Mode, in: DevicePtr, out: DevicePtr, n: usize, stream: collective.Stream) Error!void { + if (m.aborted.load(.acquire)) return error.Aborted; + if (n > m.slot_bytes) return error.Unsupported; + if (n == 0) return; + const peer = 1 - m.me; + var a: MailArgs = .{ + .in = in, + .out = out, + .peer_slots = m.dev + m.inbox(peer) + flag_bytes, + .peer_flags = m.dev + m.inbox(peer), + .my_slots = m.dev + m.inbox(m.me) + flag_bytes, + .my_flags = m.dev + m.inbox(m.me), + .state = m.state.ptr, + .ctrl = m.dev, + .nbytes = n, + .slot_bytes = m.slot_bytes, + .timeout_ns = m.timeout_ns, + .rank = m.me, + .mode = @backingInt(mode), + }; + var params = [1]?*anyopaque{@ptrCast(&a)}; + const r = m.d.api.cuLaunchKernel(m.func.handle, blocksFor(n), 1, 1, threads, 1, 1, 0, stream, ¶ms, null); + m.d.check(r, "cuLaunchKernel tp_mailbox") catch return error.BackendFailed; + } + + pub fn allGatherBytes(m: *Mailbox, s: DevicePtr, r: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return m.launch(.gather, s, r, n, stream); + } + + pub fn exchangeBytes(m: *Mailbox, s: DevicePtr, r: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return m.launch(.exchange, s, r, n, stream); + } + + /// Sums only, f32 or bf16 (what the model's partials are); the rest is `error.Unsupported`. + pub fn sum(m: *Mailbox, s: DevicePtr, r: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const mode: Mode = switch (t) { + .f32 => .sum_f32, + .bf16 => .sum_bf16, + else => return error.Unsupported, + }; + return m.launch(mode, s, r, count * t.size(), stream); + } + + /// Ends pending and later exchanges on both ranks (the abort word is shared); this rank's launches stop. + pub fn abort(m: *Mailbox) void { + m.aborted.store(true, .release); + m.ctrlWord(ctrl_abort).store(1, .release); + } + + /// A wait that timed out or was aborted on either rank. + pub fn check(m: *Mailbox) Error!void { + for (0..2) |r| { + const rec = ctrl_rank + 8 * r; + if (m.ctrlWord(rec).load(.acquire) != 0) { + log.warn("tp mailbox: rank {d}'s wait for seq {d} failed ({s} after {d} us; block {d}, flag held {d})", .{ r, m.ctrlWord(rec + 1).load(.acquire), if (m.ctrlWord(rec + 4).load(.acquire) == 1) "aborted" else "timed out", m.ctrlWord(rec + 5).load(.acquire), m.ctrlWord(rec + 3).load(.acquire), m.ctrlWord(rec + 2).load(.acquire) }); + return error.PeerFailed; + } + } + if (m.aborted.load(.acquire) or m.ctrlWord(ctrl_abort).load(.acquire) != 0) return error.Aborted; + } +}; + +test "grids grow with the shard and stay powers of two" { + try std.testing.expectEqual(@as(u32, 1), blocksFor(4096)); + try std.testing.expectEqual(@as(u32, 1), blocksFor(16 * 1024)); + try std.testing.expectEqual(@as(u32, 2), blocksFor(16 * 1024 + 1)); + try std.testing.expectEqual(@as(u32, 4), blocksFor(64 * 1024)); + try std.testing.expectEqual(@as(u32, 16), blocksFor(4 << 20)); + try std.testing.expectEqual(@as(usize, 96), @sizeOf(MailArgs)); +} diff --git a/zig/src/fabric/tp/nccl.zig b/zig/src/fabric/tp/nccl.zig new file mode 100644 index 0000000000..b641839aad --- /dev/null +++ b/zig/src/fabric/tp/nccl.zig @@ -0,0 +1,163 @@ +//! The NCCL backend: one communicator a rank, every collective one enqueue on the caller's stream. +const std = @import("std"); +const cuda = @import("cuda"); +const collective = @import("collective.zig"); +const nccl = cuda.nccl; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Op = collective.Op; +const Error = collective.Error; + +pub fn dataType(t: DType) nccl.DataType { + return switch (t) { + .u8 => .u8, + .i32 => .i32, + .i64 => .i64, + .f16 => .f16, + .bf16 => .bf16, + .f32 => .f32, + .f64 => .f64, + }; +} + +pub fn redOp(op: Op) nccl.RedOp { + return switch (op) { + .sum => .sum, + .max => .max, + .min => .min, + }; +} + +pub const Nccl = struct { + lib: *const nccl.Library, + d: *const cuda.Driver, + comm: nccl.Comm, + me: u32, + size: u32, + aborted: std.atomic.Value(bool) = .init(false), + /// The barrier's one-float buffers and its own stream (host path only). + scratch: cuda.DeviceBuffer, + stream: cuda.Stream, + + /// Joins the communicator `id` names (every rank calls it with the same id; it returns once all have). + pub fn init(lib: *const nccl.Library, d: *const cuda.Driver, id: nccl.UniqueId, rank: u32, world: u32) !Nccl { + var comm: nccl.Comm = null; + try lib.check(lib.api.ncclCommInitRank(&comm, @intCast(world), id, @intCast(rank)), "ncclCommInitRank"); + errdefer _ = lib.api.ncclCommDestroy(comm); + var scratch = try cuda.DeviceBuffer.alloc(d, 4 * (world + 1)); + errdefer scratch.free(); + try scratch.fill8(0, null); + return .{ .lib = lib, .d = d, .comm = comm, .me = rank, .size = world, .scratch = scratch, .stream = try cuda.Stream.init(d, true) }; + } + + pub fn deinit(n: *Nccl) void { + _ = n.stream.synchronize() catch {}; + n.stream.deinit(); + n.scratch.free(); + if (!n.aborted.load(.acquire)) _ = n.lib.api.ncclCommDestroy(n.comm); + } + + pub fn iface(n: *Nccl) collective.Collective { + return .{ .ptr = n, .vtable = &vtable }; + } + + const vtable: collective.Collective.VTable = .{ + .kind = .nccl, + .rank = rankOf, + .world = worldOf, + .all_reduce = allReduce, + .all_gather = allGather, + .exchange = exchange, + .send = send, + .recv = recv, + .barrier = barrier, + .check = check, + .abort = abort, + }; + + fn self_(ptr: *anyopaque) *Nccl { + return @ptrCast(@alignCast(ptr)); + } + + fn rankOf(ptr: *anyopaque) u32 { + return self_(ptr).me; + } + + fn worldOf(ptr: *anyopaque) u32 { + return self_(ptr).size; + } + + fn call(n: *Nccl, r: nccl.Result, what: []const u8) Error!void { + if (r == 0) return; + if (n.aborted.load(.acquire)) return error.Aborted; + n.lib.check(r, what) catch return error.BackendFailed; + } + + fn live(n: *Nccl) Error!void { + if (n.aborted.load(.acquire)) return error.Aborted; + } + + pub fn allReduce(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, op: Op, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try n.call(n.lib.api.ncclAllReduce(s, r, count, dataType(t), redOp(op), n.comm, stream), "ncclAllReduce"); + } + + pub fn allGather(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try n.call(n.lib.api.ncclAllGather(s, r, count, dataType(t), n.comm, stream), "ncclAllGather"); + } + + pub fn exchange(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try collective.checkPeer(n.me, n.size, peer); + try n.call(n.lib.api.ncclGroupStart(), "ncclGroupStart"); + const a = n.lib.api.ncclSend(s, count, dataType(t), @intCast(peer), n.comm, stream); + const b = n.lib.api.ncclRecv(r, count, dataType(t), @intCast(peer), n.comm, stream); + const e = n.lib.api.ncclGroupEnd(); + try n.call(a, "ncclSend"); + try n.call(b, "ncclRecv"); + try n.call(e, "ncclGroupEnd"); + } + + fn send(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try collective.checkPeer(n.me, n.size, peer); + try n.call(n.lib.api.ncclSend(buf, count, dataType(t), @intCast(peer), n.comm, stream), "ncclSend"); + } + + fn recv(ptr: *anyopaque, buf: DevicePtr, count: usize, t: DType, peer: u32, stream: collective.Stream) Error!void { + const n = self_(ptr); + try n.live(); + try collective.checkPeer(n.me, n.size, peer); + try n.call(n.lib.api.ncclRecv(buf, count, dataType(t), @intCast(peer), n.comm, stream), "ncclRecv"); + } + + /// One float all-gathered on the barrier's own stream, then a sync of it (host path, never captured). + fn barrier(ptr: *anyopaque) Error!void { + const n = self_(ptr); + try n.live(); + try n.call(n.lib.api.ncclAllGather(n.scratch.ptr, n.scratch.ptr + 4, 1, .f32, n.comm, n.stream.handle), "ncclAllGather"); + n.stream.synchronize() catch return if (n.aborted.load(.acquire)) error.Aborted else error.BackendFailed; + } + + fn check(ptr: *anyopaque) Error!void { + const n = self_(ptr); + try n.live(); + var r: nccl.Result = 0; + try n.call(n.lib.api.ncclCommGetAsyncError(n.comm, &r), "ncclCommGetAsyncError"); + if (r != 0 and r != nccl.in_progress) { + n.lib.check(r, "NCCL async error") catch {}; + return error.PeerFailed; + } + } + + fn abort(ptr: *anyopaque) void { + const n = self_(ptr); + if (n.aborted.swap(true, .acq_rel)) return; + _ = n.lib.api.ncclCommAbort(n.comm); + } +}; diff --git a/zig/src/fabric/tp/planlink.zig b/zig/src/fabric/tp/planlink.zig new file mode 100644 index 0000000000..44dbb61a0c --- /dev/null +++ b/zig/src/fabric/tp/planlink.zig @@ -0,0 +1,467 @@ +//! The round plan from rank 0 to the other ranks over its own TCP connection: a 4-byte length, then int64 words. +const std = @import("std"); +const log = @import("log.zig"); +const sock = @import("sock.zig"); +const builtin = @import("builtin"); +const config = @import("config.zig"); +const linux = std.os.linux; + +pub const max_ranks = 64; + +/// The CPU a plan thread pinned itself to (TF_TP_PLAN_PIN), -1 until one did; an engine's CPU isolation keeps off it. +pub var pinned_cpu: std.atomic.Value(i32) = .init(-1); +const more: u32 = 0x8000; // MSG_MORE: the length and the ints leave as one segment + +pub const Error = error{ PeerClosed, TooLarge, NotRankZero, NotFollower }; + +/// A batch frame's first word (a plan's op is >= 0): [batch, n, len 1, plan 1, ..., len n, plan n]. +pub const batch: i64 = std.math.minInt(i64); +/// Words rank 0 holds at most (more flushes the held plans first), and the most a batch frame carries. +pub const hold_words: usize = 1024; +pub const batch_words: usize = 2 * hold_words + 2; + +/// Largest plan either side accepts: 64M words (512 MiB), so a garbled length cannot become a 4 GiB allocation. +pub const max_frame_words: usize = 64 << 20; + +comptime { + std.debug.assert(builtin.cpu.arch.endian() == .little); // the ints go out as they are in memory +} + +pub const PlanLink = struct { + rank: u32, + world: u32, + /// Rank 0: rank i's socket at index i; the others: rank 0's at index 0 (the bootstrap's `releasePlan`). + fds: [max_ranks]sock.Fd, + sent: u64 = 0, + received: u64 = 0, + /// TF_TP_PLAN_LINK=rdma (config.plan_spin_ns): the follower polls the socket this long before it blocks. + spin_ns: u64 = 0, + /// TF_TP_PLAN_PIN: the sending / receiving thread pins itself at its first plan; `avoid`: the RoCE proxy's CPU. + pin: ?config.PlanPin = null, + avoid: ?u32 = null, + pinned: bool = false, + /// rank 0: plans held for the next send (`hold`), each as [len, plan...]; `held_n` of them + held: [hold_words]i64 = undefined, + held_len: usize = 0, + held_n: usize = 0, + /// the follower: a batch frame's plans not handed out yet, as [len, plan...] from `inbox_at` + inbox: [batch_words]i64 = undefined, + inbox_len: usize = 0, + inbox_at: usize = 0, + batches: u64 = 0, + + pub fn init(rank: u32, world: u32, fds: [max_ranks]sock.Fd) PlanLink { + return .{ .rank = rank, .world = world, .fds = fds }; + } + + /// Rank 0: one plan to every other rank (host only); held plans go first, in the same frame when they fit. + pub fn send(p: *PlanLink, ints: []const i64) Error!void { + if (p.rank != 0) return error.NotRankZero; + if (ints.len > max_frame_words) return error.TooLarge; + if (p.held_n > 0) { + if (2 + p.held_len + 1 + ints.len <= batch_words) { + var frame: [batch_words]i64 = undefined; + frame[0] = batch; + frame[1] = @intCast(p.held_n + 1); + @memcpy(frame[2..][0..p.held_len], p.held[0..p.held_len]); + frame[2 + p.held_len] = @intCast(ints.len); + @memcpy(frame[3 + p.held_len ..][0..ints.len], ints); + const n = 3 + p.held_len + ints.len; + p.held_len = 0; + p.held_n = 0; + p.batches += 1; + return p.sendFrame(frame[0..n]); + } + try p.flush(); + } + return p.sendFrame(ints); + } + + /// Rank 0: a plan the followers run before the next sent one, carried in that plan's frame (`send` or `flush`). + pub fn hold(p: *PlanLink, ints: []const i64) Error!void { + if (p.rank != 0) return error.NotRankZero; + if (1 + ints.len > hold_words) { + try p.flush(); + return p.sendFrame(ints); + } + if (p.held_len + 1 + ints.len > hold_words) try p.flush(); + p.held[p.held_len] = @intCast(ints.len); + @memcpy(p.held[p.held_len + 1 ..][0..ints.len], ints); + p.held_len += 1 + ints.len; + p.held_n += 1; + } + + /// Rank 0: the held plans now, as one batch frame (nothing held: nothing sent). + pub fn flush(p: *PlanLink) Error!void { + if (p.held_n == 0) return; + var frame: [batch_words]i64 = undefined; + frame[0] = batch; + frame[1] = @intCast(p.held_n); + @memcpy(frame[2..][0..p.held_len], p.held[0..p.held_len]); + const n = 2 + p.held_len; + p.held_len = 0; + p.held_n = 0; + p.batches += 1; + return p.sendFrame(frame[0..n]); + } + + fn sendFrame(p: *PlanLink, ints: []const i64) Error!void { + p.pinOnce(); + const bytes = std.mem.sliceAsBytes(ints); + const len: u32 = @intCast(bytes.len); + for (p.fds[1..p.world]) |fd| { + if (std.c.send(fd, &len, 4, std.posix.MSG.NOSIGNAL | more) != 4) return error.PeerClosed; + sock.sendAll(fd, bytes) catch return error.PeerClosed; + } + p.sent += 1; + } + + /// The other ranks: the next plan into `out`; `error.PeerClosed` once rank 0 is gone, `error.TooLarge` skips one. + pub fn recv(p: *PlanLink, out: []i64) Error![]i64 { + if (p.inboxNext()) |m| { + if (m.len > out.len) return error.TooLarge; + @memcpy(out[0..m.len], m); + return out[0..m.len]; + } + const n = try p.header(); + if (n > out.len) { + try p.drop(n); + return error.TooLarge; + } + const got = try p.body(out[0..n]); + if (got.len > 0 and got[0] == batch) { + try p.unpack(got); + return p.recv(out); + } + return got; + } + + /// The next plan into `buf`, grown to the frame; the buffer is kept, so steady plans allocate nothing. + pub fn recvGrow(p: *PlanLink, gpa: std.mem.Allocator, buf: *std.ArrayList(i64)) (Error || std.mem.Allocator.Error)![]i64 { + if (p.inboxNext()) |m| { + try buf.resize(gpa, m.len); + @memcpy(buf.items, m); + return buf.items; + } + const n = try p.header(); + if (n > max_frame_words) { + try p.drop(n); + return error.TooLarge; + } + try buf.resize(gpa, n); + const got = try p.body(buf.items); + if (got.len > 0 and got[0] == batch) { + try p.unpack(got); + return p.recvGrow(gpa, buf); + } + return got; + } + + /// The follower: a batch frame's plans into the inbox (checked whole first: a garbled batch ends the link). + fn unpack(p: *PlanLink, frame: []const i64) Error!void { + if (frame.len < 2 or frame.len > batch_words or frame[1] < 1) return error.PeerClosed; + var at: usize = 2; + var k: i64 = 0; + while (k < frame[1]) : (k += 1) { + if (at >= frame.len or frame[at] < 0) return error.PeerClosed; + at += 1 + @as(usize, @intCast(frame[at])); + if (at > frame.len) return error.PeerClosed; + } + if (at != frame.len) return error.PeerClosed; + @memcpy(p.inbox[0 .. frame.len - 2], frame[2..]); + p.inbox_len = frame.len - 2; + p.inbox_at = 0; + p.batches += 1; + } + + /// The follower: the next plan of the last batch frame, if any (a slice of the inbox, valid until the next recv). + fn inboxNext(p: *PlanLink) ?[]const i64 { + if (p.inbox_at >= p.inbox_len) return null; + const n: usize = @intCast(p.inbox[p.inbox_at]); + const m = p.inbox[p.inbox_at + 1 ..][0..n]; + p.inbox_at += 1 + n; + return m; + } + + /// The next frame's length in words. + fn header(p: *PlanLink) Error!usize { + if (p.rank == 0) return error.NotFollower; + p.pinOnce(); + if (p.spin_ns > 0) p.spin(); + var len: u32 = undefined; + sock.recvExact(p.fds[0], std.mem.asBytes(&len), null) catch return error.PeerClosed; + if (len % 8 != 0) return error.PeerClosed; // not a frame boundary: the stream is lost + return len / 8; + } + + fn body(p: *PlanLink, out: []i64) Error![]i64 { + sock.recvExact(p.fds[0], std.mem.sliceAsBytes(out), null) catch return error.PeerClosed; + p.received += 1; + return out; + } + + /// Reads and drops `n` words. + fn drop(p: *PlanLink, n: usize) Error!void { + var left: usize = n * 8; + var scratch: [4096]u8 = undefined; + while (left > 0) : (left -= @min(left, scratch.len)) sock.recvExact(p.fds[0], scratch[0..@min(left, scratch.len)], null) catch return error.PeerClosed; + } + + /// Polls rank 0's socket until it is readable (or hung up) or `spin_ns` passed; the read that follows blocks. + fn spin(p: *PlanLink) void { + if (comptime builtin.os.tag != .linux) return; + var fd = [1]linux.pollfd{.{ .fd = p.fds[0], .events = linux.POLL.IN, .revents = 0 }}; + const end = sock.nowNs() + p.spin_ns; + while (true) { + const rc = linux.poll(&fd, 1, 0); + if (linux.errno(rc) == .SUCCESS and rc > 0) return; + if (sock.nowNs() >= end) return; + std.atomic.spinLoopHint(); + } + } + + /// Pins the calling thread once: an explicit CPU, or auto's highest cpu_capacity. + fn pinOnce(p: *PlanLink) void { + if (p.pinned) return; + p.pinned = true; + if (comptime builtin.os.tag != .linux) return; + const want = p.pin orelse return; + const cpu = switch (want) { + .cpu => |c| c, + .auto => bestCpu(p.avoid) orelse { + log.info("[tensorfold] tp: TF_TP_PLAN_PIN=auto: no cpu_capacity differs among the allowed CPUs; the plan thread stays unpinned", .{}); + return; + }, + }; + pinTo(cpu) catch |e| { + log.warn("[tensorfold] tp: TF_TP_PLAN_PIN: CPU {d} refused ({t}); the plan thread stays unpinned", .{ cpu, e }); + return; + }; + pinned_cpu.store(@intCast(cpu), .release); + log.info("[tensorfold] tp: plan link rank {d} pinned to CPU {d}", .{ p.rank, cpu }); + } + + pub fn close(p: *PlanLink) void { + for (&p.fds) |*fd| if (fd.* >= 0) { + sock.close(fd.*); + fd.* = -1; + }; + } +}; + +/// The calling thread onto `cpu` alone. +pub fn pinTo(cpu: u32) !void { + if (cpu >= linux.CPU_SETSIZE) return error.BadCpu; + var set: linux.cpu_set_t = @splat(0); + set[cpu / @bitSizeOf(usize)] |= @as(usize, 1) << @intCast(cpu % @bitSizeOf(usize)); + try linux.sched_setaffinity(0, &set); +} + +/// Auto pin: the allowed CPU (not `avoid`) with the highest cpu_capacity, lowest on ties; null when all are equal. +pub fn bestCpu(avoid: ?u32) ?u32 { + var set: linux.cpu_set_t = undefined; + if (linux.errno(linux.sched_getaffinity(0, @sizeOf(linux.cpu_set_t), &set)) != .SUCCESS) return null; + var best: ?u32 = null; + var best_cap: u32 = 0; + var low_cap: u32 = std.math.maxInt(u32); + for (set, 0..) |word, wi| { + var bits = word; + while (bits != 0) : (bits &= bits - 1) { + const cpu: u32 = @intCast(wi * @bitSizeOf(usize) + @ctz(bits)); + if (avoid == cpu) continue; + const c = capacity(cpu) orelse continue; + low_cap = @min(low_cap, c); + if (best == null or c > best_cap) { + best = cpu; + best_cap = c; + } + } + } + return if (best != null and best_cap > low_cap) best else null; +} + +fn capacity(cpu: u32) ?u32 { + var pb: [64]u8 = undefined; + const path = std.fmt.bufPrintSentinel(&pb, "/sys/devices/system/cpu/cpu{d}/cpu_capacity", .{cpu}, 0) catch return null; + const rc = linux.open(path, .{ .ACCMODE = .RDONLY, .CLOEXEC = true }, 0); + if (linux.errno(rc) != .SUCCESS) return null; + const fd: i32 = @intCast(rc); + defer _ = linux.close(fd); + var b: [16]u8 = undefined; + const n = linux.read(fd, &b, b.len); + if (linux.errno(n) != .SUCCESS) return null; + return std.fmt.parseInt(u32, std.mem.trim(u8, b[0..n], " \n"), 10) catch null; +} + +test "plans arrive whole and in order; a closed rank 0 ends the follower" { + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + var lead = PlanLink.init(0, 2, f0); + var follow = PlanLink.init(1, 2, f1); + defer follow.close(); + var big: [5000]i64 = undefined; + for (&big, 0..) |*v, i| v.* = @as(i64, @intCast(i)) * -3 + 0x7FFF_0000_0000; + try lead.send(&.{ 0x504C414E, 2, 1 }); + try lead.send(&big); + try lead.send(&.{}); + var buf: [8192]i64 = undefined; + try std.testing.expectEqualSlices(i64, &.{ 0x504C414E, 2, 1 }, try follow.recv(&buf)); + try std.testing.expectEqualSlices(i64, &big, try follow.recv(&buf)); + try std.testing.expectEqual(@as(usize, 0), (try follow.recv(&buf)).len); + try std.testing.expectError(error.NotRankZero, follow.send(&.{1})); + var small: [2]i64 = undefined; + try lead.send(&.{ 1, 2, 3 }); + try std.testing.expectError(error.TooLarge, follow.recv(&small)); + lead.close(); + try std.testing.expectError(error.PeerClosed, follow.recv(&buf)); +} + +test "an admission-sized plan (1.2M words) arrives whole through the growable path" { + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + var lead = PlanLink.init(0, 2, f0); + defer lead.close(); + var follow = PlanLink.init(1, 2, f1); + defer follow.close(); + const gpa = std.testing.allocator; + const prompt = try gpa.alloc(i64, 1_200_000); + defer gpa.free(prompt); + for (prompt, 0..) |*v, i| v.* = @intCast((i *% 2654435761) % 163840); + // the socket buffer holds far less than 9.6 MB: rank 0 sends from its own thread, as it would + const Send = struct { + fn go(link: *PlanLink, ints: []const i64) void { + link.send(&.{ 0x504C414E, 3, 20 }) catch {}; + link.send(ints) catch {}; + link.send(&.{ 0x504C414E, 3, 21 }) catch {}; + } + }; + const t = try std.Thread.spawn(.{}, Send.go, .{ &lead, prompt }); + var buf: std.ArrayList(i64) = .empty; + defer buf.deinit(gpa); + try std.testing.expectEqualSlices(i64, &.{ 0x504C414E, 3, 20 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, prompt, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 0x504C414E, 3, 21 }, try follow.recvGrow(gpa, &buf)); + t.join(); + try std.testing.expect(buf.capacity >= prompt.len); // kept for the next admission +} + +test "TF_TP_PLAN_LINK=rdma: a spinning follower gets every plan, late or early; TF_TP_PLAN_PIN pins the thread" { + if (builtin.os.tag != .linux) return error.SkipZigTest; + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + var lead = PlanLink.init(0, 2, f0); + defer lead.close(); + var follow = PlanLink.init(1, 2, f1); + defer follow.close(); + follow.spin_ns = 2 * std.time.ns_per_ms; + // the CPU this thread may run on first: the pin keeps it allowed, so the test runner is not starved + var set: linux.cpu_set_t = undefined; + _ = linux.sched_getaffinity(0, @sizeOf(linux.cpu_set_t), &set); + var first: u32 = 0; + for (set, 0..) |w, i| if (w != 0) { + first = @intCast(i * @bitSizeOf(usize) + @ctz(w)); + break; + }; + const Follower = struct { + fn run(f: *PlanLink, cpu: u32, out: *[3]i64, ok: *bool) void { + f.pin = .{ .cpu = cpu }; + var buf: [8]i64 = undefined; + const x = (f.recv(&buf) catch return)[0]; // arrives within the spin + const y = (f.recv(&buf) catch return)[0]; // arrives after it: the blocking read + var now: linux.cpu_set_t = undefined; + _ = linux.sched_getaffinity(0, @sizeOf(linux.cpu_set_t), &now); + var count: usize = 0; + for (now) |w| count += @popCount(w); + out.* = .{ x, y, @intCast(count) }; + ok.* = f.pinned and (now[cpu / @bitSizeOf(usize)] >> @intCast(cpu % @bitSizeOf(usize))) & 1 == 1; + } + }; + var out: [3]i64 = .{ 0, 0, 0 }; + var ok = false; + const t = try std.Thread.spawn(.{}, Follower.run, .{ &follow, first, &out, &ok }); + try lead.send(&.{11}); + var ts: linux.timespec = .{ .sec = 0, .nsec = 20 * std.time.ns_per_ms }; + _ = linux.nanosleep(&ts, null); + try lead.send(&.{22}); + t.join(); + try std.testing.expectEqual([3]i64{ 11, 22, 1 }, out); + try std.testing.expect(ok); + if (bestCpu(null)) |c| try std.testing.expect((set[c / @bitSizeOf(usize)] >> @intCast(c % @bitSizeOf(usize))) & 1 == 1); +} + +test "held plans leave with the next plan, in order, as one frame; a large plan after them flushes them first" { + const bootstrap = @import("bootstrap.zig"); + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("127.0.0.1", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [max_ranks]sock.Fd = @splat(-1); + var f1: [max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + const lead = try std.testing.allocator.create(PlanLink); + defer std.testing.allocator.destroy(lead); + lead.* = PlanLink.init(0, 2, f0); + const follow = try std.testing.allocator.create(PlanLink); + defer std.testing.allocator.destroy(follow); + follow.* = PlanLink.init(1, 2, f1); + defer follow.close(); + // a round: three keeps held, the window sent: one frame + try lead.hold(&.{ 22, 0, 3 }); + try lead.hold(&.{ 22, 1, 0 }); + try lead.hold(&.{ 22, 2, 5 }); + try lead.send(&.{ 21, 4, 9, 9, 9 }); + try std.testing.expectEqual(@as(u64, 1), lead.sent); + // a held plan, then one too large for a batch: two frames, the held one first + var big: [5000]i64 = undefined; + for (&big, 0..) |*v, i| v.* = @intCast(i); + try lead.hold(&.{ 22, 3, 1 }); + try lead.send(&big); + try std.testing.expectEqual(@as(u64, 3), lead.sent); + // flush alone, an empty plan held, a plain send after + try lead.hold(&.{}); + try lead.flush(); + try lead.flush(); // nothing held: nothing sent + try lead.send(&.{7}); + try std.testing.expectEqual(@as(u64, 5), lead.sent); + var buf: std.ArrayList(i64) = .empty; + defer buf.deinit(std.testing.allocator); + const gpa = std.testing.allocator; + try std.testing.expectEqualSlices(i64, &.{ 22, 0, 3 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 22, 1, 0 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 22, 2, 5 }, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqualSlices(i64, &.{ 21, 4, 9, 9, 9 }, try follow.recvGrow(gpa, &buf)); + var out: [8192]i64 = undefined; + try std.testing.expectEqualSlices(i64, &.{ 22, 3, 1 }, try follow.recv(&out)); + try std.testing.expectEqualSlices(i64, &big, try follow.recvGrow(gpa, &buf)); + try std.testing.expectEqual(@as(usize, 0), (try follow.recv(&out)).len); + try std.testing.expectEqualSlices(i64, &.{7}, try follow.recv(&out)); + try std.testing.expectEqual(@as(u64, 3), follow.batches); + try std.testing.expectError(error.NotRankZero, follow.hold(&.{1})); + lead.close(); +} diff --git a/zig/src/fabric/tp/reference.zig b/zig/src/fabric/tp/reference.zig new file mode 100644 index 0000000000..b98966fcc8 --- /dev/null +++ b/zig/src/fabric/tp/reference.zig @@ -0,0 +1,121 @@ +//! Host references for the collectives: reductions in rank order, bf16 / f16 added in fp32 and rounded once. +const std = @import("std"); +const collective = @import("collective.zig"); +const DType = collective.DType; +const Op = collective.Op; + +/// bf16 bits to fp32 (exact). +pub fn bf16ToF32(h: u16) f32 { + return @bitCast(@as(u32, h) << 16); +} + +/// fp32 to bf16 bits, round to nearest even; NaN stays a quiet NaN. +pub fn f32ToBf16(f: f32) u16 { + const u: u32 = @bitCast(f); + if (std.math.isNan(f)) return @intCast((u >> 16) | 0x40); + const round = 0x7FFF + ((u >> 16) & 1); + return @intCast((u +% round) >> 16); +} + +fn combine(comptime T: type, a: T, b: T, op: Op) T { + return switch (op) { + .sum => if (@typeInfo(T) == .int) a +% b else a + b, + .max => @max(a, b), + .min => @min(a, b), + }; +} + +fn reduceTyped(comptime T: type, acc: []u8, src: []const u8, op: Op) void { + const n = acc.len / @sizeOf(T); + for (0..n) |i| { + const a = std.mem.readInt(@Int(.unsigned, @bitSizeOf(T)), acc[i * @sizeOf(T) ..][0..@sizeOf(T)], .little); + const b = std.mem.readInt(@Int(.unsigned, @bitSizeOf(T)), src[i * @sizeOf(T) ..][0..@sizeOf(T)], .little); + const r: T = combine(T, @bitCast(a), @bitCast(b), op); + std.mem.writeInt(@Int(.unsigned, @bitSizeOf(T)), acc[i * @sizeOf(T) ..][0..@sizeOf(T)], @bitCast(r), .little); + } +} + +fn reduceHalf(acc: []u8, src: []const u8, op: Op, bf: bool) void { + const n = acc.len / 2; + for (0..n) |i| { + const a = std.mem.readInt(u16, acc[2 * i ..][0..2], .little); + const b = std.mem.readInt(u16, src[2 * i ..][0..2], .little); + const fa: f32 = if (bf) bf16ToF32(a) else @floatCast(@as(f16, @bitCast(a))); + const fb: f32 = if (bf) bf16ToF32(b) else @floatCast(@as(f16, @bitCast(b))); + const r = combine(f32, fa, fb, op); + const out: u16 = if (bf) f32ToBf16(r) else @bitCast(@as(f16, @floatCast(r))); + std.mem.writeInt(u16, acc[2 * i ..][0..2], out, .little); + } +} + +/// acc <- acc (op) src, element-wise; `acc` holds the lower ranks' result, `src` the next rank's input. +pub fn accumulate(acc: []u8, src: []const u8, t: DType, op: Op) void { + std.debug.assert(acc.len == src.len and acc.len % t.size() == 0); + switch (t) { + .u8 => reduceTyped(u8, acc, src, op), + .i32 => reduceTyped(i32, acc, src, op), + .i64 => reduceTyped(i64, acc, src, op), + .f32 => reduceTyped(f32, acc, src, op), + .f64 => reduceTyped(f64, acc, src, op), + .bf16 => reduceHalf(acc, src, op, true), + .f16 => reduceHalf(acc, src, op, false), + } +} + +/// out <- inputs[0] (op) inputs[1] (op) ..., in rank order. +pub fn reduce(out: []u8, inputs: []const []const u8, t: DType, op: Op) void { + @memcpy(out, inputs[0]); + for (inputs[1..]) |in| accumulate(out, in, t, op); +} + +/// Deterministic test data for `rank`'s input of `bytes`: finite values of type `t` that differ per rank and position. +pub fn fill(out: []u8, t: DType, rank: u32, seed: u64) void { + var rng = std.Random.DefaultPrng.init(seed ^ (@as(u64, rank) *% 0x9E3779B97F4A7C15)); + const r = rng.random(); + const n = out.len / t.size(); + for (0..n) |i| { + const x = r.float(f32) * 8.0 - 4.0; + switch (t) { + .u8 => out[i] = r.int(u8), + .i32 => std.mem.writeInt(i32, out[4 * i ..][0..4], r.int(i32), .little), + .i64 => std.mem.writeInt(i64, out[8 * i ..][0..8], r.int(i64), .little), + .f32 => std.mem.writeInt(u32, out[4 * i ..][0..4], @bitCast(x), .little), + .f64 => std.mem.writeInt(u64, out[8 * i ..][0..8], @bitCast(@as(f64, x) * 1.000000123), .little), + .bf16 => std.mem.writeInt(u16, out[2 * i ..][0..2], f32ToBf16(x), .little), + .f16 => std.mem.writeInt(u16, out[2 * i ..][0..2], @bitCast(@as(f16, @floatCast(x))), .little), + } + } +} + +test "bf16 rounding is nearest even" { + try std.testing.expectEqual(@as(u16, 0x3F80), f32ToBf16(1.0)); + // 1 + 2^-8 is exactly halfway between bf16 1.0 and 1.0078125: rounds to the even one (1.0) + try std.testing.expectEqual(@as(u16, 0x3F80), f32ToBf16(1.00390625)); + try std.testing.expectEqual(@as(u16, 0x3F82), f32ToBf16(1.01171875)); + try std.testing.expect(std.math.isNan(bf16ToF32(f32ToBf16(std.math.nan(f32))))); +} + +test "rank-order reduction of bf16 adds in fp32 and rounds once" { + var a: [4]u8 = undefined; + var b: [4]u8 = undefined; + std.mem.writeInt(u16, a[0..2], f32ToBf16(1.0), .little); + std.mem.writeInt(u16, b[0..2], f32ToBf16(0.00390625), .little); + std.mem.writeInt(u16, a[2..4], f32ToBf16(-2.5), .little); + std.mem.writeInt(u16, b[2..4], f32ToBf16(3.0), .little); + var out: [4]u8 = undefined; + reduce(&out, &.{ &a, &b }, .bf16, .sum); + try std.testing.expectEqual(f32ToBf16(1.0), std.mem.readInt(u16, out[0..2], .little)); + try std.testing.expectEqual(f32ToBf16(0.5), std.mem.readInt(u16, out[2..4], .little)); + reduce(&out, &.{ &a, &b }, .bf16, .max); + try std.testing.expectEqual(f32ToBf16(3.0), std.mem.readInt(u16, out[2..4], .little)); +} + +test "integer sums wrap" { + var a: [4]u8 = undefined; + var b: [4]u8 = undefined; + std.mem.writeInt(i32, &a, std.math.maxInt(i32), .little); + std.mem.writeInt(i32, &b, 1, .little); + var out: [4]u8 = undefined; + reduce(&out, &.{ &a, &b }, .i32, .sum); + try std.testing.expectEqual(@as(i32, std.math.minInt(i32)), std.mem.readInt(i32, &out, .little)); +} diff --git a/zig/src/fabric/tp/roce.zig b/zig/src/fabric/tp/roce.zig new file mode 100644 index 0000000000..58f88e3441 --- /dev/null +++ b/zig/src/fabric/tp/roce.zig @@ -0,0 +1,458 @@ +//! The RoCE backend between two hosts: the tp_roce kernel stages and rings a doorbell, a proxy thread moves the slot. +const std = @import("std"); +const log = @import("log.zig"); +const cuda = @import("cuda"); +const collective = @import("collective.zig"); +const Bootstrap = @import("bootstrap.zig").Bootstrap; +const sock = @import("sock.zig"); +const rv = @import("roce_verbs.zig"); +const HostRegister = @import("mailbox.zig").HostRegister; +const posix = std.posix; +const DevicePtr = collective.DevicePtr; +const DType = collective.DType; +const Error = collective.Error; + +pub const max_blocks = 16; +const threads = 512; +const page = std.heap.page_size_min; +const flag_stride = 128; +// region: ctrl [4 KiB], flags [2 slots][2 HCAs] at 128 B, send [2][slot], recv [2][slot] (mailbox.cu's tp_roce) +const ctrl_off = 0; +const flag_off = 4096; +const send_off = 8192; +// ctrl words (mailbox.cu): doorbell, padded bytes per slot, completed, abort, proxy failure, a failed wait's record +const w_seq = 0; +const w_slot_bytes = 4; +const w_completed = 8; +const w_abort = 12; +const w_proxy_failed = 13; +const w_failed = 16; +/// Model graphs leave sub-millisecond gaps: the proxy stays hot this many idle polls, then naps 20 us at a time. +const idle_spins: u64 = 20_000_000; + +const Mode = enum(u32) { gather = 0, sum_f32 = 1, sum_bf16 = 2, exchange = 3 }; + +/// mailbox.cu's `RoceArgs`, field for field. +const RoceArgs = extern struct { + in: DevicePtr, + out: DevicePtr, + send_slots: DevicePtr, + recv_slots: DevicePtr, + flags: DevicePtr, + ctrl: DevicePtr, + state: DevicePtr, + nbytes: u64, + slot_bytes: u64, + timeout_ns: u64, + rank: u32, + mode: u32, + n_hca: u32, + /// TF_TP_ROCE_FAST (config.roce_fast; 0: the classic path) + opts: u32 = 0, + /// gather only: device int32 [2] of each rank's bytes (0: every shard is `nbytes`) + lens: DevicePtr = 0, +}; + +pub const WireKind = enum { + /// RDMA writes over the HCAs. + verbs, + /// The proxy copies into the peer's region in shared memory: the same protocol on one host without RDMA (tests). + shm, +}; + +pub const Options = struct { + wire: WireKind = .verbs, + hcas: []const rv.HcaSpec = &.{}, + traffic_class: u8 = 0, + cpu: ?u32 = null, + max_bytes: usize = 256 * 1024, + timeout_ns: u64 = 120 * std.time.ns_per_s, + /// TF_TP_ROCE_FAST: tp_roce's `opts` word (config.roce_fast) + fast: u32 = 0, + /// TF_TP_ROCE_BLOCKS: the largest grid (config.roce_blocks) + blocks: u32 = 16, +}; + +/// The shm wire: the peer's region is mapped here; a post copies the slot, then stores the flag with release order. +const ShmWire = struct { + peer: []u8, + local: []u8, + slot_bytes: usize, + + fn post(w: *ShmWire, slot: u32, seq: u32, padded: u32) void { + const recv = send_off + 2 * w.slot_bytes; + @memcpy(w.peer[recv + slot * w.slot_bytes ..][0..padded], w.local[send_off + slot * w.slot_bytes ..][0..padded]); + const flag: *std.atomic.Value(u32) = @ptrCast(@alignCast(w.peer[flag_off + slot * 2 * flag_stride ..].ptr)); + flag.store(seq, .release); + } +}; + +const Wire = union(WireKind) { + verbs: rv.VerbsWire, + shm: ShmWire, + + fn post(w: *Wire, slot: u32, seq: u32, padded: u32) !void { + switch (w.*) { + .verbs => |*v| try v.post(slot, seq, padded), + .shm => |*s| s.post(slot, seq, padded), + } + } + + fn reap(w: *Wire) !void { + switch (w.*) { + .verbs => |*v| try v.reap(), + .shm => {}, + } + } + + fn hcas(w: *const Wire) u32 { + return switch (w.*) { + .verbs => |v| v.n_hca, + .shm => 1, + }; + } +}; + +pub const Roce = struct { + me: u32, + slot_bytes: usize, + timeout_ns: u64, + /// TF_TP_ROCE_FAST: the kernel's `opts` word + fast: u32 = 0, + /// TF_TP_ROCE_BLOCKS: the largest grid + max_grid: u32 = 16, + /// This rank's region (for shm: half of the mapping in `mapping`). + region: []align(page) u8, + mapping: []align(page) u8, + wire: Wire, + cpu: ?u32, + // device side (absent in host-only use: tests and `tf-tp-test roce-host`) + d: ?*const cuda.Driver = null, + hr: ?HostRegister = null, + dev: DevicePtr = 0, + state: ?cuda.DeviceBuffer = null, + module: ?cuda.Module = null, + func: ?cuda.Function = null, + // the proxy thread + thread: ?std.Thread = null, + running: std.atomic.Value(bool) = .init(false), + last_seq: u32 = 0, + err_buf: [160]u8 = undefined, + err_len: usize = 0, + // host emulation of the kernel's side (tests) + host_epoch: u32 = 0, + aborted: std.atomic.Value(bool) = .init(false), + + pub fn regionBytes(slot: usize) usize { + return send_off + 4 * slot; + } + + /// Both ranks call it after the bootstrap; any rank failing makes all return `error.RoceUnavailable` (collective). + pub fn init(d: ?*const cuda.Driver, b: *Bootstrap, image: []const u8, o: Options) !*Roce { + if (b.world != 2) return error.Unsupported; + const slot = std.mem.alignForward(usize, o.max_bytes, 4096); + const len = regionBytes(slot); + const r = try std.heap.page_allocator.create(Roce); + errdefer std.heap.page_allocator.destroy(r); + r.* = .{ .me = b.rank, .slot_bytes = slot, .timeout_ns = o.timeout_ns, .fast = o.fast, .max_grid = o.blocks, .region = undefined, .mapping = undefined, .wire = undefined, .cpu = o.cpu }; + // 1. the regions and the wire, each rank alone; every rank learns whether all did + const ok = r.openWire(b, o, len) catch |e| blk: { + log.warn("[tensorfold] tp roce: rank {d} setup failed: {t}", .{ b.rank, e }); + break :blk false; + }; + if (!try agreeAll(b, ok, "RoCE setup")) { + if (ok) r.closeWire(); + return error.RoceUnavailable; + } + errdefer r.closeWire(); + // 2. the verbs connection records, through the bootstrap + const connected = r.connectWire(b) catch |e| blk: { + log.warn("[tensorfold] tp roce: rank {d} connect failed: {t}", .{ b.rank, e }); + break :blk false; + }; + if (!try agreeAll(b, connected, "RoCE connect")) return error.RoceUnavailable; + // 3. the device side + if (d) |drv| { + const dev_ok = r.openDevice(drv, image) catch |e| blk: { + log.warn("[tensorfold] tp roce: rank {d} device setup failed: {t}", .{ b.rank, e }); + break :blk false; + }; + if (!try agreeAll(b, dev_ok, "RoCE device setup")) { + if (dev_ok) r.closeDevice(); + return error.RoceUnavailable; + } + } + r.running.store(true, .release); + r.thread = try std.Thread.spawn(.{}, proxyLoop, .{r}); + try b.agree("proxy", "RoCE proxy start"); + return r; + } + + fn agreeAll(b: *Bootstrap, ok: bool, _: []const u8) !bool { + var blobs: [2][]const u8 = undefined; + var buf: [2]u8 = undefined; + const all = try b.allGather(&.{@intFromBool(ok)}, &buf, &blobs); + return all[0][0] == 1 and all[1][0] == 1; + } + + fn openWire(r: *Roce, b: *Bootstrap, o: Options, len: usize) !bool { + switch (o.wire) { + .verbs => { + r.mapping = try posix.mmap(null, len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .PRIVATE, .ANONYMOUS = true }, -1, 0); + r.region = r.mapping; + errdefer posix.munmap(r.mapping); + r.wire = .{ .verbs = try rv.VerbsWire.open(o.hcas, r.region, .{ .send = send_off, .recv = send_off + 2 * r.slot_bytes, .flags = flag_off, .slot = r.slot_bytes }, o.traffic_class) }; + }, + .shm => { + r.mapping = try sharedPair(b, len); + r.region = @alignCast(r.mapping[r.me * len ..][0..len]); + r.wire = .{ .shm = .{ .peer = r.mapping[(1 - r.me) * len ..][0..len], .local = r.region, .slot_bytes = r.slot_bytes } }; + }, + } + return true; + } + + fn connectWire(r: *Roce, b: *Bootstrap) !bool { + switch (r.wire) { + .shm => return true, + .verbs => |*v| { + const mine = v.record(); + var blobs: [2][]const u8 = undefined; + var buf: [2 * @sizeOf(rv.Record)]u8 = undefined; + const all = try b.allGather(std.mem.asBytes(&mine), &buf, &blobs); + var peer: rv.Record = undefined; + if (all[1 - r.me].len != @sizeOf(rv.Record)) return error.BadRecord; + @memcpy(std.mem.asBytes(&peer), all[1 - r.me]); + try v.connect(peer); + return true; + }, + } + } + + fn openDevice(r: *Roce, d: *const cuda.Driver, image: []const u8) !bool { + if (image.len == 0) return error.NoKernel; + r.d = d; + r.hr = try HostRegister.open(); + try d.check(r.hr.?.register(r.region.ptr, r.region.len, 1 | 2), "cuMemHostRegister"); // PORTABLE | DEVICEMAP + try d.check(d.api.cuMemHostGetDevicePointer_v2(&r.dev, r.region.ptr, 0), "cuMemHostGetDevicePointer"); + r.state = try cuda.DeviceBuffer.alloc(d, 16); + try r.state.?.fill8(0, null); + r.module = try cuda.Module.load(d, image); + r.func = try r.module.?.function("tp_roce"); + return true; + } + + fn closeDevice(r: *Roce) void { + if (r.module) |*m| m.unload(); + if (r.state) |*s| s.free(); + if (r.hr) |*h| { + _ = h.unregister(r.region.ptr); + h.lib.close(); + } + r.module = null; + r.state = null; + r.hr = null; + } + + fn closeWire(r: *Roce) void { + switch (r.wire) { + .verbs => |*v| v.close(), + .shm => {}, + } + posix.munmap(r.mapping); + } + + pub fn deinit(r: *Roce) void { + r.running.store(false, .release); + if (r.thread) |t| t.join(); + r.closeDevice(); + r.closeWire(); + std.heap.page_allocator.destroy(r); + } + + fn word(r: *const Roce, i: usize) *std.atomic.Value(u32) { + return @ptrCast(@alignCast(r.region[ctrl_off + 4 * i ..].ptr)); + } + + fn fail(r: *Roce, comptime fmt: []const u8, args: anytype) void { + const s = std.fmt.bufPrint(&r.err_buf, fmt, args) catch "proxy failure"; + r.err_len = s.len; + log.err("[tensorfold] tp roce proxy: {s}", .{s}); + r.word(w_proxy_failed).store(1, .release); + r.word(w_abort).store(1, .release); // waits on this rank end now instead of at the timeout + } + + /// The proxy loop: spin on the doorbell, post each new sequence (at most two pending), reap while idle, then nap. + fn proxyLoop(r: *Roce) void { + if (r.cpu) |cpu| pin(cpu); + var idle: u64 = 0; + while (r.running.load(.monotonic)) { + if (r.word(w_abort).load(.monotonic) != 0) return; + const seq = r.word(w_seq).load(.acquire); + if (seq == r.last_seq) { + idle += 1; + if (idle % 64 == 0) r.wire.reap() catch |e| return r.fail("completion: {t}", .{e}); + if (idle >= idle_spins) sock.sleepNs(20 * std.time.ns_per_us) else std.atomic.spinLoopHint(); + continue; + } + idle = 0; + var pending = seq -% r.last_seq; + if (pending > 2) return r.fail("doorbell skipped {d} ops (posted {d}, now {d})", .{ pending, r.last_seq, seq }); + var s = r.last_seq +% 1; + while (pending > 0) : ({ + s +%= 1; + pending -= 1; + }) { + const padded = r.word(w_slot_bytes + (s & 1)).load(.acquire); + if (padded == 0 or padded % 16 != 0 or padded > r.slot_bytes) return r.fail("op {d} of {d} bytes", .{ s, padded }); + r.wire.post(s & 1, s, padded) catch |e| return r.fail("post of op {d}: {t}", .{ s, e }); + r.last_seq = s; + } + r.wire.reap() catch |e| return r.fail("completion: {t}", .{e}); + } + } + + pub fn fits(r: *const Roce, n: usize) bool { + return n <= r.slot_bytes; + } + + fn launch(r: *Roce, mode: Mode, in: DevicePtr, out: DevicePtr, n: usize, lens: DevicePtr, stream: collective.Stream) Error!void { + if (r.aborted.load(.acquire)) return error.Aborted; + if (n > r.slot_bytes) return error.Unsupported; + if (n == 0) return; + const d = r.d orelse return error.Unsupported; + var a: RoceArgs = .{ + .in = in, + .out = out, + .send_slots = r.dev + send_off, + .recv_slots = r.dev + send_off + 2 * r.slot_bytes, + .flags = r.dev + flag_off, + .ctrl = r.dev + ctrl_off, + .state = r.state.?.ptr, + .nbytes = n, + .slot_bytes = r.slot_bytes, + .timeout_ns = r.timeout_ns, + .rank = r.me, + .mode = @backingInt(mode), + .n_hca = r.wire.hcas(), + .opts = r.fast, + .lens = lens, + }; + var params = [1]?*anyopaque{@ptrCast(&a)}; + const blocks = @min(@import("mailbox.zig").blocksFor(n), r.max_grid); + d.check(d.api.cuLaunchKernel(r.func.?.handle, blocks, 1, 1, threads, 1, 1, 0, stream, ¶ms, null), "cuLaunchKernel tp_roce") catch return error.BackendFailed; + } + + pub fn allGatherBytes(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return r.launch(.gather, s, o, n, 0, stream); + } + + /// `allGatherBytes` where rank r posts only lens[r] bytes (device int32 [2], padded to 16, at least 16). + pub fn allGatherV(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, lens: DevicePtr, stream: collective.Stream) Error!void { + if (lens == 0) return error.Invalid; + return r.launch(.gather, s, o, n, lens, stream); + } + + pub fn exchangeBytes(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, stream: collective.Stream) Error!void { + return r.launch(.exchange, s, o, n, 0, stream); + } + + pub fn sum(r: *Roce, s: DevicePtr, o: DevicePtr, count: usize, t: DType, stream: collective.Stream) Error!void { + const mode: Mode = switch (t) { + .f32 => .sum_f32, + .bf16 => .sum_bf16, + else => return error.Unsupported, + }; + return r.launch(mode, s, o, count * t.size(), 0, stream); + } + + /// Ends this rank's waits and the proxy (the peer learns through the fate channel and aborts its own). + pub fn abort(r: *Roce) void { + r.aborted.store(true, .release); + r.word(w_abort).store(1, .release); + } + + pub fn check(r: *Roce) Error!void { + if (r.word(w_failed).load(.acquire) != 0) { + log.warn("tp roce: rank {d}'s wait for seq {d} failed ({s} after {d} us; HCA {d}, flag held {d}); doorbell {d}, posted {d}, completed {d}, abort {d}, proxy failed {d}, flags {d} {d}", .{ r.me, r.word(w_failed + 1).load(.acquire), if (r.word(w_failed + 4).load(.acquire) == 1) "aborted" else "timed out", r.word(w_failed + 5).load(.acquire), r.word(w_failed + 3).load(.acquire), r.word(w_failed + 2).load(.acquire), r.word(w_seq).load(.acquire), r.last_seq, r.word(w_completed).load(.acquire), r.word(w_abort).load(.acquire), r.word(w_proxy_failed).load(.acquire), r.flagWord(0), r.flagWord(1) }); + return error.PeerFailed; + } + if (r.word(w_proxy_failed).load(.acquire) != 0) return error.BackendFailed; + if (r.aborted.load(.acquire) or r.word(w_abort).load(.acquire) != 0) return error.Aborted; + } + + fn flagWord(r: *const Roce, slot: usize) u32 { + const p: *std.atomic.Value(u32) = @ptrCast(@alignCast(r.region[flag_off + slot * 2 * flag_stride ..].ptr)); + return p.load(.acquire); + } + + pub fn proxyError(r: *const Roce) []const u8 { + return r.err_buf[0..r.err_len]; + } + + /// The kernel's side of one op done by the host (tests, `roce-host`): a gather through the same slots and flags. + pub fn hostGather(r: *Roce, in: []const u8, out: []u8) Error!void { + return r.hostGatherV(in, out, null); + } + + /// `hostGather` with variable lengths (`allGatherV`): rank k's first lens[k] bytes land at out + k x n. + pub fn hostGatherV(r: *Roce, in: []const u8, out: []u8, lens: ?[2]u32) Error!void { + const n = in.len; + if (n == 0 or n > r.slot_bytes or out.len < 2 * n) return error.Invalid; + if (lens) |l| if (l[0] > n or l[1] > n) return error.Invalid; + const mine: usize = if (lens) |l| l[r.me] else n; + const theirs: usize = if (lens) |l| l[1 - r.me] else n; + const seq = r.host_epoch +% 1; + const slot: usize = seq & 1; + @memcpy(r.region[send_off + slot * r.slot_bytes ..][0..mine], in[0..mine]); + r.word(w_slot_bytes + slot).store(@intCast(@max(16, std.mem.alignForward(usize, mine, 16))), .release); + r.word(w_seq).store(seq, .release); + @memcpy(out[r.me * n ..][0..mine], in[0..mine]); + const t0 = sock.nowNs(); + for (0..r.wire.hcas()) |h| { + const flag: *std.atomic.Value(u32) = @ptrCast(@alignCast(r.region[flag_off + (slot * 2 + h) * flag_stride ..].ptr)); + while (flag.load(.acquire) != seq) { + if (r.word(w_abort).load(.monotonic) != 0) return error.Aborted; + if (sock.nowNs() - t0 > r.timeout_ns) return error.Timeout; + std.atomic.spinLoopHint(); + } + } + @memcpy(out[(1 - r.me) * n ..][0..theirs], r.region[send_off + 2 * r.slot_bytes + slot * r.slot_bytes ..][0..theirs]); + r.host_epoch = seq; + } +}; + +/// One shared mapping holding both ranks' regions (rank 0 creates and sizes the file before anyone learns its name). +fn sharedPair(b: *Bootstrap, len: usize) ![]align(page) u8 { + var name: [64]u8 = undefined; + var name_len: u8 = 0; + var fd: std.c.fd_t = -1; + defer if (fd >= 0) { + _ = std.c.close(fd); + }; + if (b.rank == 0) blk: { + const s = std.fmt.bufPrint(&name, "/dev/shm/tf-tp-roce-{d}-{x}", .{ std.c.getpid(), sock.nowNs() }) catch break :blk; + name[s.len] = 0; + fd = std.c.open(name[0..s.len :0], .{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true, .CLOEXEC = true }, @as(std.c.mode_t, 0o600)); + if (fd >= 0 and std.c.ftruncate(fd, @intCast(2 * len)) == 0) name_len = @intCast(s.len); + } + // an empty name tells the other rank that rank 0 failed (it must not wait for a name that never comes) + try b.broadcast(std.mem.asBytes(&name_len), std.mem.asBytes(&name_len)); + if (name_len == 0) return error.SharedMemoryUnavailable; + try b.broadcast(name[0..name_len], name[0..name_len]); + name[name_len] = 0; + if (b.rank != 0) { + fd = std.c.open(name[0..name_len :0], .{ .ACCMODE = .RDWR, .CLOEXEC = true }, @as(std.c.mode_t, 0)); + if (fd < 0) return error.SharedMemoryUnavailable; + } + const m = try posix.mmap(null, 2 * len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); + try b.agree("mapped", "RoCE shm wire"); + if (b.rank == 0) _ = std.c.unlink(name[0..name_len :0]); + return m; +} + +fn pin(cpu: u32) void { + var set: std.os.linux.cpu_set_t = @splat(0); + set[cpu / @bitSizeOf(usize)] |= @as(usize, 1) << @intCast(cpu % @bitSizeOf(usize)); + std.os.linux.sched_setaffinity(0, &set) catch log.warn("tp roce: could not pin the proxy to CPU {d}", .{cpu}); +} diff --git a/zig/src/fabric/tp/roce_host.zig b/zig/src/fabric/tp/roce_host.zig new file mode 100644 index 0000000000..afb705d41b --- /dev/null +++ b/zig/src/fabric/tp/roce_host.zig @@ -0,0 +1,58 @@ +//! `tf-tp-test roce-host`: the RoCE wire and proxy between two hosts without a GPU, checked bit for bit, then timed. +const std = @import("std"); +const tp = @import("tp"); +const roce = tp.roce; +const rv = tp.roce_verbs; + +const sizes = [_]usize{ 16, 4096, 8192, 16384, 65536, 262144 }; + +pub fn run(cfg: tp.Config, gpa: std.mem.Allocator) !u8 { + var b = try tp.Bootstrap.open(cfg); + defer b.close(); + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var n: usize = 0; + if (cfg.roce_hca) |text| n = try rv.parseSpecs(text, &specs) else { + var v = try tp.Verbs.open(null); + defer v.close(); + n = rv.autoSpecs(&v, &specs, cfg.roce_hcas); + } + const r = try roce.Roce.init(null, &b, "", .{ .wire = if (cfg.roce_wire == .shm) .shm else .verbs, .hcas = specs[0..n], .traffic_class = cfg.roce_tc, .cpu = cfg.roce_cpu, .max_bytes = 256 * 1024, .timeout_ns = 10 * std.time.ns_per_s }); + defer r.deinit(); + const in = try gpa.alloc(u8, 256 * 1024); + defer gpa.free(in); + const peer = try gpa.alloc(u8, 256 * 1024); + defer gpa.free(peer); + const out = try gpa.alloc(u8, 512 * 1024); + defer gpa.free(out); + // 1. every size from 1 byte to the slot, both slots, bit for bit + var checked: u32 = 0; + for (0..1000) |k| { + const len = 1 + (k * 7919) % (256 * 1024); + tp.reference.fill(in[0..len], .u8, cfg.rank, k); + tp.reference.fill(peer[0..len], .u8, 1 - cfg.rank, k); + try r.hostGather(in[0..len], out[0 .. 2 * len]); + const mine_at = cfg.rank * len; + const peer_at = (1 - cfg.rank) * len; + if (!std.mem.eql(u8, out[mine_at..][0..len], in[0..len]) or !std.mem.eql(u8, out[peer_at..][0..len], peer[0..len])) { + std.debug.print("FAIL roce-host gather {d} of {d} bytes\n", .{ k, len }); + return 1; + } + checked += 1; + } + try r.check(); + std.debug.print("PASS roce-host: {d} gathers of 1 B to 256 KiB bit-exact over {s}\n", .{ checked, @tagName(cfg.roce_wire) }); + // 2. one exchange's host-to-host time by size: median and p90 of 2,000 back-to-back ops + var times: [2000]u64 = undefined; + for (sizes) |len| { + for (×) |*t| { + const t0 = tp.sock.nowNs(); + try r.hostGather(in[0..len], out[0 .. 2 * len]); + t.* = tp.sock.nowNs() - t0; + } + std.mem.sort(u64, ×, {}, std.sort.asc(u64)); + if (cfg.rank == 0) std.debug.print("RESULT {{\"backend\":\"roce-host\",\"wire\":\"{s}\",\"bytes\":{d},\"p50_us\":{d:.2},\"p90_us\":{d:.2}}}\n", .{ @tagName(cfg.roce_wire), len, @as(f64, @floatFromInt(times[1000])) / 1e3, @as(f64, @floatFromInt(times[1800])) / 1e3 }); + } + try r.check(); + try b.agree("done", "roce-host end"); + return 0; +} diff --git a/zig/src/fabric/tp/roce_test.zig b/zig/src/fabric/tp/roce_test.zig new file mode 100644 index 0000000000..51e9d79133 --- /dev/null +++ b/zig/src/fabric/tp/roce_test.zig @@ -0,0 +1,116 @@ +//! The RoCE protocol end to end on host threads: two ranks, TCP bootstrap, proxy, shm or verbs (TF_TP_ROCE_TEST_HCA). +const std = @import("std"); +const roce = @import("roce.zig"); +const rv = @import("roce_verbs.zig"); +const bootstrap = @import("bootstrap.zig"); +const sock = @import("sock.zig"); +const reference = @import("reference.zig"); +const Bootstrap = bootstrap.Bootstrap; + +const Run = struct { + port: u16, + wire: roce.WireKind, + devs: [2][]const u8 = .{ "", "" }, + ops: u32, + timeout_ns: u64 = 10 * std.time.ns_per_s, + /// Rank 1 stops after this many ops (rank 0 must then fail its next wait, by timeout or abort). + stop_after: ?u32 = null, + abort_after_ms: ?u64 = null, + errs: [2]?anyerror = .{ null, null }, + done: [2]u32 = .{ 0, 0 }, + max_bytes: usize = 16 * 1024, + /// Variable lengths (`Roce.allGatherV`'s contract): each op's lens[r] bytes of rank r's n, 0 and n included. + var_lens: bool = false, +}; + +fn rank(run: *Run, me: u32, lfd: sock.Fd) void { + run.errs[me] = body(run, me, lfd); +} + +fn body(run: *Run, me: u32, lfd: sock.Fd) ?anyerror { + var b = (if (me == 0) Bootstrap.acceptOn(lfd, 2, 5 * std.time.ns_per_s) else Bootstrap.open(.{ .world = 2, .rank = 1, .link_port = run.port, .connect_timeout_ns = 5 * std.time.ns_per_s })) catch |e| return e; + defer b.close(); + var spec: [1]rv.HcaSpec = .{.{}}; + if (run.wire == .verbs) _ = rv.parseSpecs(run.devs[me], @ptrCast(&spec)) catch |e| return e; + const r = roce.Roce.init(null, &b, "", .{ .wire = run.wire, .hcas = if (run.wire == .verbs) &spec else &.{}, .max_bytes = run.max_bytes, .timeout_ns = run.timeout_ns }) catch |e| return e; + defer r.deinit(); + var abort_thread: ?std.Thread = null; + defer if (abort_thread) |t| t.join(); + var in: [64 * 1024]u8 = undefined; + var peer: [64 * 1024]u8 = undefined; + var out: [128 * 1024]u8 = undefined; + for (0..run.ops) |k| { + if (run.stop_after) |s| if (k == s) { + // rank 1 skips op s; rank 0 waits in it until its timeout, or until the abort armed now + if (me == 1) { + sock.sleepNs(300 * std.time.ns_per_ms); + return null; + } + if (run.abort_after_ms) |ms| abort_thread = std.Thread.spawn(.{}, abortLater, .{ r, ms }) catch null; + }; + // sizes from 1 byte to the slot, odd ones included, both slots and wrap-arounds + const n = 1 + (k * 7919) % run.max_bytes; + reference.fill(in[0..n], .u8, me, k); + reference.fill(peer[0..n], .u8, 1 - me, k); + const lens: [2]u32 = if (k % 5 == 0) .{ 0, @intCast(n) } else .{ @intCast((k * 31) % (n + 1)), @intCast((k * 17 + 5) % (n + 1)) }; + const l: [2]usize = if (run.var_lens) .{ lens[0], lens[1] } else .{ n, n }; + r.hostGatherV(in[0..n], out[0 .. 2 * n], if (run.var_lens) lens else null) catch |e| return e; + const r0 = if (me == 0) in[0..n] else peer[0..n]; + const r1 = if (me == 1) in[0..n] else peer[0..n]; + if (!std.mem.eql(u8, out[0..l[0]], r0[0..l[0]]) or !std.mem.eql(u8, out[n..][0..l[1]], r1[0..l[1]])) return error.WrongBytes; + run.done[me] += 1; + } + r.check() catch |e| return e; + // both ranks finish before either tears down (the peer may still read a flag this rank's wire writes) + b.agree("done", "test end") catch |e| return e; + return null; +} + +fn abortLater(r: *roce.Roce, ms: u64) void { + sock.sleepNs(ms * std.time.ns_per_ms); + r.abort(); +} + +fn twoRanks(run: *Run) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + run.port = l.port; + const t1 = try std.Thread.spawn(.{}, rank, .{ run, 1, l.fd }); + rank(run, 0, l.fd); + t1.join(); +} + +test "the RoCE protocol over the shm wire: 2,000 gathers of every size, bit for bit" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000 }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); +} + +test "the RoCE protocol with variable lengths over the shm wire: each rank posts its own padded count, empty shards included" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000, .var_lens = true }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); +} + +test "a peer that stops makes the waiting rank time out, and abort ends a wait at once" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 50, .stop_after = 20, .timeout_ns = 100 * std.time.ns_per_ms }; + try twoRanks(&run); + try std.testing.expectEqual(@as(?anyerror, error.Timeout), run.errs[0]); + try std.testing.expectEqual(@as(u32, 20), run.done[0]); + var run2: Run = .{ .port = 0, .wire = .shm, .ops = 50, .stop_after = 20, .abort_after_ms = 50 }; + try twoRanks(&run2); + try std.testing.expectEqual(@as(?anyerror, error.Aborted), run2.errs[0]); + try std.testing.expectEqual(@as(?anyerror, null), run2.errs[1]); +} + +test "the RoCE protocol over verbs (TF_TP_ROCE_TEST_HCA=dev0,dev1: soft-RoCE or real HCAs)" { + const raw = std.c.getenv("TF_TP_ROCE_TEST_HCA") orelse return error.SkipZigTest; + var it = std.mem.splitScalar(u8, std.mem.span(raw), ','); + var run: Run = .{ .port = 0, .wire = .verbs, .ops = 2000 }; + run.devs = .{ it.next() orelse return error.SkipZigTest, it.next() orelse return error.SkipZigTest }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); +} diff --git a/zig/src/fabric/tp/roce_verbs.zig b/zig/src/fabric/tp/roce_verbs.zig new file mode 100644 index 0000000000..bc99681b5e --- /dev/null +++ b/zig/src/fabric/tp/roce_verbs.zig @@ -0,0 +1,297 @@ +//! The RoCE wire over verbs: one RC queue pair a HCA to the peer, a send slot striped over the HCAs with seq flags. +const std = @import("std"); +const log = @import("log.zig"); +const verbs = @import("verbs"); +const abi = verbs.abi; +const Verbs = verbs.Verbs; + +pub const max_hcas = 2; +/// Work requests a queue pair holds; at most a quarter outstanding (two requests a stripe). +const send_depth = 256; +const flag_stride = 128; +const link_ethernet = 2; +const record_magic: u32 = 0x5446_5244 + 1; // "TFRD", bumped once for this record's layout + +pub const Error = error{ NoHca, VerbsFailed, PostFailed, Completion, BadRecord }; + +/// One HCA as configured: its device name and RoCE v2 GID index (null: detected). +pub const HcaSpec = struct { name: [64]u8 = @splat(0), name_len: u8 = 0, gid: ?u8 = null }; + +/// What a rank tells its peer to connect (both ranks build it the same way; sent through the bootstrap). +pub const Record = extern struct { + magic: u32 = record_magic, + n_hca: u32 = 0, + region: u64 = 0, + rkey: [max_hcas]u32 = @splat(0), + lid: [max_hcas]u32 = @splat(0), + mtu: [max_hcas]u32 = @splat(0), + qpn: [max_hcas]u32 = @splat(0), + gid: [max_hcas][16]u8 = @splat(@splat(0)), +}; + +const Hca = struct { + ctx: *abi.Context, + pd: *abi.Pd, + mr: *abi.Mr, + cq: *abi.Cq, + qp: *abi.Qp, + gid_index: u8, + gid: [16]u8, + mtu: u32, + lid: u16, + outstanding: u32 = 0, +}; + +/// Parses `name[:gid],name[:gid]` (TF_TP_ROCE_HCA). +pub fn parseSpecs(text: []const u8, out: *[max_hcas]HcaSpec) !usize { + var n: usize = 0; + var it = std.mem.tokenizeScalar(u8, text, ','); + while (it.next()) |item| { + if (n == max_hcas) return error.TooManyHcas; + var parts = std.mem.splitScalar(u8, std.mem.trim(u8, item, " "), ':'); + const name = parts.next().?; + if (name.len == 0 or name.len >= 64) return error.BadHcaName; + out[n] = .{ .name_len = @intCast(name.len) }; + @memcpy(out[n].name[0..name.len], name); + if (parts.next()) |g| out[n].gid = try std.fmt.parseInt(u8, g, 10); + n += 1; + } + return n; +} + +/// RoCE v2 by sysfs (the kernel's own word for each GID); null when sysfs cannot say. +fn sysfsRoceV2(name: []const u8, index: u32) ?bool { + var path: [256]u8 = undefined; + const p = std.fmt.bufPrintSentinel(&path, "/sys/class/infiniband/{s}/ports/1/gid_attrs/types/{d}", .{ name, index }, 0) catch return null; + const fd = std.c.open(p, .{ .ACCMODE = .RDONLY }); + if (fd < 0) return null; + defer _ = std.c.close(fd); + var buf: [32]u8 = undefined; + const n = std.c.read(fd, &buf, buf.len); + if (n <= 0) return false; + return std.mem.startsWith(u8, buf[0..@intCast(n)], "RoCE v2"); +} + +/// An IPv4-mapped GID (::ffff:a.b.c.d): the address the HCA's netdev carries. +pub fn ipv4Mapped(g: [16]u8) bool { + return std.mem.allEqual(u8, g[0..10], 0) and g[10] == 0xff and g[11] == 0xff; +} + +/// The lowest RoCE v2 GID index that is IPv4-mapped (GID indices move on reboot: never hard-coded). +fn detectGid(v: *const Verbs, ctx: *abi.Context, name: []const u8, table: c_int) ?u8 { + var i: u32 = 0; + while (i < @min(@as(u32, @intCast(@max(table, 0))), 255)) : (i += 1) { + const g = v.gid(ctx, 1, i) catch continue; + if (!ipv4Mapped(g.raw)) continue; + const v2 = sysfsRoceV2(name, i) orelse (if (v.gidType(ctx, 1, i)) |t| t == 2 else true); + if (v2) return @intCast(i); + } + return null; +} + +/// Every device whose port 1 is an active Ethernet port with an IPv4 RoCE v2 GID, by name order, at most `limit`. +pub fn autoSpecs(v: *const Verbs, out: *[max_hcas]HcaSpec, limit: usize) usize { + var count: c_int = 0; + const list = v.api.ibv_get_device_list(&count) orelse return 0; + defer v.api.ibv_free_device_list(list); + var n: usize = 0; + var i: usize = 0; + while (list[i]) |dev| : (i += 1) { + if (n == limit) break; + const name = std.mem.span(v.api.ibv_get_device_name(dev) orelse continue); + const ctx = v.api.ibv_open_device(dev) orelse continue; + defer _ = v.api.ibv_close_device(ctx); + const port = v.port(ctx, 1) catch continue; + if (port.state != abi.port_active or port.link_layer != link_ethernet) continue; + const g = detectGid(v, ctx, name, port.gid_tbl_len) orelse continue; + out[n] = .{ .name_len = @intCast(@min(name.len, 63)), .gid = g }; + @memcpy(out[n].name[0..out[n].name_len], name[0..out[n].name_len]); + n += 1; + } + return n; +} + +pub const VerbsWire = struct { + v: Verbs, + hcas: [max_hcas]Hca = undefined, + n_hca: u32 = 0, + traffic_class: u8, + peer_region: u64 = 0, + peer_rkey: [max_hcas]u32 = @splat(0), + /// Offsets in both ranks' regions (the same layout): receive slots and flags. + recv_off: usize, + flag_off: usize, + send_off: usize, + slot_bytes: usize, + region: []u8, + posted: u64 = 0, + completed: u64 = 0, + + /// Opens the HCAs, registers `region` on each and creates the queue pairs (INIT); `record` then goes to the peer. + pub fn open(specs: []const HcaSpec, region: []u8, layout: struct { send: usize, recv: usize, flags: usize, slot: usize }, traffic_class: u8) !VerbsWire { + if (specs.len == 0 or specs.len > max_hcas) return error.NoHca; + var w: VerbsWire = .{ .v = try Verbs.open(null), .traffic_class = traffic_class, .recv_off = layout.recv, .flag_off = layout.flags, .send_off = layout.send, .slot_bytes = layout.slot, .region = region }; + errdefer w.close(); + for (specs) |s| { + try w.openHca(s); + } + return w; + } + + fn openHca(w: *VerbsWire, s: HcaSpec) !void { + const api = w.v.api; + const name = s.name[0..s.name_len]; + const ctx = try w.v.openDevice(name); + errdefer _ = api.ibv_close_device(ctx); + const port = try w.v.port(ctx, 1); + if (port.state != abi.port_active) { + log.warn("tp roce: {s} port 1 is not active", .{name}); + return error.NoHca; + } + const gi = s.gid orelse detectGid(&w.v, ctx, name, port.gid_tbl_len) orelse { + log.warn("tp roce: {s} has no IPv4 RoCE v2 GID", .{name}); + return error.NoHca; + }; + const gid = try w.v.gid(ctx, 1, gi); + const pd = api.ibv_alloc_pd(ctx) orelse return error.VerbsFailed; + errdefer _ = api.ibv_dealloc_pd(pd); + const mr = api.ibv_reg_mr(pd, w.region.ptr, w.region.len, abi.access_local_write | abi.access_remote_write) orelse { + log.warn("tp roce: ibv_reg_mr on {s} failed (memlock limit? ulimit -l unlimited)", .{name}); + return error.VerbsFailed; + }; + errdefer _ = api.ibv_dereg_mr(mr); + const cq = api.ibv_create_cq(ctx, send_depth, null, null, 0) orelse return error.VerbsFailed; + errdefer _ = api.ibv_destroy_cq(cq); + var init: abi.QpInitAttr = .{ .send_cq = cq, .recv_cq = cq, .cap = .{ .max_send_wr = send_depth, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 16 }, .qp_type = abi.qpt_rc }; + const qp = api.ibv_create_qp(pd, &init) orelse return error.VerbsFailed; + errdefer _ = api.ibv_destroy_qp(qp); + var a: abi.QpAttr = .{ .qp_state = abi.qps_init, .port_num = 1, .qp_access_flags = abi.access_remote_write }; + if (api.ibv_modify_qp(qp, &a, abi.mask.state | abi.mask.pkey_index | abi.mask.port | abi.mask.access_flags) != 0) return error.VerbsFailed; + w.hcas[w.n_hca] = .{ .ctx = ctx, .pd = pd, .mr = mr, .cq = cq, .qp = qp, .gid_index = gi, .gid = gid.raw, .mtu = @intCast(port.active_mtu), .lid = port.lid }; + w.n_hca += 1; + log.info("[tensorfold] tp roce: {s} GID {d} ({d}.{d}.{d}.{d}), MTU enum {d}", .{ name, gi, gid.raw[12], gid.raw[13], gid.raw[14], gid.raw[15], port.active_mtu }); + } + + pub fn close(w: *VerbsWire) void { + const api = w.v.api; + for (w.hcas[0..w.n_hca]) |h| { + _ = api.ibv_destroy_qp(h.qp); + _ = api.ibv_destroy_cq(h.cq); + _ = api.ibv_dereg_mr(h.mr); + _ = api.ibv_dealloc_pd(h.pd); + _ = api.ibv_close_device(h.ctx); + } + w.n_hca = 0; + w.v.close(); + } + + pub fn record(w: *const VerbsWire) Record { + var r: Record = .{ .n_hca = w.n_hca, .region = @intFromPtr(w.region.ptr) }; + for (w.hcas[0..w.n_hca], 0..) |h, i| { + r.rkey[i] = h.mr.rkey; + r.lid[i] = h.lid; + r.mtu[i] = h.mtu; + r.qpn[i] = h.qp.qp_num; + r.gid[i] = h.gid; + } + return r; + } + + /// The peer HCA for each of ours: the same IPv4 /24 (functions pair by subnet, names may differ), else position. + pub fn pairing(mine: Record, peer: Record) ![max_hcas]u32 { + if (peer.magic != record_magic) return error.BadRecord; + if (peer.n_hca != mine.n_hca) return error.BadRecord; + var out: [max_hcas]u32 = undefined; + for (0..mine.n_hca) |i| { + out[i] = @intCast(i); + var hits: u32 = 0; + for (0..peer.n_hca) |j| if (std.mem.eql(u8, mine.gid[i][12..15], peer.gid[j][12..15])) { + out[i] = @intCast(j); + hits += 1; + }; + if (hits > 1) out[i] = @intCast(i); // one subnet for both: keep the order + } + return out; + } + + /// INIT to RTR to RTS against the peer's record. + pub fn connect(w: *VerbsWire, peer: Record) !void { + const pairs = try pairing(w.record(), peer); + const api = w.v.api; + const m = abi.mask; + w.peer_region = peer.region; + for (w.hcas[0..w.n_hca], 0..) |*h, i| { + const j = pairs[i]; + w.peer_rkey[i] = peer.rkey[j]; + var rtr: abi.QpAttr = .{ .qp_state = abi.qps_rtr, .path_mtu = @intCast(@min(peer.mtu[j], h.mtu)), .dest_qp_num = peer.qpn[j], .rq_psn = 0, .max_dest_rd_atomic = 1, .min_rnr_timer = 12 }; + rtr.ah_attr = .{ .grh = .{ .dgid = .{ .raw = peer.gid[j] }, .sgid_index = h.gid_index, .hop_limit = 64, .traffic_class = w.traffic_class }, .dlid = @intCast(peer.lid[j]), .is_global = 1, .port_num = 1 }; + if (api.ibv_modify_qp(h.qp, &rtr, m.state | m.av | m.path_mtu | m.dest_qpn | m.rq_psn | m.max_dest_rd_atomic | m.min_rnr_timer) != 0) return error.VerbsFailed; + var rts: abi.QpAttr = .{ .qp_state = abi.qps_rts, .timeout = 14, .retry_cnt = 7, .rnr_retry = 7, .sq_psn = 0, .max_rd_atomic = 1 }; + if (api.ibv_modify_qp(h.qp, &rts, m.state | m.timeout | m.retry_cnt | m.rnr_retry | m.sq_psn | m.max_qp_rd_atomic) != 0) return error.VerbsFailed; + } + } + + /// Slot `slot`'s first `padded` bytes to the peer, striped over the HCAs, each stripe's `seq` flag on the same QP. + pub fn post(w: *VerbsWire, slot: u32, seq: u32, padded: u32) !void { + const packs = padded / 16; + var offset: u64 = 0; + var seq_copy = seq; + for (w.hcas[0..w.n_hca], 0..) |*h, i| { + while (h.outstanding >= send_depth / 4) try w.reapHca(h); + var stripe = packs / w.n_hca; + if (i < packs % w.n_hca) stripe += 1; + const bytes: u32 = stripe * 16; + var flag_sge: abi.Sge = .{ .addr = @intFromPtr(&seq_copy), .length = 4, .lkey = 0 }; + var flag_wr: abi.SendWr = .{ .wr_id = seq, .sg_list = @ptrCast(&flag_sge), .num_sge = 1, .opcode = abi.wr_rdma_write, .send_flags = abi.send_signaled | abi.send_inline, .remote_addr = w.peer_region + w.flag_off + (@as(u64, slot) * max_hcas + i) * flag_stride, .rkey = w.peer_rkey[i] }; + var data_sge: abi.Sge = .{ .addr = @intFromPtr(w.region.ptr) + w.send_off + slot * w.slot_bytes + offset, .length = bytes, .lkey = h.mr.lkey }; + var data_wr: abi.SendWr = .{ .wr_id = seq, .next = &flag_wr, .sg_list = @ptrCast(&data_sge), .num_sge = 1, .opcode = abi.wr_rdma_write, .send_flags = 0, .remote_addr = w.peer_region + w.recv_off + slot * w.slot_bytes + offset, .rkey = w.peer_rkey[i] }; + var bad: ?*abi.SendWr = null; + const first: *abi.SendWr = if (bytes > 0) &data_wr else &flag_wr; + if (h.ctx.ops.post_send(h.qp, first, &bad) != 0) return error.PostFailed; + h.outstanding += 1; + offset += bytes; + } + w.posted += 1; + } + + fn reapHca(w: *VerbsWire, h: *Hca) !void { + var wc: [32]abi.Wc = undefined; + const n = h.ctx.ops.poll_cq(h.cq, 32, &wc); + if (n < 0) return error.Completion; + for (wc[0..@intCast(n)]) |c| if (c.status != abi.wc_success) { + log.err("tp roce: RDMA write of seq {d} failed: {s} (vendor 0x{x})", .{ c.wr_id, w.v.statusText(c.status), c.vendor_err }); + return error.Completion; + }; + h.outstanding -= @intCast(n); + w.completed += @intCast(n); + } + + /// Reap every HCA's completions (after each op, and in the proxy's idle loop). + pub fn reap(w: *VerbsWire) !void { + for (w.hcas[0..w.n_hca]) |*h| try w.reapHca(h); + } +}; + +test "HCA lists parse and subnets pair the functions" { + var specs: [max_hcas]HcaSpec = undefined; + const n = try parseSpecs("rocep1s0f1:3, roceP2p1s0f1", &specs); + try std.testing.expectEqual(@as(usize, 2), n); + try std.testing.expectEqualStrings("rocep1s0f1", specs[0].name[0..specs[0].name_len]); + try std.testing.expectEqual(@as(?u8, 3), specs[0].gid); + try std.testing.expectEqual(@as(?u8, null), specs[1].gid); + var a: Record = .{ .n_hca = 2 }; + var b: Record = .{ .n_hca = 2 }; + const sub100 = [16]u8{ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff, 192, 168, 100, 10 }; + var sub101 = sub100; + sub101[14] = 101; + a.gid = .{ sub100, sub101 }; + b.gid = .{ sub101, sub100 }; // the peer lists its functions the other way round + b.gid[0][15] = 11; + b.gid[1][15] = 11; + const p = try VerbsWire.pairing(a, b); + try std.testing.expectEqual(@as(u32, 1), p[0]); + try std.testing.expectEqual(@as(u32, 0), p[1]); + try std.testing.expect(ipv4Mapped(sub100)); + b.n_hca = 1; + try std.testing.expectError(error.BadRecord, VerbsWire.pairing(a, b)); +} diff --git a/zig/src/fabric/tp/run_hosts.sh b/zig/src/fabric/tp/run_hosts.sh new file mode 100644 index 0000000000..17a4ccbe91 --- /dev/null +++ b/zig/src/fabric/tp/run_hosts.sh @@ -0,0 +1,24 @@ +#!/usr/bin/env bash +# One rank of the two-host TP=2 gate over RoCE: `run_hosts.sh RANK RANK0_ADDR ./tf-tp-test OUTDIR` on both hosts. +set -u +RANK=${1:?rank}; MASTER=${2:?master address}; BIN=${3:?tf-tp-test}; OUT=${4:?outdir} +mkdir -p "$OUT" +[ -n "${NCCL_SOCKET_IFNAME:-}" ] && export NCCL_SOCKET_IFNAME # else NCCL picks the interface +ulimit -l unlimited 2> /dev/null || echo "warning: memlock is $(ulimit -l) KiB: ibv_reg_mr of the RoCE region may fail" >&2 +{ hostname; nvidia-smi -L; ibv_devices 2>&1; ulimit -l; } > "$OUT/host.txt" 2>&1 +PORT=29700 +step() { # NAME BACKEND ARGS...: this host's rank; log OUT/NAME.log ending with an EXIT line + local name=$1 backend=$2; shift 2 + PORT=$((PORT + 10)) + ( TF_TP_RANK=$RANK TF_TP_MASTER=$MASTER TF_TP_PORT=$PORT TF_COMM_BACKEND=$backend TF_TP_ROCE_FALLBACK=error \ + timeout 900 "$BIN" "$@"; echo "EXIT rank $RANK rc $? at $(date +%s.%N)" ) > "$OUT/$name.log" 2>&1 + echo "$name: $(grep -h '^SUMMARY\|^PASS roce-host\|^INJECT\|^EXIT' "$OUT/$name.log" | tr '\n' ' ')" +} +step roce-host roce roce-host +step suite-roce roce suite +step suite-nccl nccl suite +for backend in roce nccl; do + for mode in fatal kill; do + for victim in 0 1; do step "ff-$backend-$mode-$victim" "$backend" failfast "$victim" "$mode"; done + done +done diff --git a/zig/src/fabric/tp/run_pair.sh b/zig/src/fabric/tp/run_pair.sh new file mode 100644 index 0000000000..1cf73c6261 --- /dev/null +++ b/zig/src/fabric/tp/run_pair.sh @@ -0,0 +1,52 @@ +#!/usr/bin/env bash +# Two ranks of tf-tp-test on one two-GPU host: `[SUITES=..] [FAILFAST=..] run_pair.sh BIN OUTDIR`. +set -u +BIN=${1:?bin}; OUT=${2:?outdir} +mkdir -p "$OUT" +nvidia-smi > "$OUT/smi.txt" 2>&1 +nvidia-smi topo -m > "$OUT/topo.txt" 2>&1 +ldconfig -p | grep -E "libnccl|libcuda" > "$OUT/libs.txt" 2>&1 +PORT=29500 +pair() { # NAME BACKEND ARGS...: both ranks, logs OUT/NAME.rR.log ending with an EXIT line + local name=$1 backend=$2 r; shift 2 + PORT=$((PORT + 10)) + for r in 0 1; do + ( TF_TP_RANK=$r TF_TP_DEVICE=$r TF_TP_PORT=$PORT TF_COMM_BACKEND=$backend NCCL_DEBUG=${NCCL_DEBUG:-WARN} \ + timeout 600 "$BIN" "$@"; echo "EXIT rank $r rc $? at $(date +%s.%N)" ) > "$OUT/$name.r$r.log" 2>&1 & + done + wait +} +# ffsum OUT NAME VICTIM: a fail-fast pair's verdict line (exit codes, exit times after the injection) +ffsum() { + [ -f "$1/$2.r0.log" ] && [ -f "$1/$2.r1.log" ] || { echo "ffsum: $1/$2.r0.log or .r1.log missing" >&2; return 1; } + awk -v name="$2" -v victim="$3" ' + FNR == 1 { r = (FILENAME ~ /\.r1\.log$/) ? 1 : 0 } + !(r in ex) && match($0, /EXIT rank [0-9] rc [0-9]+ at [0-9.]+/) { + split(substr($0, RSTART, RLENGTH), f, " "); rc[r] = f[5]; at[r] = f[7]; ex[r] = 1 + } + r == victim && !inj && match($0, /INJECT [A-Za-z0-9_]+ rank [0-9]+ at [0-9.]+/) { + split(substr($0, RSTART, RLENGTH), g, " "); t0 = g[6]; inj = 1 + } + END { + if (!inj || !(0 in ex) || !(1 in ex)) { print name ": FAIL (no inject / exit line)"; exit } + for (i = 0; i < 2; i++) { want[i] = (i == victim && name ~ /kill/) ? 137 : 70; dt[i] = at[i] - t0 } + ok = rc[0] == want[0] && rc[1] == want[1] && dt[1 - victim] < 3.0 + printf "%s: %s rc [%d, %d] (want [%d, %d]), exit after injection: rank0 %.3f s, rank1 %.3f s\n", name, + ok ? "PASS" : "FAIL", rc[0], rc[1], want[0], want[1], dt[0], dt[1] + }' "$1/$2.r0.log" "$1/$2.r1.log" +} +for backend in ${SUITES:-mailbox nccl}; do + pair "suite-$backend" "$backend" suite +done +: > "$OUT/failfast.txt" +for backend in ${FAILFAST:-nccl mailbox}; do + for mode in fatal kill; do + for victim in 0 1; do + name=ff-$backend-$mode-$victim + pair "$name" "$backend" failfast "$victim" "$mode" + ffsum "$OUT" "$name" "$victim" >> "$OUT/failfast.txt" + done + done +done +grep -h "^SUMMARY\|^FAIL" "$OUT"/suite*.log +cat "$OUT/failfast.txt" diff --git a/zig/src/fabric/tp/sock.zig b/zig/src/fabric/tp/sock.zig new file mode 100644 index 0000000000..b382884034 --- /dev/null +++ b/zig/src/fabric/tp/sock.zig @@ -0,0 +1,169 @@ +//! TCP for the bootstrap and the fate channel: listen, accept and connect with deadlines, whole sends and receives. +const std = @import("std"); +const c = std.c; +const posix = std.posix; + +pub const Fd = c.fd_t; + +pub const Error = error{ SocketFailed, AddressUnresolved, AddressInUse, Timeout, Closed, Interrupted }; + +pub fn nowNs() u64 { + var ts: c.timespec = undefined; + _ = c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); +} + +pub fn sleepNs(ns: u64) void { + var ts: c.timespec = .{ .sec = @intCast(ns / std.time.ns_per_s), .nsec = @intCast(ns % std.time.ns_per_s) }; + while (c.nanosleep(&ts, &ts) != 0) {} +} + +/// TCP_NODELAY (frames are tiny and latency matters) and keepalive (a vanished host fails a read in ~90 s, not never). +fn tune(fd: Fd) void { + const one: c_int = 1; + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.NODELAY, &one, @sizeOf(c_int)); + _ = c.setsockopt(fd, posix.SOL.SOCKET, posix.SO.KEEPALIVE, &one, @sizeOf(c_int)); + const idle: c_int = 30; + const intvl: c_int = 10; + const cnt: c_int = 6; + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.KEEPIDLE, &idle, @sizeOf(c_int)); + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.KEEPINTVL, &intvl, @sizeOf(c_int)); + _ = c.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.KEEPCNT, &cnt, @sizeOf(c_int)); +} + +/// Rank 0's listening socket on every IPv4 address at `port`. +pub fn listen(port: u16) Error!Fd { + const fd = c.socket(posix.AF.INET, posix.SOCK.STREAM | posix.SOCK.CLOEXEC, 0); + if (fd < 0) return error.SocketFailed; + errdefer _ = c.close(fd); + const one: c_int = 1; + _ = c.setsockopt(fd, posix.SOL.SOCKET, posix.SO.REUSEADDR, &one, @sizeOf(c_int)); + var addr: posix.sockaddr.in = .{ .port = std.mem.nativeToBig(u16, port), .addr = 0 }; + if (c.bind(fd, @ptrCast(&addr), @sizeOf(posix.sockaddr.in)) != 0) return error.AddressInUse; + if (c.listen(fd, 64) != 0) return error.SocketFailed; + return fd; +} + +/// One connection before `deadline_ns` (monotonic). +pub fn accept(fd: Fd, deadline_ns: u64) Error!Fd { + while (true) { + try waitReadable(fd, deadline_ns); + const conn = c.accept(fd, null, null); + if (conn >= 0) { + tune(conn); + return conn; + } + } +} + +/// Connects to `host:port`, retrying every 50 ms (the master may not listen yet) until `deadline_ns`. +pub fn connect(host: []const u8, port: u16, deadline_ns: u64) Error!Fd { + var name: [256]u8 = undefined; + if (host.len >= name.len) return error.AddressUnresolved; + @memcpy(name[0..host.len], host); + name[host.len] = 0; + var service: [8]u8 = undefined; + const svc = std.fmt.bufPrintSentinel(&service, "{d}", .{port}, 0) catch unreachable; + const hints: c.addrinfo = .{ .flags = .{}, .family = posix.AF.INET, .socktype = posix.SOCK.STREAM, .protocol = 0, .addrlen = 0, .addr = null, .canonname = null, .next = null }; + while (true) { + var res: ?*c.addrinfo = null; + if (@backingInt(c.getaddrinfo(name[0..host.len :0], svc, &hints, &res)) == 0) { + defer c.freeaddrinfo(res.?); + const fd = c.socket(posix.AF.INET, posix.SOCK.STREAM | posix.SOCK.CLOEXEC, 0); + if (fd < 0) return error.SocketFailed; + if (c.connect(fd, res.?.addr.?, res.?.addrlen) == 0) { + tune(fd); + return fd; + } + _ = c.close(fd); + } + if (nowNs() >= deadline_ns) return error.Timeout; + sleepNs(50 * std.time.ns_per_ms); + } +} + +fn waitReadable(fd: Fd, deadline_ns: u64) Error!void { + while (true) { + const t = nowNs(); + if (t >= deadline_ns) return error.Timeout; + const ms: c_int = @intCast(@min((deadline_ns - t) / std.time.ns_per_ms + 1, 1000)); + var p = [_]c.pollfd{.{ .fd = fd, .events = posix.POLL.IN, .revents = 0 }}; + const n = c.poll(&p, 1, ms); + if (n > 0) return; + } +} + +/// Every byte of `bytes`, or `error.Closed`. +pub fn sendAll(fd: Fd, bytes: []const u8) Error!void { + var off: usize = 0; + while (off < bytes.len) { + const n = c.send(fd, bytes[off..].ptr, bytes.len - off, posix.MSG.NOSIGNAL); + if (n <= 0) { + if (n < 0 and posix.errno(n) == .INTR) continue; + return error.Closed; + } + off += @intCast(n); + } +} + +/// Fills `out` before `deadline_ns` (null: no deadline), or `error.Closed` at the end of the stream. +pub fn recvExact(fd: Fd, out: []u8, deadline_ns: ?u64) Error!void { + var off: usize = 0; + while (off < out.len) { + if (deadline_ns) |d| try waitReadable(fd, d); + const n = c.recv(fd, out[off..].ptr, out.len - off, 0); + if (n <= 0) { + if (n < 0 and posix.errno(n) == .INTR) continue; + return error.Closed; + } + off += @intCast(n); + } +} + +/// A non-blocking send of the whole of `bytes` if the socket buffer has room (false: dropped, or the peer is gone). +pub fn trySend(fd: Fd, bytes: []const u8) bool { + const n = c.send(fd, bytes.ptr, bytes.len, posix.MSG.NOSIGNAL | posix.MSG.DONTWAIT); + return n == @as(isize, @intCast(bytes.len)); +} + +pub fn hangup(fd: Fd) void { + _ = c.shutdown(fd, posix.SHUT.RDWR); +} + +pub fn close(fd: Fd) void { + _ = c.close(fd); +} + +/// This host's identity (the kernel's boot id, else the host name): ranks with equal ids share memory. +pub fn hostId(out: *[64]u8) []const u8 { + @memset(out, 0); + const f = c.open("/proc/sys/kernel/random/boot_id", .{ .ACCMODE = .RDONLY }); + if (f >= 0) { + defer _ = c.close(f); + const n = c.read(f, out, out.len); + if (n > 0) return std.mem.trim(u8, out[0..@intCast(n)], " \n"); + } + if (c.gethostname(out, out.len) != 0) return ""; + return std.mem.sliceTo(out, 0); +} + +test "a listener and a client trade bytes and see the hang-up" { + const l = try listen(0); + defer close(l); + var addr: posix.sockaddr.in = undefined; + var len: posix.socklen_t = @sizeOf(posix.sockaddr.in); + try std.testing.expectEqual(@as(c_int, 0), c.getsockname(l, @ptrCast(&addr), &len)); + const port = std.mem.bigToNative(u16, addr.port); + const deadline = nowNs() + 5 * std.time.ns_per_s; + const a = try connect("localhost", port, deadline); + defer close(a); + const b = try accept(l, deadline); + defer close(b); + try sendAll(a, "hello"); + var got: [5]u8 = undefined; + try recvExact(b, &got, deadline); + try std.testing.expectEqualStrings("hello", &got); + try std.testing.expectError(error.Timeout, recvExact(b, &got, nowNs() + 20 * std.time.ns_per_ms)); + hangup(a); + try std.testing.expectError(error.Closed, recvExact(b, &got, deadline)); +} diff --git a/zig/src/fabric/tp/tp.zig b/zig/src/fabric/tp/tp.zig new file mode 100644 index 0000000000..9094429c29 --- /dev/null +++ b/zig/src/fabric/tp/tp.zig @@ -0,0 +1,212 @@ +//! Tensor parallelism across processes, one GPU each: bootstrap, collectives, fail-fast; `Session.open` starts a rank. +const std = @import("std"); +const log = @import("log.zig"); +const cuda = @import("cuda"); + +pub const collective = @import("collective.zig"); +pub const reference = @import("reference.zig"); +pub const config = @import("config.zig"); +pub const sock = @import("sock.zig"); +pub const bootstrap = @import("bootstrap.zig"); +pub const fate = @import("fate.zig"); +pub const host = @import("host.zig"); +pub const nccl = @import("nccl.zig"); +pub const mailbox = @import("mailbox.zig"); +pub const hybrid = @import("hybrid.zig"); +pub const roce = @import("roce.zig"); +pub const roce_verbs = @import("roce_verbs.zig"); +pub const planlink = @import("planlink.zig"); +pub const Verbs = @import("verbs").Verbs; + +pub const Collective = collective.Collective; +pub const DType = collective.DType; +pub const Op = collective.Op; +pub const Config = config.Config; +pub const Bootstrap = bootstrap.Bootstrap; +pub const Fate = fate.Fate; + +pub const Session = struct { + cfg: Config, + lib: cuda.nccl.Library, + nccl: nccl.Nccl, + mail: ?mailbox.Mailbox = null, + roce: ?*roce.Roce = null, + route: hybrid.Hybrid = undefined, + /// The round plan's link (rank 0 sends, the others receive), on its own TCP connection. + plan: planlink.PlanLink = undefined, + fate: ?Fate = null, + + /// Brings this rank up in place; `kernel_image` is mailbox.cu's fatbin, empty sends those backends to NCCL. + pub fn open(s: *Session, d: *const cuda.Driver, cfg: Config, kernel_image: []const u8) !void { + var b = try Bootstrap.open(cfg); + defer b.close(); + try b.agree(&cfg.digest(), "settings (TF_TP_WORLD, TF_COMM_BACKEND, TF_TP_MAILBOX_MAX_KB, TF_TP_FAILFAST, TF_TP_ROCE_WIRE / _FALLBACK)"); + var lib = try cuda.nccl.Library.openPath(cfg.nccl_lib); + errdefer lib.close(); + const v = try lib.version(); + try b.agree(std.mem.asBytes(&v), "NCCL version"); + var id: cuda.nccl.UniqueId = undefined; + if (cfg.rank == 0) try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + try b.broadcast(&id.internal, &id.internal); + s.* = .{ .cfg = cfg, .lib = lib, .nccl = undefined }; + s.nccl = try nccl.Nccl.init(&s.lib, d, id, cfg.rank, cfg.world); + errdefer s.nccl.deinit(); + switch (cfg.backend) { + .nccl => {}, + .mailbox => if (try mailboxPossible(&b, kernel_image)) { + s.mail = try mailbox.Mailbox.init(d, &b, kernel_image, cfg.mailbox_max_bytes, cfg.mailbox_timeout_ns); + s.route = .{ .small = .{ .mailbox = &s.mail.? }, .nccl = &s.nccl, .max_bytes = cfg.mailbox_max_bytes }; + } else if (cfg.rank == 0) { + log.warn("[tensorfold] tp: TF_COMM_BACKEND=mailbox needs two ranks on one host and the kernel; every rank uses NCCL", .{}); + }, + .roce => try s.openRoce(d, &b, kernel_image), + } + try s.agreeVar(&b); + s.plan = planlink.PlanLink.init(cfg.rank, cfg.world, b.releasePlan()); + s.plan.spin_ns = cfg.plan_spin_ns; + s.plan.pin = cfg.plan_pin; + s.plan.avoid = cfg.roce_cpu; + if (cfg.failfast) { + s.fate = Fate.init(cfg, b.release()); + s.fate.?.onFail(.{ .ctx = s, .call = abortHook }); + try s.fate.?.start(); + } + log.info("[tensorfold] tp: rank {d} of {d} up ({t}, NCCL {d}, fail-fast {s}, roce fast {s}: {d} poller(s), stagger {d} ns)", .{ cfg.rank, cfg.world, s.comm().kind(), v, if (cfg.failfast) "on" else "off", if (cfg.roce_fast & 1 != 0) "on" else "off", (cfg.roce_fast >> 8) & 0xff, cfg.roce_fast >> 16 }); + } + + /// RoCE for the small exchanges, probed against NCCL at load; on failure every rank obeys TF_TP_ROCE_FALLBACK. + fn openRoce(s: *Session, d: *const cuda.Driver, b: *Bootstrap, image: []const u8) !void { + const cfg = s.cfg; + var specs: [roce_verbs.max_hcas]roce_verbs.HcaSpec = undefined; + var n_specs: usize = 0; + if (cfg.roce_wire == .verbs) { + if (cfg.roce_hca) |text| { + n_specs = try roce_verbs.parseSpecs(text, &specs); + } else if (Verbs.open(null)) |vb| { + var v = vb; + defer v.close(); + n_specs = roce_verbs.autoSpecs(&v, &specs, cfg.roce_hcas); + } else |_| {} + } + const opts: roce.Options = .{ .wire = if (cfg.roce_wire == .shm) .shm else .verbs, .hcas = specs[0..n_specs], .traffic_class = cfg.roce_tc, .cpu = cfg.roce_cpu, .max_bytes = cfg.mailbox_max_bytes, .timeout_ns = cfg.mailbox_timeout_ns, .fast = cfg.roce_fast, .blocks = cfg.roce_blocks }; + const r = roce.Roce.init(d, b, image, opts) catch |e| return s.roceFallback(e); + s.roce = r; + s.route = .{ .small = .{ .roce = r }, .nccl = &s.nccl, .max_bytes = cfg.mailbox_max_bytes }; + const good = probe(s, d) catch false; + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + var buf: [bootstrap.max_ranks]u8 = undefined; + const all = try b.allGather(&.{@intFromBool(good)}, &buf, &blobs); + for (all) |x| if (x[0] != 1) { + r.deinit(); + s.roce = null; + return s.roceFallback(error.ProbeMismatch); + }; + } + + /// `Collective.varAgreed`: whether every rank's transport moves device lengths, agreed once a session. + fn agreeVar(s: *Session, b: *Bootstrap) !void { + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + var buf: [bootstrap.max_ranks]u8 = undefined; + const all = try b.allGather(&.{@intFromBool(s.comm().varFits(16))}, &buf, &blobs); + var ok = true; + for (all) |x| ok = ok and x[0] == 1; + if (s.mail != null or s.roce != null) s.route.var_agreed = ok; + } + + fn roceFallback(s: *Session, e: anyerror) !void { + if (s.cfg.roce_fallback == .@"error") return e; + if (s.cfg.rank == 0) log.warn("[tensorfold] tp: RoCE unavailable ({t}); every rank uses NCCL (TF_TP_ROCE_FALLBACK=error refuses instead)", .{e}); + } + + /// One 4 KiB all-gather and one bf16 sum through RoCE and through NCCL: the bytes must be equal. + fn probe(s: *Session, d: *const cuda.Driver) !bool { + const n = 4096; + // input [0, n), gathers [n, 3n) and [3n, 5n), sums [5n, 6n) and [6n, 7n) + var a = try cuda.DeviceBuffer.alloc(d, 7 * n); + defer a.free(); + var stream = try cuda.Stream.init(d, true); + defer stream.deinit(); + var bytes: [n]u8 = undefined; + reference.fill(&bytes, .bf16, s.cfg.rank, 0x7072_6f62); + try a.uploadAsync(0, &bytes, stream.handle); + const fast = s.route.iface(); + const slow = s.nccl.iface(); + try fast.allGather(a.ptr, a.ptr + n, n, .u8, stream.handle); + try slow.allGather(a.ptr, a.ptr + 3 * n, n, .u8, stream.handle); + try fast.allReduce(a.ptr, a.ptr + 5 * n, n / 2, .bf16, .sum, stream.handle); + try slow.allReduce(a.ptr, a.ptr + 6 * n, n / 2, .bf16, .sum, stream.handle); + try stream.synchronize(); + try fast.check(); + var got: [6 * n]u8 = undefined; + try a.download(n, &got); + const gathers = std.mem.eql(u8, got[0 .. 2 * n], got[2 * n .. 4 * n]); + const sums = std.mem.eql(u8, got[4 * n .. 5 * n], got[5 * n .. 6 * n]); + if (!gathers or !sums) log.warn("[tensorfold] tp roce probe: gather {s}, sum {s} against NCCL", .{ if (gathers) "equal" else "DIFFERS", if (sums) "equal" else "DIFFERS" }); + return gathers and sums; + } + + fn mailboxPossible(b: *Bootstrap, image: []const u8) !bool { + if (b.world != 2) return false; + var id_buf: [64]u8 = undefined; + const id = sock.hostId(&id_buf); + var mine: [65]u8 = @splat(0); + mine[0] = @intFromBool(image.len > 0); + @memcpy(mine[1..][0..id.len], id); + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + var buf: [2 * 65]u8 = undefined; + const all = try b.allGather(&mine, &buf, &blobs); + return all[0][0] == 1 and std.mem.eql(u8, all[0], all[1]); + } + + fn abortHook(ctx: *anyopaque, _: []const u8) void { + const s: *Session = @ptrCast(@alignCast(ctx)); + s.comm().abort(); + } + + /// The communicator the model uses: NCCL, or the hybrid route over the mailbox / RoCE. + pub fn comm(s: *Session) Collective { + if (s.mail != null or s.roce != null) return s.route.iface(); + return s.nccl.iface(); + } + + /// This rank failed inside a round: every rank goes down (exit 70). Without fail-fast it only aborts. + pub fn fatal(s: *Session, where: []const u8, err: anyerror) void { + if (s.fate) |*f| return f.fatal(where, err); + s.comm().abort(); + } + + /// A clean stop announced by rank 0 before it closes: followers end with no fail-fast exit. + pub fn announceStop(s: *Session) void { + if (s.fate) |*f| f.announceStop(); + } + + /// This rank failed after `open`, before serving: the other ranks learn why and go down (fail-fast). + pub fn bootFailed(s: *Session, err: anyerror) void { + if (s.fate) |*f| f.failBoot("boot", err); + s.close(); + } + + /// A normal stop: every rank calls it (rank 0 first marks the stop as expected with `fate.expectStop`). + pub fn close(s: *Session) void { + if (s.fate) |*f| f.close(); + s.plan.close(); + if (s.roce) |r| r.deinit(); + if (s.mail) |*m| m.deinit(); + s.nccl.deinit(); + s.lib.close(); + } +}; + +test { + _ = collective; + _ = reference; + _ = config; + _ = sock; + _ = bootstrap; + _ = mailbox; + _ = roce_verbs; + _ = planlink; + _ = @import("roce_test.zig"); + _ = @import("host_test.zig"); + _ = @import("fate_test.zig"); +} diff --git a/zig/src/fabric/verbs.zig b/zig/src/fabric/verbs.zig index 1da50d386d..f3af36587b 100644 --- a/zig/src/fabric/verbs.zig +++ b/zig/src/fabric/verbs.zig @@ -1,7 +1,7 @@ //! The system's verbs library by dlopen (Apple's librdma, Linux libibverbs): every entry point resolved by name. const std = @import("std"); const builtin = @import("builtin"); -const abi = @import("verbs_abi.zig"); +pub const abi = @import("verbs_abi.zig"); pub const Error = error{ LibraryUnavailable, MissingSymbol, NoDeviceList, NoSuchDevice, VerbsFailed, BadDeviceList, NoUsableGid }; pub const max_links = 8; From 9af2ea2330935ec98cd72abddc92d785db5b5814 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:24:56 +0700 Subject: [PATCH 08/43] TP tests: a drain on rank 0 then its clean stop ends rank 1's loop at the stop word with no fail-fast exit; a rank killed mid-drain still fails the other fast. --- zig/src/fabric/tp/fate_test.zig | 95 +++++++++++++++++++++++++++++++++ 1 file changed, 95 insertions(+) diff --git a/zig/src/fabric/tp/fate_test.zig b/zig/src/fabric/tp/fate_test.zig index 6019055f50..728f4bfdf3 100644 --- a/zig/src/fabric/tp/fate_test.zig +++ b/zig/src/fabric/tp/fate_test.zig @@ -144,3 +144,98 @@ test "rank 0's clean stop ends every rank with no fail-fast exit, whatever order try std.testing.expectEqual(@as(u32, 0), exits.load(.acquire)); try std.testing.expectEqual(@as(u32, 0), hooks.load(.acquire)); } + +// -- a server's drain on two ranks: the plan link and the fate channel in the order a family's stop uses them -- + +const planlink = @import("planlink.zig"); + +/// Plan words a family's rank 0 sends: a round's operation, and a stop word that ends a follower's loop. +const op_round: i64 = 1; +const op_stop: i64 = 4; + +fn linkPair(links: *[2]planlink.PlanLink) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + const deadline = sock.nowNs() + 5 * std.time.ns_per_s; + const a = try sock.connect("localhost", l.port, deadline); + const b = try sock.accept(l.fd, deadline); + var f0: [planlink.max_ranks]sock.Fd = @splat(-1); + var f1: [planlink.max_ranks]sock.Fd = @splat(-1); + f0[1] = b; + f1[0] = a; + links[0] = planlink.PlanLink.init(0, 2, f0); + links[1] = planlink.PlanLink.init(1, 2, f1); +} + +/// Rank 1's loop: each round as it comes, until the stop word, or until the link ends. +const Follower = struct { + link: *planlink.PlanLink, + rounds: u32 = 0, + ended: enum { running, stopped, peer_closed, failed } = .running, + + fn run(f: *Follower) void { + var buf: [16]i64 = undefined; + while (true) { + const msg = f.link.recv(&buf) catch |e| { + f.ended = if (e == error.PeerClosed) .peer_closed else .failed; + return; + }; + if (msg.len == 0) continue; + if (msg[0] == op_stop) { + f.ended = .stopped; + return; + } + if (msg[0] == op_round) f.rounds += 1; + } + } +}; + +test "a drain on rank 0, then its clean stop at a round boundary: rank 1 ends its loop there and nobody exits 70" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + var links: [2]planlink.PlanLink = undefined; + try linkPair(&links); + var follower: Follower = .{ .link = &links[1] }; + const t = try std.Thread.spawn(.{}, Follower.run, .{&follower}); + // serving, then draining (rank 0 refuses new work, rounds go on; the fate channel is untouched), past a grace + var sent: u32 = 0; + for (0..30) |_| { + try links[0].send(&.{ op_round, sent }); + sent += 1; + sock.sleepNs(5 * std.time.ns_per_ms); + } + try std.testing.expect(!fates[0].closing.load(.acquire)); + // the stop's order: the stop word first, so the follower ends its loop before the fate STOP and the hang-up + try links[0].send(&.{op_stop}); + fates[0].announceStop(); + links[0].close(); + fates[0].close(); + t.join(); + // rank 1's loop returned: its model closes and the process exits 0 + fates[1].close(); + links[1].close(); + sock.sleepNs(2 * grace + 100 * std.time.ns_per_ms); + try std.testing.expectEqual(.stopped, follower.ended); + try std.testing.expectEqual(sent, follower.rounds); + try std.testing.expectEqual(@as(u32, 0), exits.load(.acquire)); + try std.testing.expectEqual(@as(u32, 0), hooks.load(.acquire)); +} + +test "a drain on rank 0 does not disarm fail-fast: rank 1 killed mid-drain still ends rank 0 (exit 70)" { + reset(); + var fates: [2]Fate = undefined; + try pair(&fates); + defer fates[0].close(); + var links: [2]planlink.PlanLink = undefined; + try linkPair(&links); + defer links[0].close(); + for (0..5) |i| try links[0].send(&.{ op_round, @intCast(i) }); + // rank 0 is draining (no stop announced); rank 1 is killed: its sockets close with no word + links[1].close(); + fates[1].expectStop(); + fates[1].close(); + _ = try waitExits(1); + try std.testing.expect(std.mem.indexOf(u8, fates[0].reason(), "rank 1 closed the fate channel") != null); + try std.testing.expectEqual(@as(u32, 1), exits.load(.acquire)); +} From 7c9a26e98040aa33298ad5e738cec03655a3d174 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:09:44 +0700 Subject: [PATCH 09/43] TP RoCE: with no TF_TP_ROCE_HCA, each rank picks the ports on an IPv4 /24 the peer also has The automatic pick took the first N active RoCE ports in device-list order and never checked that the peer can reach them. On a host with two neighbours, one picked port can sit on a subnet the peer does not have, and the queue pair's RTR step then fails, which reports RoCE unavailable. Every rank now sends its usable ports' IPv4 addresses over the bootstrap. Each rank then keeps the ports whose /24 the peer also lists, ordered by subnet and each /24 no more often than both sides have it, so both ranks pick the same number. With no shared /24 the old first-N order is kept, with a warning to set TF_TP_ROCE_HCA. The exchange runs on every rank, including ranks that set TF_TP_ROCE_HCA, so the bootstrap stays in step. --- zig/src/fabric/tp/config.zig | 2 +- zig/src/fabric/tp/roce_host.zig | 6 +- zig/src/fabric/tp/roce_pick.zig | 186 +++++++++++++++++++++++++++++++ zig/src/fabric/tp/roce_verbs.zig | 23 ++-- zig/src/fabric/tp/tp.zig | 12 +- 5 files changed, 207 insertions(+), 22 deletions(-) create mode 100644 zig/src/fabric/tp/roce_pick.zig diff --git a/zig/src/fabric/tp/config.zig b/zig/src/fabric/tp/config.zig index 13647f5921..ddc0c2a1d6 100644 --- a/zig/src/fabric/tp/config.zig +++ b/zig/src/fabric/tp/config.zig @@ -42,7 +42,7 @@ pub const Config = struct { /// Rank 1's memory report period on the fate channel (also how fast it notices rank 0 hanging up). report_ns: u64 = std.time.ns_per_s / 2, nccl_lib: []const u8 = "libnccl.so.2", - /// `name[:gid],...` RDMA devices for `roce` (null: every active Ethernet port with an IPv4 RoCE v2 GID). + /// `name[:gid],...` RDMA devices for `roce` (null: active RoCE v2 ports on an IPv4 /24 the peer has, roce_pick). roce_hca: ?[]const u8 = null, /// Stripe over at most this many HCAs (auto detection only). roce_hcas: u8 = 2, diff --git a/zig/src/fabric/tp/roce_host.zig b/zig/src/fabric/tp/roce_host.zig index afb705d41b..1dec07e9ba 100644 --- a/zig/src/fabric/tp/roce_host.zig +++ b/zig/src/fabric/tp/roce_host.zig @@ -11,11 +11,7 @@ pub fn run(cfg: tp.Config, gpa: std.mem.Allocator) !u8 { defer b.close(); var specs: [rv.max_hcas]rv.HcaSpec = undefined; var n: usize = 0; - if (cfg.roce_hca) |text| n = try rv.parseSpecs(text, &specs) else { - var v = try tp.Verbs.open(null); - defer v.close(); - n = rv.autoSpecs(&v, &specs, cfg.roce_hcas); - } + if (cfg.roce_wire == .verbs) n = try tp.roce_pick.resolve(cfg, &b, &specs); const r = try roce.Roce.init(null, &b, "", .{ .wire = if (cfg.roce_wire == .shm) .shm else .verbs, .hcas = specs[0..n], .traffic_class = cfg.roce_tc, .cpu = cfg.roce_cpu, .max_bytes = 256 * 1024, .timeout_ns = 10 * std.time.ns_per_s }); defer r.deinit(); const in = try gpa.alloc(u8, 256 * 1024); diff --git a/zig/src/fabric/tp/roce_pick.zig b/zig/src/fabric/tp/roce_pick.zig new file mode 100644 index 0000000000..627d0aefe1 --- /dev/null +++ b/zig/src/fabric/tp/roce_pick.zig @@ -0,0 +1,186 @@ +//! Which RoCE devices a rank stripes over: TF_TP_ROCE_HCA, else the usable ports on an IPv4 /24 the peer also has. +const std = @import("std"); +const log = @import("log.zig"); +const rv = @import("roce_verbs.zig"); +const bootstrap = @import("bootstrap.zig"); +const Config = @import("config.zig").Config; +const Verbs = @import("verbs").Verbs; + +/// Ports a rank lists to its peers (a DGX Spark's ConnectX-7 shows four RDMA devices). +pub const max_ports = 8; +const blob_bytes = 1 + 4 * max_ports; + +pub const Pick = struct { n: usize, shared: bool }; + +fn usable(p: rv.Port) bool { + return p.active and p.ethernet and p.ipv4 != null; +} + +fn subnet(a: [4]u8) u32 { + return @as(u32, a[0]) << 16 | @as(u32, a[1]) << 8 | a[2]; +} + +fn countIn(list: []const [4]u8, net: u32) usize { + var k: usize = 0; + for (list) |a| k += @intFromBool(subnet(a) == net); + return k; +} + +/// Usable ports on a /24 the peers list, by subnet, each /24 as often as both sides have it; else the first usable. +pub fn choose(ports: []const rv.Port, peer: []const [4]u8, out: *[rv.max_hcas]rv.HcaSpec, limit: usize) Pick { + const cap = @min(limit, rv.max_hcas); + var order: [max_ports]usize = undefined; + var m: usize = 0; + for (ports[0..@min(ports.len, max_ports)], 0..) |p, i| if (usable(p)) { + order[m] = i; + m += 1; + }; + var n: usize = 0; + for (order[0..m]) |i| n += @intFromBool(countIn(peer, subnet(ports[i].ipv4.?)) > 0); + if (n == 0) { + for (order[0..@min(m, cap)], 0..) |i, k| out[k] = ports[i].spec; + return .{ .n = @min(m, cap), .shared = false }; + } + // stable insertion sort by subnet (at most max_ports entries) + for (1..m) |a| { + var j = a; + while (j > 0 and subnet(ports[order[j - 1]].ipv4.?) > subnet(ports[order[j]].ipv4.?)) : (j -= 1) { + std.mem.swap(usize, &order[j - 1], &order[j]); + } + } + n = 0; + var taken: [max_ports]u32 = undefined; + for (order[0..m]) |i| { + if (n == cap) break; + const net = subnet(ports[i].ipv4.?); + var mine: usize = 0; + for (taken[0..n]) |t| mine += @intFromBool(t == net); + if (mine >= countIn(peer, net)) continue; + taken[n] = net; + out[n] = ports[i].spec; + n += 1; + } + return .{ .n = n, .shared = true }; +} + +/// `[count][count x IPv4]` of the usable ports, the blob each rank sends. +fn encode(ports: []const rv.Port, out: *[blob_bytes]u8) []const u8 { + var n: u8 = 0; + for (ports[0..@min(ports.len, max_ports)]) |p| if (usable(p)) { + out[1 + 4 * @as(usize, n) ..][0..4].* = p.ipv4.?; + n += 1; + }; + out[0] = n; + return out[0 .. 1 + 4 * @as(usize, n)]; +} + +/// A peer's blob into `out`; a malformed one counts as no ports. +fn decode(blob: []const u8, out: [][4]u8) usize { + if (blob.len == 0 or blob[0] > max_ports or blob.len != 1 + 4 * @as(usize, blob[0]) or out.len < blob[0]) return 0; + for (0..blob[0]) |k| out[k] = blob[1 + 4 * k ..][0..4].*; + return blob[0]; +} + +/// The devices for the verbs wire. A collective: every rank calls it, including ranks that set TF_TP_ROCE_HCA. +pub fn resolve(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSpec) !usize { + var ports: [max_ports]rv.Port = undefined; + var n_ports: usize = 0; + if (Verbs.open(null)) |opened| { + var v = opened; + defer v.close(); + n_ports = rv.scanPorts(&v, &ports); + } else |_| {} + var mine: [blob_bytes]u8 = undefined; + var buf: [bootstrap.max_ranks * blob_bytes]u8 = undefined; + var blobs: [bootstrap.max_ranks][]const u8 = undefined; + const all = try b.allGather(encode(ports[0..n_ports], &mine), &buf, &blobs); + if (cfg.roce_hca) |text| return rv.parseSpecs(text, out); + var peer: [bootstrap.max_ranks * max_ports][4]u8 = undefined; + var n_peer: usize = 0; + for (all, 0..) |blob, r| if (r != b.rank) { + n_peer += decode(blob, peer[n_peer..]); + }; + const pick = choose(ports[0..n_ports], peer[0..n_peer], out, cfg.roce_hcas); + if (!pick.shared and pick.n > 0) log.warn("tp roce: no RoCE port shares an IPv4 /24 with the peer's; using the first {d} active (set TF_TP_ROCE_HCA)", .{pick.n}); + return pick.n; +} + +fn port(name: []const u8, active: bool, ipv4: ?[4]u8) rv.Port { + var p: rv.Port = .{ .active = active, .ethernet = true, .ipv4 = ipv4 }; + p.spec.name_len = @intCast(name.len); + @memcpy(p.spec.name[0..name.len], name); + p.spec.gid = 3; + return p; +} + +fn names(specs: []const rv.HcaSpec, buf: []u8) []const u8 { + var w: std.Io.Writer = .fixed(buf); + for (specs, 0..) |s, i| w.print("{s}{s}", .{ if (i == 0) "" else ",", s.name[0..s.name_len] }) catch unreachable; + return w.buffered(); +} + +// A host with two neighbours: port f0 cabled to one Spark (both functions on one /24), port f1 to the other. +const a_ports = [_]rv.Port{ + port("rocep1s0f0", true, .{ 203, 0, 113, 1 }), + port("rocep1s0f1", true, .{ 192, 0, 2, 1 }), + port("roceP2p1s0f0", true, .{ 203, 0, 113, 65 }), + port("roceP2p1s0f1", true, .{ 198, 51, 100, 1 }), +}; +// Its peer for this run: only port f1 is cabled, f0 is down. +const b_ports = [_]rv.Port{ + port("rocep1s0f0", false, null), + port("rocep1s0f1", true, .{ 192, 0, 2, 2 }), + port("roceP2p1s0f0", false, null), + port("roceP2p1s0f1", true, .{ 198, 51, 100, 2 }), +}; + +fn listed(ports: []const rv.Port, out: *[max_ports][4]u8) []const [4]u8 { + var blob: [blob_bytes]u8 = undefined; + return out[0..decode(encode(ports, &blob), out)]; +} + +test "a two-neighbour host picks the rails its peer shares, not its first two active ports" { + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var buf: [128]u8 = undefined; + // the old rule (no peer list): the first two active ports, one of them cabled to the other neighbour + const old = choose(&a_ports, &.{}, &specs, 2); + try std.testing.expect(!old.shared); + try std.testing.expectEqualStrings("rocep1s0f0,rocep1s0f1", names(specs[0..old.n], &buf)); + var peer_buf: [max_ports][4]u8 = undefined; + const b_addrs = listed(&b_ports, &peer_buf); + try std.testing.expectEqual(@as(usize, 0), countIn(b_addrs, subnet(.{ 203, 0, 113, 1 }))); + // with the peer's addresses: the two functions of the cabled port, on each side + const a = choose(&a_ports, b_addrs, &specs, 2); + try std.testing.expect(a.shared); + try std.testing.expectEqualStrings("rocep1s0f1,roceP2p1s0f1", names(specs[0..a.n], &buf)); + var a_buf: [max_ports][4]u8 = undefined; + const b = choose(&b_ports, listed(&a_ports, &a_buf), &specs, 2); + try std.testing.expectEqual(@as(usize, 2), b.n); + try std.testing.expectEqualStrings("rocep1s0f1,roceP2p1s0f1", names(specs[0..b.n], &buf)); + // one rail asked for: the lower /24 on both sides + try std.testing.expectEqualStrings("rocep1s0f1", names(specs[0..choose(&a_ports, b_addrs, &specs, 1).n], &buf)); +} + +test "two rails on one /24 count once per peer port; no shared /24 keeps the list order" { + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var buf: [128]u8 = undefined; + // the neighbour cabled to port f0 lists one address on that /24: one rail each, equal counts + const one = [_][4]u8{.{ 203, 0, 113, 2 }}; + const p = choose(&a_ports, &one, &specs, 2); + try std.testing.expectEqual(@as(usize, 1), p.n); + try std.testing.expectEqualStrings("rocep1s0f0", names(specs[0..p.n], &buf)); + const two = [_][4]u8{ .{ 203, 0, 113, 2 }, .{ 203, 0, 113, 66 } }; + try std.testing.expectEqualStrings("rocep1s0f0,roceP2p1s0f0", names(specs[0..choose(&a_ports, &two, &specs, 2).n], &buf)); + // inactive and non-Ethernet ports never count; a peer on other subnets falls back to the first usable ports + var mixed = b_ports; + mixed[1].ethernet = false; + const other = [_][4]u8{.{ 198, 51, 100, 9 }}; + const q = choose(&mixed, &.{.{ 203, 0, 113, 9 }}, &specs, 2); + try std.testing.expect(!q.shared); + try std.testing.expectEqualStrings("roceP2p1s0f1", names(specs[0..q.n], &buf)); + try std.testing.expectEqualStrings("roceP2p1s0f1", names(specs[0..choose(&mixed, &other, &specs, 2).n], &buf)); + // a malformed blob is no ports; the detected GID index travels with the pick + var out: [max_ports][4]u8 = undefined; + try std.testing.expectEqual(@as(usize, 0), decode(&.{ 2, 1, 2, 3, 4 }, &out)); + try std.testing.expectEqual(@as(?u8, 3), specs[0].gid); +} diff --git a/zig/src/fabric/tp/roce_verbs.zig b/zig/src/fabric/tp/roce_verbs.zig index bc99681b5e..a540a4ce02 100644 --- a/zig/src/fabric/tp/roce_verbs.zig +++ b/zig/src/fabric/tp/roce_verbs.zig @@ -89,23 +89,32 @@ fn detectGid(v: *const Verbs, ctx: *abi.Context, name: []const u8, table: c_int) return null; } -/// Every device whose port 1 is an active Ethernet port with an IPv4 RoCE v2 GID, by name order, at most `limit`. -pub fn autoSpecs(v: *const Verbs, out: *[max_hcas]HcaSpec, limit: usize) usize { +/// A device's port 1 as `roce_pick.choose` sees it: up, Ethernet, and its IPv4 RoCE v2 GID (index and address). +pub const Port = struct { spec: HcaSpec = .{}, active: bool = false, ethernet: bool = false, ipv4: ?[4]u8 = null }; + +/// Port 1 of every device that opens, in the device list's order, at most `out.len`. +pub fn scanPorts(v: *const Verbs, out: []Port) usize { var count: c_int = 0; const list = v.api.ibv_get_device_list(&count) orelse return 0; defer v.api.ibv_free_device_list(list); var n: usize = 0; var i: usize = 0; while (list[i]) |dev| : (i += 1) { - if (n == limit) break; + if (n == out.len) break; const name = std.mem.span(v.api.ibv_get_device_name(dev) orelse continue); const ctx = v.api.ibv_open_device(dev) orelse continue; defer _ = v.api.ibv_close_device(ctx); const port = v.port(ctx, 1) catch continue; - if (port.state != abi.port_active or port.link_layer != link_ethernet) continue; - const g = detectGid(v, ctx, name, port.gid_tbl_len) orelse continue; - out[n] = .{ .name_len = @intCast(@min(name.len, 63)), .gid = g }; - @memcpy(out[n].name[0..out[n].name_len], name[0..out[n].name_len]); + var p: Port = .{ .active = port.state == abi.port_active, .ethernet = port.link_layer == link_ethernet }; + p.spec.name_len = @intCast(@min(name.len, 63)); + @memcpy(p.spec.name[0..p.spec.name_len], name[0..p.spec.name_len]); + if (p.active and p.ethernet) if (detectGid(v, ctx, name, port.gid_tbl_len)) |gi| { + if (v.gid(ctx, 1, gi)) |g| { + p.spec.gid = gi; + p.ipv4 = g.raw[12..16].*; + } else |_| {} + }; + out[n] = p; n += 1; } return n; diff --git a/zig/src/fabric/tp/tp.zig b/zig/src/fabric/tp/tp.zig index 9094429c29..3fe6d3fab5 100644 --- a/zig/src/fabric/tp/tp.zig +++ b/zig/src/fabric/tp/tp.zig @@ -15,6 +15,7 @@ pub const mailbox = @import("mailbox.zig"); pub const hybrid = @import("hybrid.zig"); pub const roce = @import("roce.zig"); pub const roce_verbs = @import("roce_verbs.zig"); +pub const roce_pick = @import("roce_pick.zig"); pub const planlink = @import("planlink.zig"); pub const Verbs = @import("verbs").Verbs; @@ -79,15 +80,7 @@ pub const Session = struct { const cfg = s.cfg; var specs: [roce_verbs.max_hcas]roce_verbs.HcaSpec = undefined; var n_specs: usize = 0; - if (cfg.roce_wire == .verbs) { - if (cfg.roce_hca) |text| { - n_specs = try roce_verbs.parseSpecs(text, &specs); - } else if (Verbs.open(null)) |vb| { - var v = vb; - defer v.close(); - n_specs = roce_verbs.autoSpecs(&v, &specs, cfg.roce_hcas); - } else |_| {} - } + if (cfg.roce_wire == .verbs) n_specs = try roce_pick.resolve(cfg, b, &specs); const opts: roce.Options = .{ .wire = if (cfg.roce_wire == .shm) .shm else .verbs, .hcas = specs[0..n_specs], .traffic_class = cfg.roce_tc, .cpu = cfg.roce_cpu, .max_bytes = cfg.mailbox_max_bytes, .timeout_ns = cfg.mailbox_timeout_ns, .fast = cfg.roce_fast, .blocks = cfg.roce_blocks }; const r = roce.Roce.init(d, b, image, opts) catch |e| return s.roceFallback(e); s.roce = r; @@ -205,6 +198,7 @@ test { _ = bootstrap; _ = mailbox; _ = roce_verbs; + _ = roce_pick; _ = planlink; _ = @import("roce_test.zig"); _ = @import("host_test.zig"); From 8ebb0be41afddfc33b4dacea44b285bc01796eb3 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:09:44 +0700 Subject: [PATCH 10/43] TP README: when -Dtp-nvcc, TF_TP_ROCE_HCA, NCCL_IB_HCA, NCCL_SOCKET_IFNAME and TF_TP_MASTER are needed, with a two-host example --- zig/src/fabric/tp/README.md | 53 +++++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) create mode 100644 zig/src/fabric/tp/README.md diff --git a/zig/src/fabric/tp/README.md b/zig/src/fabric/tp/README.md new file mode 100644 index 0000000000..5d5c57358f --- /dev/null +++ b/zig/src/fabric/tp/README.md @@ -0,0 +1,53 @@ +# Tensor parallelism across processes + +`zig/src/fabric/tp` runs tensor parallelism with one GPU per process, either on one host or on two hosts over RoCE. +`TF_COMM_BACKEND` picks the transport for small exchanges: `nccl` (the default), `mailbox` (two GPUs on one host) +or `roce` (two hosts). NCCL carries everything larger, so it must work on its own as well. + +## Building `tf-tp-test` + +The `mailbox` and `roce` backends need the mailbox kernel (`mailbox.cu`) built into the binary: + +```sh +zig build tf-tp-test -Doptimize=ReleaseSafe -Dtp-nvcc=/usr/local/cuda/bin/nvcc +# or embed a fatbin built elsewhere: +zig build tf-tp-test -Doptimize=ReleaseSafe -Dtp-fatbin=/abs/path/tp_mailbox.fatbin +``` + +A build without `-Dtp-nvcc` or `-Dtp-fatbin` still runs `nccl` and `roce-host`, which needs no GPU. +In that build the `mailbox` and `roce` suites fail at load with `NoKernel` (an engine falls back to NCCL instead). +The `INFO` line `tf-tp-test` prints at start shows the embedded kernel's size; `0 B` means none was built in. + +## Choosing the RDMA devices + +| Variable | Used by | Needed when | +| --- | --- | --- | +| `TF_TP_ROCE_HCA` | the RoCE wire | the host has more than one active RoCE port that is not cabled to the peer | +| `NCCL_IB_HCA` | NCCL | the same as above; NCCL otherwise picks its own devices | +| `NCCL_SOCKET_IFNAME` | NCCL bootstrap | the host has several IP interfaces and NCCL picks one the peer cannot reach | +| `TF_TP_MASTER` | the TCP bootstrap | always on two hosts: rank 0's address as the other rank reaches it | + +`TF_TP_ROCE_HCA` is `name[:gid],name[:gid]`: up to two devices, and optionally each one's RoCE v2 GID index. +Without it the rank picks the devices itself, at most `TF_TP_ROCE_HCAS` (default 2). +A host cabled to more than one neighbour should always set it explicitly. +The two hosts may list their devices in a different order and use different names. +Each device pairs with the peer device on the same IPv4 /24. + +`NCCL_IB_HCA` should name the same devices without the `:gid` suffix. +In NCCL, `:N` after a device name selects port `N`, not a GID index. + +Two DGX Spark nodes, with the second ConnectX-7 port of each cabled to the other (both PCIe functions of that +port), and the first port cabled to a third node: + +```sh +# on both hosts; rank 0's link address is 192.0.2.1 +export TF_TP_ROCE_HCA=rocep1s0f1,roceP2p1s0f1 +export NCCL_IB_HCA=rocep1s0f1,roceP2p1s0f1 +export NCCL_SOCKET_IFNAME=enp1s0f1np1 +zig/src/fabric/tp/run_hosts.sh 0 192.0.2.1 ./zig-out/bin/tf-tp-test out-rank0 # host A +zig/src/fabric/tp/run_hosts.sh 1 192.0.2.1 ./zig-out/bin/tf-tp-test out-rank1 # host B +``` + +`ibv_devices` lists the RDMA devices; `ibv_devinfo -d NAME` shows a port's state and link layer. +`/sys/class/infiniband/NAME/ports/1/gids/N` and `gid_attrs/ndevs/N` show each GID and its netdev. +`ulimit -l` must allow the RoCE region to be pinned (`unlimited` is simplest). From e9e005656edad761d71c7bb2d501702bb46da8bc Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:09:44 +0700 Subject: [PATCH 11/43] TP test runner: a binary built without the mailbox kernel refuses the kernel backends up front; run_hosts.sh passes HCA picks - `tf-tp-test kernel` reports the embedded kernel and exits 1 without one. `suite` and `failfast` on the mailbox or RoCE backend stop at once with a message that names `-Dtp-nvcc`, instead of failing later on a missing kernel. - `run_hosts.sh` stops before any step when the kernel is missing. It exports TF_TP_ROCE_HCA and NCCL_IB_HCA (derived from TF_TP_ROCE_HCA when unset), warns when a host has more than one RDMA device and neither is set, and records all three variables in host.txt. --- zig/src/fabric/tp/gpu_main.zig | 13 +++++++++++-- zig/src/fabric/tp/run_hosts.sh | 11 ++++++++++- 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/zig/src/fabric/tp/gpu_main.zig b/zig/src/fabric/tp/gpu_main.zig index c5668a7b06..e3b9a420e7 100644 --- a/zig/src/fabric/tp/gpu_main.zig +++ b/zig/src/fabric/tp/gpu_main.zig @@ -1,4 +1,4 @@ -//! `tf-tp-test suite | failfast RANK fatal|kill | roce-host`: one rank of a two-process TP=2 test on real GPUs. +//! `tf-tp-test suite | failfast RANK fatal|kill | roce-host | kernel`: one rank of a two-process TP=2 test. const std = @import("std"); const cuda = @import("cuda"); const tp = @import("tp"); @@ -11,11 +11,20 @@ const max_bytes = (4 << 20) + 4096; pub fn main(init: std.process.Init) !u8 { const args = try init.minimal.args.toSlice(init.arena.allocator()); if (args.len < 2) { - std.debug.print("usage: tf-tp-test suite | failfast RANK fatal|kill | roce-host\n", .{}); + std.debug.print("usage: tf-tp-test suite | failfast RANK fatal|kill | roce-host | kernel\n", .{}); return 2; } + // `kernel`: whether the build embeds the mailbox kernel (-Dtp-nvcc or -Dtp-fatbin), for scripts to check first + if (std.mem.eql(u8, args[1], "kernel")) { + std.debug.print("mailbox kernel {d} B\n", .{tp_kernel.mailbox.len}); + return if (tp_kernel.mailbox.len > 0) 0 else 1; + } const cfg = try tp.Config.fromEnv(); if (std.mem.eql(u8, args[1], "roce-host")) return @import("roce_host.zig").run(cfg, init.gpa); + if (cfg.backend != .nccl and tp_kernel.mailbox.len == 0) { + std.debug.print("FAIL TF_COMM_BACKEND={t} needs the mailbox kernel: rebuild with -Dtp-nvcc or -Dtp-fatbin\n", .{cfg.backend}); + return 2; + } var driver = try cuda.Driver.open(); defer driver.close(); var ctx = try cuda.Context.init(&driver, @intCast(cfg.device)); diff --git a/zig/src/fabric/tp/run_hosts.sh b/zig/src/fabric/tp/run_hosts.sh index 17a4ccbe91..5f00168f83 100644 --- a/zig/src/fabric/tp/run_hosts.sh +++ b/zig/src/fabric/tp/run_hosts.sh @@ -4,8 +4,17 @@ set -u RANK=${1:?rank}; MASTER=${2:?master address}; BIN=${3:?tf-tp-test}; OUT=${4:?outdir} mkdir -p "$OUT" [ -n "${NCCL_SOCKET_IFNAME:-}" ] && export NCCL_SOCKET_IFNAME # else NCCL picks the interface +# NCCL uses the RoCE wire's devices unless NCCL_IB_HCA is set: the names without `:gid` (NCCL reads `:N` as a port) +[ -n "${TF_TP_ROCE_HCA:-}" ] && [ -z "${NCCL_IB_HCA:-}" ] && NCCL_IB_HCA=$(echo "$TF_TP_ROCE_HCA" | sed 's/:[0-9]*//g; s/ //g') +[ -n "${TF_TP_ROCE_HCA:-}" ] && export TF_TP_ROCE_HCA +[ -n "${NCCL_IB_HCA:-}" ] && export NCCL_IB_HCA +"$BIN" kernel > /dev/null 2>&1 || { echo "error: $BIN has no mailbox kernel: rebuild with -Dtp-nvcc=PATH or -Dtp-fatbin=PATH" >&2; exit 2; } ulimit -l unlimited 2> /dev/null || echo "warning: memlock is $(ulimit -l) KiB: ibv_reg_mr of the RoCE region may fail" >&2 -{ hostname; nvidia-smi -L; ibv_devices 2>&1; ulimit -l; } > "$OUT/host.txt" 2>&1 +DEVICES=$(ibv_devices 2> /dev/null | tail -n +3 | grep -c .) +if [ "$DEVICES" -gt 1 ] && [ -z "${TF_TP_ROCE_HCA:-}${NCCL_IB_HCA:-}" ]; then + echo "warning: $DEVICES RDMA devices and no TF_TP_ROCE_HCA / NCCL_IB_HCA: name the peer's ports (README.md)" >&2 +fi +{ hostname; nvidia-smi -L; ibv_devices 2>&1; ulimit -l; env | grep -E '^(TF_TP_ROCE_HCA|NCCL_IB_HCA|NCCL_SOCKET_IFNAME)='; } > "$OUT/host.txt" 2>&1 PORT=29700 step() { # NAME BACKEND ARGS...: this host's rank; log OUT/NAME.log ending with an EXIT line local name=$1 backend=$2; shift 2 From 0472cbf9a2ea642520dd1ba5b90589cfaba1e472 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:39:14 +0700 Subject: [PATCH 12/43] TP RoCE: every RDMA post is bounds-checked before the verb, and peers with another region layout are refused roce_guard.stripes computes each HCA's share of a post (its send stripe, the peer's receive range and its flag word) and refuses with error.OutOfRegion any range past its slot, its flag area or the registered region, before post_send. The verbs wire and the shm wire both post through it. The shm wire moves to roce_shm.zig and can stripe like the verbs wire (Options.shm_stripes), with an optional tap that sees every post for tests. The connection record carries the slot and registered length, so a peer with another layout is refused at connect; its magic is bumped. Tests: stripes tile the slot over two HCAs and refuse a post past the slot, a third slot, an unpadded or empty post, no HCA and a short region. On fake verbs a post past its slot reaches no verb, and a full slot stays inside the region. Over the shm wire striped as two HCAs, 2,000 variable-length gathers stay inside slot and region, bit for bit. --- zig/src/fabric/tp/roce.zig | 60 ++++------------- zig/src/fabric/tp/roce_fake.zig | 96 ++++++++++++++++++++++++++++ zig/src/fabric/tp/roce_guard.zig | 63 ++++++++++++++++++ zig/src/fabric/tp/roce_post_test.zig | 20 ++++++ zig/src/fabric/tp/roce_shm.zig | 63 ++++++++++++++++++ zig/src/fabric/tp/roce_test.zig | 41 +++++++++++- zig/src/fabric/tp/roce_verbs.zig | 44 ++++++++----- zig/src/fabric/tp/tp.zig | 2 + 8 files changed, 323 insertions(+), 66 deletions(-) create mode 100644 zig/src/fabric/tp/roce_fake.zig create mode 100644 zig/src/fabric/tp/roce_guard.zig create mode 100644 zig/src/fabric/tp/roce_post_test.zig create mode 100644 zig/src/fabric/tp/roce_shm.zig diff --git a/zig/src/fabric/tp/roce.zig b/zig/src/fabric/tp/roce.zig index 58f88e3441..7806fd60d5 100644 --- a/zig/src/fabric/tp/roce.zig +++ b/zig/src/fabric/tp/roce.zig @@ -6,6 +6,8 @@ const collective = @import("collective.zig"); const Bootstrap = @import("bootstrap.zig").Bootstrap; const sock = @import("sock.zig"); const rv = @import("roce_verbs.zig"); +const shm = @import("roce_shm.zig"); +const guard = @import("roce_guard.zig"); const HostRegister = @import("mailbox.zig").HostRegister; const posix = std.posix; const DevicePtr = collective.DevicePtr; @@ -71,30 +73,20 @@ pub const Options = struct { fast: u32 = 0, /// TF_TP_ROCE_BLOCKS: the largest grid (config.roce_blocks) blocks: u32 = 16, -}; - -/// The shm wire: the peer's region is mapped here; a post copies the slot, then stores the flag with release order. -const ShmWire = struct { - peer: []u8, - local: []u8, - slot_bytes: usize, - - fn post(w: *ShmWire, slot: u32, seq: u32, padded: u32) void { - const recv = send_off + 2 * w.slot_bytes; - @memcpy(w.peer[recv + slot * w.slot_bytes ..][0..padded], w.local[send_off + slot * w.slot_bytes ..][0..padded]); - const flag: *std.atomic.Value(u32) = @ptrCast(@alignCast(w.peer[flag_off + slot * 2 * flag_stride ..].ptr)); - flag.store(seq, .release); - } + /// The shm wire's stripes, as HCAs: its posts take the verbs wire's split and bounds. + shm_stripes: u32 = 1, + /// Tests: sees every shm post's stripes before the copy. + tap: ?*shm.PostTap = null, }; const Wire = union(WireKind) { verbs: rv.VerbsWire, - shm: ShmWire, + shm: shm.ShmWire, fn post(w: *Wire, slot: u32, seq: u32, padded: u32) !void { switch (w.*) { .verbs => |*v| try v.post(slot, seq, padded), - .shm => |*s| s.post(slot, seq, padded), + .shm => |*s| try s.post(slot, seq, padded), } } @@ -108,7 +100,7 @@ const Wire = union(WireKind) { fn hcas(w: *const Wire) u32 { return switch (w.*) { .verbs => |v| v.n_hca, - .shm => 1, + .shm => |s| s.stripes, }; } }; @@ -204,9 +196,10 @@ pub const Roce = struct { r.wire = .{ .verbs = try rv.VerbsWire.open(o.hcas, r.region, .{ .send = send_off, .recv = send_off + 2 * r.slot_bytes, .flags = flag_off, .slot = r.slot_bytes }, o.traffic_class) }; }, .shm => { - r.mapping = try sharedPair(b, len); + r.mapping = try shm.sharedPair(b, len); r.region = @alignCast(r.mapping[r.me * len ..][0..len]); - r.wire = .{ .shm = .{ .peer = r.mapping[(1 - r.me) * len ..][0..len], .local = r.region, .slot_bytes = r.slot_bytes } }; + const layout: guard.Layout = .{ .send = send_off, .recv = send_off + 2 * r.slot_bytes, .flags = flag_off, .slot = r.slot_bytes, .len = len }; + r.wire = .{ .shm = .{ .peer = r.mapping[(1 - r.me) * len ..][0..len], .layout = layout, .local = r.region, .stripes = o.shm_stripes, .tap = o.tap } }; }, } return true; @@ -422,35 +415,6 @@ pub const Roce = struct { } }; -/// One shared mapping holding both ranks' regions (rank 0 creates and sizes the file before anyone learns its name). -fn sharedPair(b: *Bootstrap, len: usize) ![]align(page) u8 { - var name: [64]u8 = undefined; - var name_len: u8 = 0; - var fd: std.c.fd_t = -1; - defer if (fd >= 0) { - _ = std.c.close(fd); - }; - if (b.rank == 0) blk: { - const s = std.fmt.bufPrint(&name, "/dev/shm/tf-tp-roce-{d}-{x}", .{ std.c.getpid(), sock.nowNs() }) catch break :blk; - name[s.len] = 0; - fd = std.c.open(name[0..s.len :0], .{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true, .CLOEXEC = true }, @as(std.c.mode_t, 0o600)); - if (fd >= 0 and std.c.ftruncate(fd, @intCast(2 * len)) == 0) name_len = @intCast(s.len); - } - // an empty name tells the other rank that rank 0 failed (it must not wait for a name that never comes) - try b.broadcast(std.mem.asBytes(&name_len), std.mem.asBytes(&name_len)); - if (name_len == 0) return error.SharedMemoryUnavailable; - try b.broadcast(name[0..name_len], name[0..name_len]); - name[name_len] = 0; - if (b.rank != 0) { - fd = std.c.open(name[0..name_len :0], .{ .ACCMODE = .RDWR, .CLOEXEC = true }, @as(std.c.mode_t, 0)); - if (fd < 0) return error.SharedMemoryUnavailable; - } - const m = try posix.mmap(null, 2 * len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); - try b.agree("mapped", "RoCE shm wire"); - if (b.rank == 0) _ = std.c.unlink(name[0..name_len :0]); - return m; -} - fn pin(cpu: u32) void { var set: std.os.linux.cpu_set_t = @splat(0); set[cpu / @bitSizeOf(usize)] |= @as(usize, 1) << @intCast(cpu % @bitSizeOf(usize)); diff --git a/zig/src/fabric/tp/roce_fake.zig b/zig/src/fabric/tp/roce_fake.zig new file mode 100644 index 0000000000..e556c1990a --- /dev/null +++ b/zig/src/fabric/tp/roce_fake.zig @@ -0,0 +1,96 @@ +//! Fake verbs for host tests of VerbsWire: each QP's posted writes counted from the requests, completions that lag. +const std = @import("std"); +const verbs = @import("verbs"); +const abi = verbs.abi; +const rv = @import("roce_verbs.zig"); + +pub const slot: u64 = 1 << 20; +pub const send_off = 8192; + +/// What the HCAs hold (posted, not completed); every count comes from the work requests themselves. +pub const Fake = struct { + pub var region: []const u8 = &.{}; + pub var lag: u32 = 1; + pub var polls: u64 = 0; + pub var queue: [2][512]u64 = undefined; + pub var head: [2]u32 = .{ 0, 0 }; + pub var len: [2]u32 = .{ 0, 0 }; + pub var inflight: u64 = 0; + pub var max_inflight: u64 = 0; + pub var max_write: u64 = 0; + pub var outside: u64 = 0; + pub var bytes: u64 = 0; + pub var posts: u64 = 0; + + pub fn reset(l: u32) void { + lag = l; + polls = 0; + head = .{ 0, 0 }; + len = .{ 0, 0 }; + inflight = 0; + max_inflight = 0; + max_write = 0; + outside = 0; + bytes = 0; + posts = 0; + } + + pub fn postSend(qp: *abi.Qp, wr: *abi.SendWr, _: *?*abi.SendWr) callconv(.c) c_int { + const q = qp.qp_num; + posts += 1; + var sum: u64 = 0; + var at: ?*abi.SendWr = wr; + while (at) |w| : (at = w.next) { + if (w.send_flags & abi.send_inline == 0) { + const sge = w.sg_list[0]; + const lo = @intFromPtr(region.ptr); + if (sge.addr < lo + send_off or sge.addr + sge.length > lo + send_off + 2 * slot) outside += 1; + max_write = @max(max_write, sge.length); + sum += sge.length; + } + if (w.send_flags & abi.send_signaled != 0) { + queue[q][(head[q] + len[q]) % 512] = sum; + len[q] += 1; + inflight += sum; + bytes += sum; + max_inflight = @max(max_inflight, inflight); + sum = 0; + } + } + return 0; + } + + /// One completion every `lag` polls of a QP's CQ: the HCA finishes writes far slower than they are posted. + pub fn pollCq(cq: *abi.Cq, _: c_int, wc: [*]abi.Wc) callconv(.c) c_int { + const q = cq.handle; + polls += 1; + if (polls % lag != 0 or len[q] == 0) return 0; + inflight -= queue[q][head[q]]; + head[q] = (head[q] + 1) % 512; + len[q] -= 1; + wc[0] = std.mem.zeroes(abi.Wc); + wc[0].status = abi.wc_success; + return 1; + } +}; + +/// Two fake HCAs behind one VerbsWire, connected to a peer whose region has this rank's layout. +pub const Hw = struct { + ctx: [2]abi.Context = undefined, + qp: [2]abi.Qp = undefined, + cq: [2]abi.Cq = undefined, + mr: [2]abi.Mr = undefined, + + pub fn wire(hw: *Hw, region: []u8) rv.VerbsWire { + var w: rv.VerbsWire = .{ .v = undefined, .traffic_class = 0, .recv_off = send_off + 2 * slot, .flag_off = 4096, .send_off = send_off, .slot_bytes = slot, .region = region }; + for (0..2) |i| { + hw.ctx[i].ops = .{ ._before_poll = @splat(null), .poll_cq = Fake.pollCq, ._before_post = @splat(null), .post_send = Fake.postSend, .post_recv = undefined, ._after = @splat(null) }; + hw.qp[i].qp_num = @intCast(i); + hw.cq[i].handle = @intCast(i); + hw.mr[i].lkey = 7; + w.hcas[i] = .{ .ctx = &hw.ctx[i], .pd = undefined, .mr = &hw.mr[i], .cq = &hw.cq[i], .qp = &hw.qp[i], .gid_index = 0, .gid = @splat(0), .mtu = 0, .lid = 0 }; + } + w.n_hca = 2; + return w; + } +}; diff --git a/zig/src/fabric/tp/roce_guard.zig b/zig/src/fabric/tp/roce_guard.zig new file mode 100644 index 0000000000..6851f9a6a9 --- /dev/null +++ b/zig/src/fabric/tp/roce_guard.zig @@ -0,0 +1,63 @@ +//! Every RoCE post's bounds before the verb: each HCA's stripe inside its slot and the registered region. +const std = @import("std"); + +pub const max_hcas = 2; +pub const flag_stride = 128; + +pub const Error = error{OutOfRegion}; + +/// Both ranks' region layout (roce.zig's): byte offsets and the registered length every post must stay inside. +pub const Layout = struct { send: u64, recv: u64, flags: u64, slot: u64, len: u64 }; + +/// One HCA's share of a post: its data stripe in this rank's send slot and the peer's receive slot, and its flag word. +pub const Stripe = struct { local: u64, remote: u64, len: u32, flag: u64 }; + +/// Slot `slot`'s first `padded` bytes striped over `n_hca` HCAs; error.OutOfRegion before any verb past a slot or MR. +pub fn stripes(l: Layout, n_hca: u32, slot: u32, padded: u32, out: *[max_hcas]Stripe) Error!void { + if (n_hca == 0 or n_hca > max_hcas or slot > 1) return error.OutOfRegion; + if (padded == 0 or padded % 16 != 0 or padded > l.slot) return error.OutOfRegion; + const packs = padded / 16; + var offset: u64 = 0; + for (out[0..n_hca], 0..) |*o, i| { + var stripe = packs / n_hca; + if (i < packs % n_hca) stripe += 1; + o.* = .{ .local = l.send + slot * l.slot + offset, .remote = l.recv + slot * l.slot + offset, .len = stripe * 16, .flag = l.flags + (@as(u64, slot) * max_hcas + i) * flag_stride }; + if (!within(o.local, o.len, l.send + slot * l.slot, l.slot) or !within(o.remote, o.len, l.recv + slot * l.slot, l.slot)) return error.OutOfRegion; + if (!within(o.flag, 4, l.flags, l.send - l.flags) or !within(o.local, o.len, 0, l.len) or !within(o.remote, o.len, 0, l.len)) return error.OutOfRegion; + offset += o.len; + } +} + +fn within(at: u64, len: u64, lo: u64, size: u64) bool { + return at >= lo and at + len <= lo + size; +} + +/// roce.zig's layout for a 64 KiB slot: ctrl 4 KiB, flags 4 KiB, send [2][slot], recv [2][slot]. +fn testLayout() Layout { + const slot: u64 = 64 << 10; + return .{ .send = 8192, .recv = 8192 + 2 * slot, .flags = 4096, .slot = slot, .len = 8192 + 4 * slot }; +} + +test "a post striped over two HCAs: the stripes tile the slot's first bytes, each flag on its own word" { + const l = testLayout(); + var s: [max_hcas]Stripe = undefined; + try stripes(l, 2, 1, 48, &s); + try std.testing.expectEqual(Stripe{ .local = l.send + l.slot, .remote = l.recv + l.slot, .len = 32, .flag = l.flags + 2 * flag_stride }, s[0]); + try std.testing.expectEqual(Stripe{ .local = l.send + l.slot + 32, .remote = l.recv + l.slot + 32, .len = 16, .flag = l.flags + 3 * flag_stride }, s[1]); + try stripes(l, 2, 1, @intCast(l.slot), &s); + try std.testing.expectEqual(l.recv + 2 * l.slot, s[1].remote + s[1].len); +} + +test "a post past its slot, a third slot, an unpadded or empty post, or no HCA is refused" { + const l = testLayout(); + var s: [max_hcas]Stripe = undefined; + try std.testing.expectError(error.OutOfRegion, stripes(l, 2, 0, @intCast(l.slot + 16), &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 2, 2, 16, &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 2, 0, 24, &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 2, 0, 0, &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 0, 0, 16, &s)); + try std.testing.expectError(error.OutOfRegion, stripes(l, 3, 0, 16, &s)); + var short = l; + short.len -= 16; + try std.testing.expectError(error.OutOfRegion, stripes(short, 2, 1, @intCast(l.slot), &s)); +} diff --git a/zig/src/fabric/tp/roce_post_test.zig b/zig/src/fabric/tp/roce_post_test.zig new file mode 100644 index 0000000000..fe931e1e2d --- /dev/null +++ b/zig/src/fabric/tp/roce_post_test.zig @@ -0,0 +1,20 @@ +//! VerbsWire.post on fake verbs: a post that would leave its slot or region is refused before any verb. +const std = @import("std"); +const fake = @import("roce_fake.zig"); +const Roce = @import("roce.zig").Roce; +const Fake = fake.Fake; + +test "a post past its slot is refused before any verb; a full slot posts inside the region" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(fake.slot)); + defer std.testing.allocator.free(region); + Fake.region = region; + Fake.reset(1); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + try std.testing.expectError(error.OutOfRegion, w.post(1, 1, fake.slot + 16)); + try std.testing.expectError(error.OutOfRegion, w.post(2, 1, 16)); + try std.testing.expectEqual(@as(u64, 0), Fake.posts); + try w.post(1, 1, fake.slot); + try std.testing.expectEqual(@as(u64, 2), Fake.posts); + try std.testing.expectEqual(@as(u64, 0), Fake.outside); +} diff --git a/zig/src/fabric/tp/roce_shm.zig b/zig/src/fabric/tp/roce_shm.zig new file mode 100644 index 0000000000..3924be3e96 --- /dev/null +++ b/zig/src/fabric/tp/roce_shm.zig @@ -0,0 +1,63 @@ +//! The RoCE protocol's shm wire: both ranks' regions in one shared mapping, posts through the guard's stripes. +const std = @import("std"); +const Bootstrap = @import("bootstrap.zig").Bootstrap; +const sock = @import("sock.zig"); +const guard = @import("roce_guard.zig"); +const posix = std.posix; +const page = std.heap.page_size_min; + +/// A test's view of each post: slot, padded bytes and the stripes the wire is about to move. +pub const PostTap = struct { + ctx: *anyopaque, + seen: *const fn (ctx: *anyopaque, l: guard.Layout, slot: u32, padded: u32, s: []const guard.Stripe) void, +}; + +/// The peer's region is mapped here; a post copies each stripe, then stores its flag with release order. +pub const ShmWire = struct { + peer: []u8, + layout: guard.Layout, + local: []u8, + /// Stripes a post makes, as HCAs would: the verbs wire's split, flags and bounds on one host. + stripes: u32 = 1, + tap: ?*PostTap = null, + + pub fn post(w: *ShmWire, slot: u32, seq: u32, padded: u32) guard.Error!void { + var plan: [guard.max_hcas]guard.Stripe = undefined; + try guard.stripes(w.layout, w.stripes, slot, padded, &plan); + if (w.tap) |t| t.seen(t.ctx, w.layout, slot, padded, plan[0..w.stripes]); + for (plan[0..w.stripes]) |s| { + @memcpy(w.peer[s.remote..][0..s.len], w.local[s.local..][0..s.len]); + const flag: *std.atomic.Value(u32) = @ptrCast(@alignCast(w.peer[s.flag..].ptr)); + flag.store(seq, .release); + } + } +}; + +/// One shared mapping holding both ranks' regions (rank 0 creates and sizes the file before anyone learns its name). +pub fn sharedPair(b: *Bootstrap, len: usize) ![]align(page) u8 { + var name: [64]u8 = undefined; + var name_len: u8 = 0; + var fd: std.c.fd_t = -1; + defer if (fd >= 0) { + _ = std.c.close(fd); + }; + if (b.rank == 0) blk: { + const s = std.fmt.bufPrint(&name, "/dev/shm/tf-tp-roce-{d}-{x}", .{ std.c.getpid(), sock.nowNs() }) catch break :blk; + name[s.len] = 0; + fd = std.c.open(name[0..s.len :0], .{ .ACCMODE = .RDWR, .CREAT = true, .EXCL = true, .CLOEXEC = true }, @as(std.c.mode_t, 0o600)); + if (fd >= 0 and std.c.ftruncate(fd, @intCast(2 * len)) == 0) name_len = @intCast(s.len); + } + // an empty name tells the other rank that rank 0 failed (it must not wait for a name that never comes) + try b.broadcast(std.mem.asBytes(&name_len), std.mem.asBytes(&name_len)); + if (name_len == 0) return error.SharedMemoryUnavailable; + try b.broadcast(name[0..name_len], name[0..name_len]); + name[name_len] = 0; + if (b.rank != 0) { + fd = std.c.open(name[0..name_len :0], .{ .ACCMODE = .RDWR, .CLOEXEC = true }, @as(std.c.mode_t, 0)); + if (fd < 0) return error.SharedMemoryUnavailable; + } + const m = try posix.mmap(null, 2 * len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .SHARED }, fd, 0); + try b.agree("mapped", "RoCE shm wire"); + if (b.rank == 0) _ = std.c.unlink(name[0..name_len :0]); + return m; +} diff --git a/zig/src/fabric/tp/roce_test.zig b/zig/src/fabric/tp/roce_test.zig index 51e9d79133..d8cf3fc03f 100644 --- a/zig/src/fabric/tp/roce_test.zig +++ b/zig/src/fabric/tp/roce_test.zig @@ -2,6 +2,8 @@ const std = @import("std"); const roce = @import("roce.zig"); const rv = @import("roce_verbs.zig"); +const shm = @import("roce_shm.zig"); +const guard = @import("roce_guard.zig"); const bootstrap = @import("bootstrap.zig"); const sock = @import("sock.zig"); const reference = @import("reference.zig"); @@ -21,6 +23,30 @@ const Run = struct { max_bytes: usize = 16 * 1024, /// Variable lengths (`Roce.allGatherV`'s contract): each op's lens[r] bytes of rank r's n, 0 and n included. var_lens: bool = false, + /// The shm wire's stripes (2: the verbs wire's split over two HCAs). + stripes: u32 = 1, + /// Each rank's shm posts as the tap saw them: how many, and how many left their slot or the region. + taps: [2]Tap = .{ .{}, .{} }, +}; + +/// Counts an shm wire's posts and any stripe outside its slot, its flag words or the region. +const Tap = struct { + posts: u32 = 0, + outside: u32 = 0, + tap: shm.PostTap = undefined, + + fn seen(ctx: *anyopaque, l: guard.Layout, slot: u32, padded: u32, s: []const guard.Stripe) void { + const t: *Tap = @ptrCast(@alignCast(ctx)); + t.posts += 1; + var sum: u64 = 0; + for (s) |x| { + sum += x.len; + const in_slot = x.local >= l.send + slot * l.slot and x.local + x.len <= l.send + (slot + 1) * l.slot; + const in_peer = x.remote >= l.recv + slot * l.slot and x.remote + x.len <= l.recv + (slot + 1) * l.slot; + if (!in_slot or !in_peer or x.remote + x.len > l.len or x.flag < l.flags or x.flag + 4 > l.send) t.outside += 1; + } + if (sum != padded) t.outside += 1; + } }; fn rank(run: *Run, me: u32, lfd: sock.Fd) void { @@ -32,7 +58,9 @@ fn body(run: *Run, me: u32, lfd: sock.Fd) ?anyerror { defer b.close(); var spec: [1]rv.HcaSpec = .{.{}}; if (run.wire == .verbs) _ = rv.parseSpecs(run.devs[me], @ptrCast(&spec)) catch |e| return e; - const r = roce.Roce.init(null, &b, "", .{ .wire = run.wire, .hcas = if (run.wire == .verbs) &spec else &.{}, .max_bytes = run.max_bytes, .timeout_ns = run.timeout_ns }) catch |e| return e; + const mine = &run.taps[me]; + mine.tap = .{ .ctx = mine, .seen = Tap.seen }; + const r = roce.Roce.init(null, &b, "", .{ .wire = run.wire, .hcas = if (run.wire == .verbs) &spec else &.{}, .max_bytes = run.max_bytes, .timeout_ns = run.timeout_ns, .shm_stripes = run.stripes, .tap = &mine.tap }) catch |e| return e; defer r.deinit(); var abort_thread: ?std.Thread = null; defer if (abort_thread) |t| t.join(); @@ -94,6 +122,17 @@ test "the RoCE protocol with variable lengths over the shm wire: each rank posts try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); } +test "the shm wire striped as over two HCAs: every post inside its slot and the region, bytes and lengths exact" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000, .var_lens = true, .stripes = 2 }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); + for (run.taps) |t| { + try std.testing.expectEqual(@as(u32, 2000), t.posts); + try std.testing.expectEqual(@as(u32, 0), t.outside); + } +} + test "a peer that stops makes the waiting rank time out, and abort ends a wait at once" { var run: Run = .{ .port = 0, .wire = .shm, .ops = 50, .stop_after = 20, .timeout_ns = 100 * std.time.ns_per_ms }; try twoRanks(&run); diff --git a/zig/src/fabric/tp/roce_verbs.zig b/zig/src/fabric/tp/roce_verbs.zig index a540a4ce02..2502aa0757 100644 --- a/zig/src/fabric/tp/roce_verbs.zig +++ b/zig/src/fabric/tp/roce_verbs.zig @@ -4,15 +4,15 @@ const log = @import("log.zig"); const verbs = @import("verbs"); const abi = verbs.abi; const Verbs = verbs.Verbs; +const guard = @import("roce_guard.zig"); -pub const max_hcas = 2; +pub const max_hcas = guard.max_hcas; /// Work requests a queue pair holds; at most a quarter outstanding (two requests a stripe). const send_depth = 256; -const flag_stride = 128; const link_ethernet = 2; -const record_magic: u32 = 0x5446_5244 + 1; // "TFRD", bumped once for this record's layout +const record_magic: u32 = 0x5446_5244 + 2; // "TFRD", bumped for each change of this record's layout -pub const Error = error{ NoHca, VerbsFailed, PostFailed, Completion, BadRecord }; +pub const Error = error{ NoHca, VerbsFailed, PostFailed, Completion, BadRecord } || guard.Error; /// One HCA as configured: its device name and RoCE v2 GID index (null: detected). pub const HcaSpec = struct { name: [64]u8 = @splat(0), name_len: u8 = 0, gid: ?u8 = null }; @@ -22,6 +22,9 @@ pub const Record = extern struct { magic: u32 = record_magic, n_hca: u32 = 0, region: u64 = 0, + /// The region's slot and registered length: a peer whose layout differs is refused at connect. + slot_bytes: u64 = 0, + region_len: u64 = 0, rkey: [max_hcas]u32 = @splat(0), lid: [max_hcas]u32 = @splat(0), mtu: [max_hcas]u32 = @splat(0), @@ -195,7 +198,7 @@ pub const VerbsWire = struct { } pub fn record(w: *const VerbsWire) Record { - var r: Record = .{ .n_hca = w.n_hca, .region = @intFromPtr(w.region.ptr) }; + var r: Record = .{ .n_hca = w.n_hca, .region = @intFromPtr(w.region.ptr), .slot_bytes = w.slot_bytes, .region_len = w.region.len }; for (w.hcas[0..w.n_hca], 0..) |h, i| { r.rkey[i] = h.mr.rkey; r.lid[i] = h.lid; @@ -209,7 +212,7 @@ pub const VerbsWire = struct { /// The peer HCA for each of ours: the same IPv4 /24 (functions pair by subnet, names may differ), else position. pub fn pairing(mine: Record, peer: Record) ![max_hcas]u32 { if (peer.magic != record_magic) return error.BadRecord; - if (peer.n_hca != mine.n_hca) return error.BadRecord; + if (peer.n_hca != mine.n_hca or peer.slot_bytes != mine.slot_bytes or peer.region_len != mine.region_len) return error.BadRecord; var out: [max_hcas]u32 = undefined; for (0..mine.n_hca) |i| { out[i] = @intCast(i); @@ -242,27 +245,28 @@ pub const VerbsWire = struct { /// Slot `slot`'s first `padded` bytes to the peer, striped over the HCAs, each stripe's `seq` flag on the same QP. pub fn post(w: *VerbsWire, slot: u32, seq: u32, padded: u32) !void { - const packs = padded / 16; - var offset: u64 = 0; + var plan: [max_hcas]guard.Stripe = undefined; + try guard.stripes(w.regionLayout(), w.n_hca, slot, padded, &plan); var seq_copy = seq; - for (w.hcas[0..w.n_hca], 0..) |*h, i| { + for (w.hcas[0..w.n_hca], plan[0..w.n_hca], 0..) |*h, s, i| { while (h.outstanding >= send_depth / 4) try w.reapHca(h); - var stripe = packs / w.n_hca; - if (i < packs % w.n_hca) stripe += 1; - const bytes: u32 = stripe * 16; var flag_sge: abi.Sge = .{ .addr = @intFromPtr(&seq_copy), .length = 4, .lkey = 0 }; - var flag_wr: abi.SendWr = .{ .wr_id = seq, .sg_list = @ptrCast(&flag_sge), .num_sge = 1, .opcode = abi.wr_rdma_write, .send_flags = abi.send_signaled | abi.send_inline, .remote_addr = w.peer_region + w.flag_off + (@as(u64, slot) * max_hcas + i) * flag_stride, .rkey = w.peer_rkey[i] }; - var data_sge: abi.Sge = .{ .addr = @intFromPtr(w.region.ptr) + w.send_off + slot * w.slot_bytes + offset, .length = bytes, .lkey = h.mr.lkey }; - var data_wr: abi.SendWr = .{ .wr_id = seq, .next = &flag_wr, .sg_list = @ptrCast(&data_sge), .num_sge = 1, .opcode = abi.wr_rdma_write, .send_flags = 0, .remote_addr = w.peer_region + w.recv_off + slot * w.slot_bytes + offset, .rkey = w.peer_rkey[i] }; + var flag_wr: abi.SendWr = .{ .wr_id = seq, .sg_list = @ptrCast(&flag_sge), .num_sge = 1, .opcode = abi.wr_rdma_write, .send_flags = abi.send_signaled | abi.send_inline, .remote_addr = w.peer_region + s.flag, .rkey = w.peer_rkey[i] }; + var data_sge: abi.Sge = .{ .addr = @intFromPtr(w.region.ptr) + s.local, .length = s.len, .lkey = h.mr.lkey }; + var data_wr: abi.SendWr = .{ .wr_id = seq, .next = &flag_wr, .sg_list = @ptrCast(&data_sge), .num_sge = 1, .opcode = abi.wr_rdma_write, .send_flags = 0, .remote_addr = w.peer_region + s.remote, .rkey = w.peer_rkey[i] }; var bad: ?*abi.SendWr = null; - const first: *abi.SendWr = if (bytes > 0) &data_wr else &flag_wr; + const first: *abi.SendWr = if (s.len > 0) &data_wr else &flag_wr; if (h.ctx.ops.post_send(h.qp, first, &bad) != 0) return error.PostFailed; h.outstanding += 1; - offset += bytes; } w.posted += 1; } + /// This rank's layout; the MR covers `region` exactly and the peer's record matched it at connect. + pub fn regionLayout(w: *const VerbsWire) guard.Layout { + return .{ .send = w.send_off, .recv = w.recv_off, .flags = w.flag_off, .slot = w.slot_bytes, .len = w.region.len }; + } + fn reapHca(w: *VerbsWire, h: *Hca) !void { var wc: [32]abi.Wc = undefined; const n = h.ctx.ops.poll_cq(h.cq, 32, &wc); @@ -303,4 +307,10 @@ test "HCA lists parse and subnets pair the functions" { try std.testing.expect(ipv4Mapped(sub100)); b.n_hca = 1; try std.testing.expectError(error.BadRecord, VerbsWire.pairing(a, b)); + b.n_hca = 2; + b.slot_bytes = 1 << 20; + try std.testing.expectError(error.BadRecord, VerbsWire.pairing(a, b)); + b.slot_bytes = 0; + b.region_len = 4096; + try std.testing.expectError(error.BadRecord, VerbsWire.pairing(a, b)); } diff --git a/zig/src/fabric/tp/tp.zig b/zig/src/fabric/tp/tp.zig index 3fe6d3fab5..8e5d5b0d34 100644 --- a/zig/src/fabric/tp/tp.zig +++ b/zig/src/fabric/tp/tp.zig @@ -200,7 +200,9 @@ test { _ = roce_verbs; _ = roce_pick; _ = planlink; + _ = @import("roce_guard.zig"); _ = @import("roce_test.zig"); + _ = @import("roce_post_test.zig"); _ = @import("host_test.zig"); _ = @import("fate_test.zig"); } From 968839899608146dab06e8ea7c420cbb274b4942 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:03:59 +0700 Subject: [PATCH 13/43] TP RoCE: no RDMA write past 1 MiB, and boot refuses any setting that would ask for one A post of about 9.6 MB faulted a DGX Spark's HCA. roce_guard.max_write (1 MiB, a constant) is checked in stripes() before the slot bounds, so the verbs wire, the shm wire and the proxy refuse a larger post with error.WriteTooLarge before any verb, whatever the slot size. Roce.init refuses a slot past it. Config.parse fails boot with a clear message when TF_TP_MAILBOX_MAX_KB would make the RoCE slot pass 1 MiB; the value is never clamped. Tests: a 9.57 MB and a 1 MiB + 16 post are refused in a 10 MiB slot and reach no verb, a 1 MiB post goes out; boot refuses TF_TP_MAILBOX_MAX_KB 1028, 2048 and 16384 and keeps 1024. --- zig/src/fabric/tp/config.zig | 7 ++++ zig/src/fabric/tp/roce.zig | 1 + zig/src/fabric/tp/roce_fake.zig | 7 ++-- zig/src/fabric/tp/roce_guard.zig | 12 ++++++- zig/src/fabric/tp/roce_post_test.zig | 48 ++++++++++++++++++++++++++-- 5 files changed, 69 insertions(+), 6 deletions(-) diff --git a/zig/src/fabric/tp/config.zig b/zig/src/fabric/tp/config.zig index ddc0c2a1d6..d823297743 100644 --- a/zig/src/fabric/tp/config.zig +++ b/zig/src/fabric/tp/config.zig @@ -16,6 +16,9 @@ pub const Fallback = enum { nccl, @"error" }; pub const Error = error{BadSetting}; +/// RoCE's per-write cap (roce_guard.max_write): on the roce backend a slot past it fails boot, never clamped. +pub const max_rdma_write: usize = @import("roce_guard.zig").max_write; + /// TF_TP_PLAN_PIN: the CPU the plan threads pin to at their first plan (auto: the highest cpu_capacity). pub const PlanPin = union(enum) { auto, cpu: u32 }; @@ -76,6 +79,10 @@ pub const Config = struct { return error.BadSetting; }; c.mailbox_max_bytes = 1024 * try int(usize, get, "TF_TP_MAILBOX_MAX_KB", 256, 4, 16 * 1024); + if (c.backend == .roce and c.mailbox_max_bytes > max_rdma_write) { + log.warn("TF_TP_MAILBOX_MAX_KB={d}: a RoCE write is capped at {d} KiB; refusing to start", .{ c.mailbox_max_bytes / 1024, max_rdma_write / 1024 }); + return error.BadSetting; + } c.mailbox_timeout_ns = try seconds(get, "TF_TP_MAILBOX_TIMEOUT_S", 120, 0.001, 86400); if (get("TF_TP_FAILFAST")) |v| c.failfast = !(std.mem.eql(u8, v, "0") or std.mem.eql(u8, v, "off")); c.grace_ns = try seconds(get, "TF_TP_FAILFAST_GRACE_S", 1.0, 0, 60); diff --git a/zig/src/fabric/tp/roce.zig b/zig/src/fabric/tp/roce.zig index 7806fd60d5..e3695929f1 100644 --- a/zig/src/fabric/tp/roce.zig +++ b/zig/src/fabric/tp/roce.zig @@ -142,6 +142,7 @@ pub const Roce = struct { /// Both ranks call it after the bootstrap; any rank failing makes all return `error.RoceUnavailable` (collective). pub fn init(d: ?*const cuda.Driver, b: *Bootstrap, image: []const u8, o: Options) !*Roce { if (b.world != 2) return error.Unsupported; + if (o.max_bytes > guard.max_write) return error.WriteTooLarge; // a slot no post could fill in one write const slot = std.mem.alignForward(usize, o.max_bytes, 4096); const len = regionBytes(slot); const r = try std.heap.page_allocator.create(Roce); diff --git a/zig/src/fabric/tp/roce_fake.zig b/zig/src/fabric/tp/roce_fake.zig index e556c1990a..a6b02db77d 100644 --- a/zig/src/fabric/tp/roce_fake.zig +++ b/zig/src/fabric/tp/roce_fake.zig @@ -10,6 +10,8 @@ pub const send_off = 8192; /// What the HCAs hold (posted, not completed); every count comes from the work requests themselves. pub const Fake = struct { pub var region: []const u8 = &.{}; + /// The wire's slot (`Hw.wire`), every data write's bound with the send slots. + pub var slot_bytes: u64 = slot; pub var lag: u32 = 1; pub var polls: u64 = 0; pub var queue: [2][512]u64 = undefined; @@ -44,7 +46,7 @@ pub const Fake = struct { if (w.send_flags & abi.send_inline == 0) { const sge = w.sg_list[0]; const lo = @intFromPtr(region.ptr); - if (sge.addr < lo + send_off or sge.addr + sge.length > lo + send_off + 2 * slot) outside += 1; + if (sge.addr < lo + send_off or sge.addr + sge.length > lo + send_off + 2 * slot_bytes) outside += 1; max_write = @max(max_write, sge.length); sum += sge.length; } @@ -82,7 +84,8 @@ pub const Hw = struct { mr: [2]abi.Mr = undefined, pub fn wire(hw: *Hw, region: []u8) rv.VerbsWire { - var w: rv.VerbsWire = .{ .v = undefined, .traffic_class = 0, .recv_off = send_off + 2 * slot, .flag_off = 4096, .send_off = send_off, .slot_bytes = slot, .region = region }; + const s = Fake.slot_bytes; + var w: rv.VerbsWire = .{ .v = undefined, .traffic_class = 0, .recv_off = send_off + 2 * s, .flag_off = 4096, .send_off = send_off, .slot_bytes = s, .region = region }; for (0..2) |i| { hw.ctx[i].ops = .{ ._before_poll = @splat(null), .poll_cq = Fake.pollCq, ._before_post = @splat(null), .post_send = Fake.postSend, .post_recv = undefined, ._after = @splat(null) }; hw.qp[i].qp_num = @intCast(i); diff --git a/zig/src/fabric/tp/roce_guard.zig b/zig/src/fabric/tp/roce_guard.zig index 6851f9a6a9..8f85247afd 100644 --- a/zig/src/fabric/tp/roce_guard.zig +++ b/zig/src/fabric/tp/roce_guard.zig @@ -4,7 +4,10 @@ const std = @import("std"); pub const max_hcas = 2; pub const flag_stride = 128; -pub const Error = error{OutOfRegion}; +pub const Error = error{ OutOfRegion, WriteTooLarge }; + +/// The largest RDMA write a post may make, whatever the slot (a 9.6 MB post faulted a DGX Spark's HCA); never raised. +pub const max_write: u32 = 1 << 20; /// Both ranks' region layout (roce.zig's): byte offsets and the registered length every post must stay inside. pub const Layout = struct { send: u64, recv: u64, flags: u64, slot: u64, len: u64 }; @@ -15,6 +18,7 @@ pub const Stripe = struct { local: u64, remote: u64, len: u32, flag: u64 }; /// Slot `slot`'s first `padded` bytes striped over `n_hca` HCAs; error.OutOfRegion before any verb past a slot or MR. pub fn stripes(l: Layout, n_hca: u32, slot: u32, padded: u32, out: *[max_hcas]Stripe) Error!void { if (n_hca == 0 or n_hca > max_hcas or slot > 1) return error.OutOfRegion; + if (padded > max_write) return error.WriteTooLarge; if (padded == 0 or padded % 16 != 0 or padded > l.slot) return error.OutOfRegion; const packs = padded / 16; var offset: u64 = 0; @@ -57,6 +61,12 @@ test "a post past its slot, a third slot, an unpadded or empty post, or no HCA i try std.testing.expectError(error.OutOfRegion, stripes(l, 2, 0, 0, &s)); try std.testing.expectError(error.OutOfRegion, stripes(l, 0, 0, 16, &s)); try std.testing.expectError(error.OutOfRegion, stripes(l, 3, 0, 16, &s)); + var big = l; + big.slot = 10 << 20; + big.recv = big.send + 2 * big.slot; + big.len = big.send + 4 * big.slot; + try std.testing.expectError(error.WriteTooLarge, stripes(big, 2, 0, max_write + 16, &s)); + try stripes(big, 2, 0, max_write, &s); var short = l; short.len -= 16; try std.testing.expectError(error.OutOfRegion, stripes(short, 2, 1, @intCast(l.slot), &s)); diff --git a/zig/src/fabric/tp/roce_post_test.zig b/zig/src/fabric/tp/roce_post_test.zig index fe931e1e2d..3e815c2c8f 100644 --- a/zig/src/fabric/tp/roce_post_test.zig +++ b/zig/src/fabric/tp/roce_post_test.zig @@ -2,19 +2,61 @@ const std = @import("std"); const fake = @import("roce_fake.zig"); const Roce = @import("roce.zig").Roce; +const Config = @import("config.zig").Config; const Fake = fake.Fake; test "a post past its slot is refused before any verb; a full slot posts inside the region" { - const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(fake.slot)); + const small: u64 = 64 << 10; + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(small)); defer std.testing.allocator.free(region); Fake.region = region; + Fake.slot_bytes = small; + defer Fake.slot_bytes = fake.slot; Fake.reset(1); var hw: fake.Hw = .{}; var w = hw.wire(region); - try std.testing.expectError(error.OutOfRegion, w.post(1, 1, fake.slot + 16)); + try std.testing.expectError(error.OutOfRegion, w.post(1, 1, small + 16)); try std.testing.expectError(error.OutOfRegion, w.post(2, 1, 16)); try std.testing.expectEqual(@as(u64, 0), Fake.posts); - try w.post(1, 1, fake.slot); + try w.post(1, 1, small); try std.testing.expectEqual(@as(u64, 2), Fake.posts); try std.testing.expectEqual(@as(u64, 0), Fake.outside); } + +test "no RDMA write past 1 MiB, whatever the slot: a 9.57 MB and a 1 MiB + 16 post refused in a 10 MiB slot" { + const big: u64 = 10 << 20; + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(big)); + defer std.testing.allocator.free(region); + Fake.region = region; + Fake.slot_bytes = big; + defer Fake.slot_bytes = fake.slot; + Fake.reset(1); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + try std.testing.expectError(error.WriteTooLarge, w.post(0, 1, 9_568_256)); + try std.testing.expectError(error.WriteTooLarge, w.post(1, 1, (1 << 20) + 16)); + try std.testing.expectEqual(@as(u64, 0), Fake.posts); + try w.post(0, 1, 1 << 20); + try std.testing.expectEqual(@as(u64, 1 << 20), Fake.max_write * 2); +} + +const Env = struct { + var pairs: []const [2][]const u8 = &.{}; + + fn get(name: []const u8) ?[]const u8 { + for (pairs) |p| if (std.mem.eql(u8, p[0], name)) return p[1]; + return null; + } +}; + +test "boot refuses a RoCE slot past 1 MiB (TF_TP_MAILBOX_MAX_KB), never clamping it" { + for ([_][]const u8{ "1028", "2048", "16384" }) |kb| { + Env.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_MAILBOX_MAX_KB", kb } }; + try std.testing.expectError(error.BadSetting, Config.parse(&Env.get)); + } + Env.pairs = &.{ .{ "TF_COMM_BACKEND", "roce" }, .{ "TF_TP_MAILBOX_MAX_KB", "1024" } }; + try std.testing.expectEqual(@as(usize, 1 << 20), (try Config.parse(&Env.get)).mailbox_max_bytes); + // the cap is RoCE's: the same-host mailbox keeps its larger slots + Env.pairs = &.{ .{ "TF_COMM_BACKEND", "mailbox" }, .{ "TF_TP_MAILBOX_MAX_KB", "4096" } }; + try std.testing.expectEqual(@as(usize, 4 << 20), (try Config.parse(&Env.get)).mailbox_max_bytes); +} From 72132f08eca54c36e13dab4fa7dfb853182646fc Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:04:52 +0700 Subject: [PATCH 14/43] TP RoCE: TF_TP_ROCE_CHUNKED_KB gathers strides past the slot in slot-sized ops, at most 2 MiB in flight Default off. allGatherV strides past the slot, up to TF_TP_ROCE_CHUNKED_KB, go RoCE as ceil(n / slot) launches of the tp_roce kernel: op c moves at most a slot of each shard from c x slot (RoceArgs.chunk), and an op empty on both ranks rings no doorbell and leaves the epoch. Every RDMA write keeps the 1 MiB cap. The kernel poisons the transport rather than stage past a slot, and its four failure-record blocks fold into fail_record (same stores and fences). VerbsWire.post waits for completions while the bytes posted and not completed over both HCAs would pass max_inflight (2 MiB, a constant in the post path); a wait that sees no completion for 10 s fails the proxy. Boot refuses TF_TP_ROCE_CHUNKED_KB when the kernel image lacks the tp_roce_chunk_abi marker. The setting is part of the digest both ranks agree on. Tests: chunked strides from 300 KB to a 128 MiB union posted back to back on fake verbs whose completions lag keep at most 2 MiB in flight and no write past 1 MiB, and reach the cap; the wire's own in-flight count and peak match the fake HCAs'. Over the shm wire, 2,000 chunked variable-length gathers of up to four slots arrive bit for bit. --- zig/src/fabric/tp/config.zig | 4 ++ zig/src/fabric/tp/hybrid.zig | 24 +++++-- zig/src/fabric/tp/mailbox.cu | 80 +++++++++++++----------- zig/src/fabric/tp/roce.zig | 48 +++++++++++--- zig/src/fabric/tp/roce_inflight_test.zig | 26 ++++++++ zig/src/fabric/tp/roce_post_test.zig | 42 +++++++++++++ zig/src/fabric/tp/roce_test.zig | 21 ++++++- zig/src/fabric/tp/roce_verbs.zig | 41 +++++++++++- zig/src/fabric/tp/tp.zig | 11 +++- 9 files changed, 243 insertions(+), 54 deletions(-) create mode 100644 zig/src/fabric/tp/roce_inflight_test.zig diff --git a/zig/src/fabric/tp/config.zig b/zig/src/fabric/tp/config.zig index d823297743..442f82ed0e 100644 --- a/zig/src/fabric/tp/config.zig +++ b/zig/src/fabric/tp/config.zig @@ -37,6 +37,8 @@ pub const Config = struct { device: u32 = 0, backend: Backend = .nccl, mailbox_max_bytes: usize = 256 * 1024, + /// TF_TP_ROCE_CHUNKED_KB: `allGatherV` strides up to this go RoCE in ops of at most a slot (0: off). + roce_chunked_bytes: usize = 0, /// A mailbox wait for the peer gives up after this long and poisons the transport (fail-stop). mailbox_timeout_ns: u64 = 120 * std.time.ns_per_s, failfast: bool = true, @@ -83,6 +85,7 @@ pub const Config = struct { log.warn("TF_TP_MAILBOX_MAX_KB={d}: a RoCE write is capped at {d} KiB; refusing to start", .{ c.mailbox_max_bytes / 1024, max_rdma_write / 1024 }); return error.BadSetting; } + c.roce_chunked_bytes = 1024 * try int(usize, get, "TF_TP_ROCE_CHUNKED_KB", 0, 0, 64 * 1024); c.mailbox_timeout_ns = try seconds(get, "TF_TP_MAILBOX_TIMEOUT_S", 120, 0.001, 86400); if (get("TF_TP_FAILFAST")) |v| c.failfast = !(std.mem.eql(u8, v, "0") or std.mem.eql(u8, v, "off")); c.grace_ns = try seconds(get, "TF_TP_FAILFAST_GRACE_S", 1.0, 0, 60); @@ -141,6 +144,7 @@ pub const Config = struct { h.update(std.mem.asBytes(&c.world)); h.update(std.mem.asBytes(&c.backend)); h.update(std.mem.asBytes(&c.mailbox_max_bytes)); + h.update(std.mem.asBytes(&c.roce_chunked_bytes)); h.update(std.mem.asBytes(&c.failfast)); h.update(std.mem.asBytes(&c.roce_wire)); h.update(std.mem.asBytes(&c.roce_fallback)); diff --git a/zig/src/fabric/tp/hybrid.zig b/zig/src/fabric/tp/hybrid.zig index 2cdeb0e4a4..e5f6d72c31 100644 --- a/zig/src/fabric/tp/hybrid.zig +++ b/zig/src/fabric/tp/hybrid.zig @@ -51,11 +51,18 @@ pub const Small = union(enum) { } }; +/// Whether a variable-length gather of `n`-byte strides one op cannot take goes RoCE in slot-sized ops (needs the ABI). +pub fn chunkRoute(chunked_bytes: usize, n: usize, single: bool, abi: bool) bool { + return abi and !single and n > 0 and n <= chunked_bytes; +} + pub const Hybrid = struct { small: Small, nccl: *Nccl, /// Largest shard a rank sends through the small path. max_bytes: usize, + /// Largest stride a variable-length gather takes on RoCE in slot-sized ops (TF_TP_ROCE_CHUNKED_KB; 0: none). + chunked_bytes: usize = 0, /// Every rank's `varFits(16)`, agreed at `Session.open` (the bootstrap's all-gather): `varAgreed`'s answer. var_agreed: bool = false, @@ -117,19 +124,28 @@ pub const Hybrid = struct { fn allGatherV(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, stride: usize, lens: DevicePtr, t: DType, stream: collective.Stream) Error!void { const h = self_(ptr); const n = stride * t.size(); - if (varFits(ptr, n)) return h.small.roce.allGatherV(s, r, n, lens, stream); + if (singleFits(h, n)) return h.small.roce.allGatherV(s, r, n, lens, stream); + if (varFits(ptr, n)) return h.small.roce.allGatherVChunked(s, r, n, lens, stream); return allGather(ptr, s, r, stride, t, stream); } - /// RoCE with its device side, for a stride the small path takes (the route is the size's alone: every rank alike). - pub fn varFits(ptr: *anyopaque, n: usize) bool { - const h = self_(ptr); + fn singleFits(h: *const Hybrid, n: usize) bool { return switch (h.small) { .roce => |p| n > 0 and h.routed(n) and p.d != null, .mailbox => false, }; } + /// RoCE with its device side, in one op or in slot-sized ops (the route is the size's alone: every rank alike). + pub fn varFits(ptr: *anyopaque, n: usize) bool { + const h = self_(ptr); + const single = singleFits(h, n); + return single or switch (h.small) { + .roce => |p| chunkRoute(h.chunked_bytes, n, single, p.d != null and p.chunk_abi), + .mailbox => false, + }; + } + fn varAgreed(ptr: *anyopaque) Error!bool { return self_(ptr).var_agreed; } diff --git a/zig/src/fabric/tp/mailbox.cu b/zig/src/fabric/tp/mailbox.cu index 63d53b2484..843b4f1af2 100644 --- a/zig/src/fabric/tp/mailbox.cu +++ b/zig/src/fabric/tp/mailbox.cu @@ -270,6 +270,21 @@ __device__ __noinline__ bool wait_flag_shared(const uint32_t *flag, uint32_t seq } } +// One failure record a rank (the first failing wait's): its fields, then its valid word, each fenced at system scope. +__device__ __forceinline__ void fail_record(uint32_t *poison, uint32_t *rec, uint32_t seq, uint32_t seen, uint32_t who, + uint32_t why, uint32_t waited) { + if (atom_cas_gpu(poison, 0u, 1u) == 0u) { + st_relaxed_sys(rec + 1, seq); + st_relaxed_sys(rec + 2, seen); + st_relaxed_sys(rec + 3, who); + st_relaxed_sys(rec + 4, why); + st_relaxed_sys(rec + 5, waited); + fence_sc_sys(); + st_relaxed_sys(rec, 1u); + fence_sc_sys(); + } +} + } // namespace tftp using namespace tftp; @@ -313,17 +328,7 @@ extern "C" __global__ void __launch_bounds__(512) tp_mailbox(MailArgs a) { uint32_t seen = 0u, why = 0u, waited = 0u; if (!wait_flag(a.my_flags + flag_at, seq, a.ctrl, a.timeout_ns, &seen, &why, &waited)) { s_bad = 1u; - if (atom_cas_gpu(a.state + 2, 0u, 1u) == 0u) { // one failure record a rank - uint32_t *rec = a.ctrl + kCtrlRank + 8u * a.rank; - st_relaxed_sys(rec + 1, seq); - st_relaxed_sys(rec + 2, seen); - st_relaxed_sys(rec + 3, blockIdx.x); - st_relaxed_sys(rec + 4, why); - st_relaxed_sys(rec + 5, waited); - fence_sc_sys(); - st_relaxed_sys(rec, 1u); - fence_sc_sys(); - } + fail_record(a.state + 2, a.ctrl + kCtrlRank + 8u * a.rank, seq, seen, blockIdx.x, why, waited); } } __syncthreads(); @@ -391,6 +396,7 @@ struct RoceArgs { uint32_t opts; // TF_TP_ROCE_FAST: bit 0 fast, bits 8-15 pollers a flag, bits 16-31 start stagger ns; 0: classic // gather only (nullptr: every shard is nbytes): lens[r] = rank r's bytes, padded to 16 on the wire const int32_t *lens; + uint64_t chunk; // chunked gathers: this op's offset into each shard (at most a slot of each from there; else 0) }; // Rank r's bytes this op: ``nbytes``, or lens[r] capped at nbytes. @@ -411,8 +417,17 @@ __device__ __forceinline__ uint64_t roce_span(const RoceArgs &a, uint64_t n) { return (n + 15u) & ~uint64_t(15u); } +// A shard's bytes past `off`, at most a slot (chunked gathers; off 0 and a shard within the slot: the shard). +__device__ __forceinline__ uint64_t roce_part(const RoceArgs &a, uint64_t shard, uint64_t off) { + const uint64_t rest = shard > off ? shard - off : 0u; + return rest < a.slot_bytes ? rest : a.slot_bytes; +} + } // namespace tftp +// Present in images whose RoceArgs carry ``chunk`` (roce.zig refuses TF_TP_ROCE_CHUNKED_KB without it); never launched. +extern "C" __global__ void tp_roce_chunk_abi() {} + extern "C" __global__ void __launch_bounds__(512) tp_roce(RoceArgs a) { __shared__ uint32_t s_head[2]; // poison, seq __shared__ uint32_t s_bad; @@ -438,12 +453,23 @@ extern "C" __global__ void __launch_bounds__(512) tp_roce(RoceArgs a) { const uint32_t seq = s_head[1]; const uint64_t slot = seq & 1u; const bool var = a.mode == kGather && a.lens != nullptr; - const uint64_t mine = var ? roce_shard(a, a.rank) : a.nbytes; - const uint64_t theirs = var ? roce_shard(a, 1u - a.rank) : a.nbytes; + const uint64_t off = var ? a.chunk : 0u; + const uint64_t mine = var ? roce_part(a, roce_shard(a, a.rank), off) : a.nbytes; + const uint64_t theirs = var ? roce_part(a, roce_shard(a, 1u - a.rank), off) : a.nbytes; const uint64_t padded = var ? (mine < 16u ? 16u : (mine + 15u) & ~uint64_t(15u)) : (a.nbytes + 15u) & ~uint64_t(15u); + if (off != 0u && mine == 0u && theirs == 0u) { + return; // nothing past this offset on either rank: no op, the epoch stays (both ranks skip alike) + } + if (padded > a.slot_bytes) { // never staged past a slot: a refused wait (why 3), the transport poisoned + if (tid == 0 && blockIdx.x == 0) { + fail_record(a.state + 2, a.ctrl + kRoceFailed, seq, 0u, 0u, 3u, 0u); + } + return; + } // 1. stage this block's share into the send slot; the last block to finish rings the doorbell - copy_any(fast, a.send_slots + slot * a.slot_bytes, a.in, var ? roce_span(a, mine) : a.nbytes, index, stride); + copy_any(fast, a.send_slots + slot * a.slot_bytes, a.in + off, var ? roce_span(a, mine) : a.nbytes, index, + stride); __syncthreads(); if (tid == 0 && fast) { // lean doorbell: release RMW per block arrival, the last acquires, then a release store (no SC system fences) @@ -468,7 +494,7 @@ extern "C" __global__ void __launch_bounds__(512) tp_roce(RoceArgs a) { } } if (a.mode == kGather) { - copy_any(fast, a.out + static_cast(a.rank) * a.nbytes, a.in, + copy_any(fast, a.out + static_cast(a.rank) * a.nbytes + off, a.in + off, var ? roce_span(a, mine) : a.nbytes, index, stride); } @@ -486,16 +512,7 @@ extern "C" __global__ void __launch_bounds__(512) tp_roce(RoceArgs a) { if (!wait_flag_shared(a.flags + (slot * kRoceHcas + lane) * kFlagStride, seq, a.ctrl + kRoceAbort, a.timeout_ns, &seen, &why, &waited, &s_done[lane])) { s_bad = 1u; - if (atom_cas_gpu(a.state + 2, 0u, 1u) == 0u) { - st_relaxed_sys(a.ctrl + kRoceFailed + 1, seq); - st_relaxed_sys(a.ctrl + kRoceFailed + 2, seen); - st_relaxed_sys(a.ctrl + kRoceFailed + 3, lane); - st_relaxed_sys(a.ctrl + kRoceFailed + 4, why); - st_relaxed_sys(a.ctrl + kRoceFailed + 5, waited); - fence_sc_sys(); - st_relaxed_sys(a.ctrl + kRoceFailed, 1u); - fence_sc_sys(); - } + fail_record(a.state + 2, a.ctrl + kRoceFailed, seq, seen, lane, why, waited); } } } else if (tid == 0) { @@ -504,16 +521,7 @@ extern "C" __global__ void __launch_bounds__(512) tp_roce(RoceArgs a) { if (!wait_flag(a.flags + (slot * kRoceHcas + h) * kFlagStride, seq, a.ctrl + kRoceAbort, a.timeout_ns, &seen, &why, &waited)) { s_bad = 1u; - if (atom_cas_gpu(a.state + 2, 0u, 1u) == 0u) { - st_relaxed_sys(a.ctrl + kRoceFailed + 1, seq); - st_relaxed_sys(a.ctrl + kRoceFailed + 2, seen); - st_relaxed_sys(a.ctrl + kRoceFailed + 3, h); - st_relaxed_sys(a.ctrl + kRoceFailed + 4, why); - st_relaxed_sys(a.ctrl + kRoceFailed + 5, waited); - fence_sc_sys(); - st_relaxed_sys(a.ctrl + kRoceFailed, 1u); - fence_sc_sys(); - } + fail_record(a.state + 2, a.ctrl + kRoceFailed, seq, seen, h, why, waited); break; } } @@ -525,7 +533,7 @@ extern "C" __global__ void __launch_bounds__(512) tp_roce(RoceArgs a) { const uint8_t *inbox = a.recv_slots + slot * a.slot_bytes; switch (a.mode) { case kGather: - copy_any(fast, a.out + static_cast(1u - a.rank) * a.nbytes, inbox, + copy_any(fast, a.out + static_cast(1u - a.rank) * a.nbytes + off, inbox, var ? roce_span(a, theirs) : a.nbytes, index, stride); break; case kExchange: diff --git a/zig/src/fabric/tp/roce.zig b/zig/src/fabric/tp/roce.zig index e3695929f1..3e7c1e644a 100644 --- a/zig/src/fabric/tp/roce.zig +++ b/zig/src/fabric/tp/roce.zig @@ -53,6 +53,8 @@ const RoceArgs = extern struct { opts: u32 = 0, /// gather only: device int32 [2] of each rank's bytes (0: every shard is `nbytes`) lens: DevicePtr = 0, + /// Chunked gathers: this op's byte offset into each rank's shard (it moves at most a slot from there; 0 otherwise). + chunk: u64 = 0, }; pub const WireKind = enum { @@ -125,6 +127,8 @@ pub const Roce = struct { state: ?cuda.DeviceBuffer = null, module: ?cuda.Module = null, func: ?cuda.Function = null, + /// The kernel image takes `RoceArgs.chunk` (it exports tp_roce_chunk_abi): TF_TP_ROCE_CHUNKED_KB needs it. + chunk_abi: bool = false, // the proxy thread thread: ?std.Thread = null, running: std.atomic.Value(bool) = .init(false), @@ -233,6 +237,7 @@ pub const Roce = struct { try r.state.?.fill8(0, null); r.module = try cuda.Module.load(d, image); r.func = try r.module.?.function("tp_roce"); + r.chunk_abi = if (r.module.?.function("tp_roce_chunk_abi")) |_| true else |_| false; return true; } @@ -311,8 +316,12 @@ pub const Roce = struct { } fn launch(r: *Roce, mode: Mode, in: DevicePtr, out: DevicePtr, n: usize, lens: DevicePtr, stream: collective.Stream) Error!void { + return r.launchAt(mode, in, out, n, lens, 0, stream); + } + + fn launchAt(r: *Roce, mode: Mode, in: DevicePtr, out: DevicePtr, n: usize, lens: DevicePtr, chunk: u64, stream: collective.Stream) Error!void { if (r.aborted.load(.acquire)) return error.Aborted; - if (n > r.slot_bytes) return error.Unsupported; + if (n > r.slot_bytes and (lens == 0 or !r.chunk_abi)) return error.Unsupported; if (n == 0) return; const d = r.d orelse return error.Unsupported; var a: RoceArgs = .{ @@ -331,9 +340,10 @@ pub const Roce = struct { .n_hca = r.wire.hcas(), .opts = r.fast, .lens = lens, + .chunk = chunk, }; var params = [1]?*anyopaque{@ptrCast(&a)}; - const blocks = @min(@import("mailbox.zig").blocksFor(n), r.max_grid); + const blocks = @min(@import("mailbox.zig").blocksFor(@min(n, r.slot_bytes)), r.max_grid); d.check(d.api.cuLaunchKernel(r.func.?.handle, blocks, 1, 1, threads, 1, 1, 0, stream, ¶ms, null), "cuLaunchKernel tp_roce") catch return error.BackendFailed; } @@ -347,6 +357,14 @@ pub const Roce = struct { return r.launch(.gather, s, o, n, lens, stream); } + /// `allGatherV` past the slot: op c moves at most a slot of each shard from c x slot (empty on both ranks: no op). + pub fn allGatherVChunked(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, lens: DevicePtr, stream: collective.Stream) Error!void { + if (lens == 0) return error.Invalid; + if (!r.chunk_abi) return error.Unsupported; + var off: usize = 0; + while (off < n) : (off += r.slot_bytes) try r.launchAt(.gather, s, o, n, lens, off, stream); + } + pub fn exchangeBytes(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, stream: collective.Stream) Error!void { return r.launch(.exchange, s, o, n, 0, stream); } @@ -391,17 +409,31 @@ pub const Roce = struct { /// `hostGather` with variable lengths (`allGatherV`): rank k's first lens[k] bytes land at out + k x n. pub fn hostGatherV(r: *Roce, in: []const u8, out: []u8, lens: ?[2]u32) Error!void { + if (in.len > r.slot_bytes) return error.Invalid; + return r.hostGatherAt(in, out, lens, 0); + } + + /// `allGatherVChunked`'s ops done by the host (the kernel's side of each, as `hostGatherV`). + pub fn hostGatherVChunked(r: *Roce, in: []const u8, out: []u8, lens: [2]u32) Error!void { + var off: usize = 0; + while (off < in.len) : (off += r.slot_bytes) try r.hostGatherAt(in, out, lens, off); + } + + /// One op from shard offset `off`: each rank's bytes past it, at most a slot; an op empty on both ranks is skipped. + fn hostGatherAt(r: *Roce, in: []const u8, out: []u8, lens: ?[2]u32, off: usize) Error!void { const n = in.len; - if (n == 0 or n > r.slot_bytes or out.len < 2 * n) return error.Invalid; + if (n == 0 or out.len < 2 * n) return error.Invalid; if (lens) |l| if (l[0] > n or l[1] > n) return error.Invalid; - const mine: usize = if (lens) |l| l[r.me] else n; - const theirs: usize = if (lens) |l| l[1 - r.me] else n; + const full: [2]usize = if (lens) |l| .{ l[0], l[1] } else .{ n, n }; + const mine: usize = @min(full[r.me] -| off, r.slot_bytes); + const theirs: usize = @min(full[1 - r.me] -| off, r.slot_bytes); + if (off > 0 and mine == 0 and theirs == 0) return; const seq = r.host_epoch +% 1; const slot: usize = seq & 1; - @memcpy(r.region[send_off + slot * r.slot_bytes ..][0..mine], in[0..mine]); + @memcpy(r.region[send_off + slot * r.slot_bytes ..][0..mine], in[off..][0..mine]); r.word(w_slot_bytes + slot).store(@intCast(@max(16, std.mem.alignForward(usize, mine, 16))), .release); r.word(w_seq).store(seq, .release); - @memcpy(out[r.me * n ..][0..mine], in[0..mine]); + @memcpy(out[r.me * n + off ..][0..mine], in[off..][0..mine]); const t0 = sock.nowNs(); for (0..r.wire.hcas()) |h| { const flag: *std.atomic.Value(u32) = @ptrCast(@alignCast(r.region[flag_off + (slot * 2 + h) * flag_stride ..].ptr)); @@ -411,7 +443,7 @@ pub const Roce = struct { std.atomic.spinLoopHint(); } } - @memcpy(out[(1 - r.me) * n ..][0..theirs], r.region[send_off + 2 * r.slot_bytes + slot * r.slot_bytes ..][0..theirs]); + @memcpy(out[(1 - r.me) * n + off ..][0..theirs], r.region[send_off + 2 * r.slot_bytes + slot * r.slot_bytes ..][0..theirs]); r.host_epoch = seq; } }; diff --git a/zig/src/fabric/tp/roce_inflight_test.zig b/zig/src/fabric/tp/roce_inflight_test.zig new file mode 100644 index 0000000000..a8b83a1564 --- /dev/null +++ b/zig/src/fabric/tp/roce_inflight_test.zig @@ -0,0 +1,26 @@ +//! VerbsWire on fake verbs whose completions lag: what it holds in flight and what a failed transport leaves behind. +const std = @import("std"); +const rv = @import("roce_verbs.zig"); +const fake = @import("roce_fake.zig"); +const Roce = @import("roce.zig").Roce; +const Fake = fake.Fake; +const slot = fake.slot; + +test "chunked posts: the wire's own in-flight count and peak match the HCAs', and reach the 2 MiB cap" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(slot)); + defer std.testing.allocator.free(region); + Fake.region = region; + Fake.slot_bytes = slot; + Fake.reset(3); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + for (1..9) |seq| { + try w.post(@intCast(seq & 1), @intCast(seq), slot); + try std.testing.expectEqual(Fake.inflight, w.inflight()); + try std.testing.expect(w.inflight() <= rv.max_inflight); + } + while (w.inflight() > 0) try w.reap(); + try std.testing.expectEqual(@as(u64, 0), Fake.inflight); + try std.testing.expectEqual(Fake.max_inflight, w.peak.load(.monotonic)); + try std.testing.expectEqual(@as(u64, 2 << 20), w.peak.load(.monotonic)); +} diff --git a/zig/src/fabric/tp/roce_post_test.zig b/zig/src/fabric/tp/roce_post_test.zig index 3e815c2c8f..a4f7971fed 100644 --- a/zig/src/fabric/tp/roce_post_test.zig +++ b/zig/src/fabric/tp/roce_post_test.zig @@ -60,3 +60,45 @@ test "boot refuses a RoCE slot past 1 MiB (TF_TP_MAILBOX_MAX_KB), never clamping Env.pairs = &.{ .{ "TF_COMM_BACKEND", "mailbox" }, .{ "TF_TP_MAILBOX_MAX_KB", "4096" } }; try std.testing.expectEqual(@as(usize, 4 << 20), (try Config.parse(&Env.get)).mailbox_max_bytes); } + +test "chunked strides posted back to back: at most 2 MiB in flight over both HCAs and no write past 1 MiB" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(fake.slot)); + defer std.testing.allocator.free(region); + Fake.region = region; + Fake.slot_bytes = fake.slot; + // gathers of R x 512 x 584 B strides and 584 B rows up to a 128 MiB union, each in <= 1 MiB ops + var strides: [16]u64 = undefined; + var n: usize = 0; + for ([_]u64{ 1, 2, 3, 4, 6, 16, 32, 64 }) |r| { + strides[n] = r * 512 * 584; + n += 1; + } + for ([_]u64{ 1795, 1796, (128 << 20) / (3 * 584) }) |m| { + strides[n] = m * 584; + n += 1; + } + for ([_]u32{ 1, 3, 17 }) |lag| { + Fake.reset(lag); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + var seq: u32 = 0; + var want: u64 = 0; + for (strides[0..n]) |stride| { + var off: u64 = 0; + while (off < stride) : (off += fake.slot) { + seq += 1; + const padded: u32 = @intCast(std.mem.alignForward(u64, @min(fake.slot, stride - off), 16)); + try w.post(seq & 1, seq, padded); + want += padded; + } + } + while (Fake.inflight > 0) try w.reap(); + try std.testing.expectEqual(want, Fake.bytes); + try std.testing.expectEqual(@as(u64, 0), Fake.outside); + // the cap's numbers, not roce_verbs.zig's constants: 1 MiB a write, 2 MiB in flight + try std.testing.expect(Fake.max_write <= 1 << 20); + try std.testing.expect(Fake.max_inflight <= 2 << 20); + // the bound is reached, not avoided: posts outran completions + try std.testing.expect(Fake.max_inflight > 1 << 20); + } +} diff --git a/zig/src/fabric/tp/roce_test.zig b/zig/src/fabric/tp/roce_test.zig index d8cf3fc03f..dff1ab7eb9 100644 --- a/zig/src/fabric/tp/roce_test.zig +++ b/zig/src/fabric/tp/roce_test.zig @@ -23,6 +23,8 @@ const Run = struct { max_bytes: usize = 16 * 1024, /// Variable lengths (`Roce.allGatherV`'s contract): each op's lens[r] bytes of rank r's n, 0 and n included. var_lens: bool = false, + /// Strides up to four slots through `hostGatherVChunked` (variable lengths, slot-sized ops). + chunked: bool = false, /// The shm wire's stripes (2: the verbs wire's split over two HCAs). stripes: u32 = 1, /// Each rank's shm posts as the tap saw them: how many, and how many left their slot or the region. @@ -77,12 +79,14 @@ fn body(run: *Run, me: u32, lfd: sock.Fd) ?anyerror { if (run.abort_after_ms) |ms| abort_thread = std.Thread.spawn(.{}, abortLater, .{ r, ms }) catch null; }; // sizes from 1 byte to the slot, odd ones included, both slots and wrap-arounds - const n = 1 + (k * 7919) % run.max_bytes; + const n = 1 + (k * 7919) % (if (run.chunked) 4 * run.max_bytes else run.max_bytes); reference.fill(in[0..n], .u8, me, k); reference.fill(peer[0..n], .u8, 1 - me, k); const lens: [2]u32 = if (k % 5 == 0) .{ 0, @intCast(n) } else .{ @intCast((k * 31) % (n + 1)), @intCast((k * 17 + 5) % (n + 1)) }; - const l: [2]usize = if (run.var_lens) .{ lens[0], lens[1] } else .{ n, n }; - r.hostGatherV(in[0..n], out[0 .. 2 * n], if (run.var_lens) lens else null) catch |e| return e; + const l: [2]usize = if (run.var_lens or run.chunked) .{ lens[0], lens[1] } else .{ n, n }; + if (run.chunked) { + r.hostGatherVChunked(in[0..n], out[0 .. 2 * n], lens) catch |e| return e; + } else r.hostGatherV(in[0..n], out[0 .. 2 * n], if (run.var_lens) lens else null) catch |e| return e; const r0 = if (me == 0) in[0..n] else peer[0..n]; const r1 = if (me == 1) in[0..n] else peer[0..n]; if (!std.mem.eql(u8, out[0..l[0]], r0[0..l[0]]) or !std.mem.eql(u8, out[n..][0..l[1]], r1[0..l[1]])) return error.WrongBytes; @@ -133,6 +137,17 @@ test "the shm wire striped as over two HCAs: every post inside its slot and the } } +test "chunked gathers over the shm wire: strides up to four slots in slot-sized ops, tails empty on both ranks skipped" { + var run: Run = .{ .port = 0, .wire = .shm, .ops = 2000, .chunked = true, .stripes = 2 }; + try twoRanks(&run); + for (run.errs) |e| if (e) |err| return err; + try std.testing.expectEqual([2]u32{ 2000, 2000 }, run.done); + for (run.taps) |t| { + try std.testing.expect(t.posts > 2000); + try std.testing.expectEqual(@as(u32, 0), t.outside); + } +} + test "a peer that stops makes the waiting rank time out, and abort ends a wait at once" { var run: Run = .{ .port = 0, .wire = .shm, .ops = 50, .stop_after = 20, .timeout_ns = 100 * std.time.ns_per_ms }; try twoRanks(&run); diff --git a/zig/src/fabric/tp/roce_verbs.zig b/zig/src/fabric/tp/roce_verbs.zig index 2502aa0757..cfcb2d143d 100644 --- a/zig/src/fabric/tp/roce_verbs.zig +++ b/zig/src/fabric/tp/roce_verbs.zig @@ -5,6 +5,7 @@ const verbs = @import("verbs"); const abi = verbs.abi; const Verbs = verbs.Verbs; const guard = @import("roce_guard.zig"); +const sock = @import("sock.zig"); pub const max_hcas = guard.max_hcas; /// Work requests a queue pair holds; at most a quarter outstanding (two requests a stripe). @@ -14,6 +15,12 @@ const record_magic: u32 = 0x5446_5244 + 2; // "TFRD", bumped for each change of pub const Error = error{ NoHca, VerbsFailed, PostFailed, Completion, BadRecord } || guard.Error; +/// Bytes posted and not yet completed, every HCA together: a post waits for completions past it; never raised. +pub const max_inflight: u64 = 2 * @as(u64, guard.max_write); +/// A post that waits on completions this long fails the proxy (a peer or HCA that stopped answering). +const inflight_wait_ns: u64 = 10 * std.time.ns_per_s; +const ring_len = send_depth / 4; + /// One HCA as configured: its device name and RoCE v2 GID index (null: detected). pub const HcaSpec = struct { name: [64]u8 = @splat(0), name_len: u8 = 0, gid: ?u8 = null }; @@ -43,6 +50,10 @@ const Hca = struct { mtu: u32, lid: u16, outstanding: u32 = 0, + /// Each outstanding post's stripe bytes, oldest at `head` (completions arrive in order). + ring: [ring_len]u32 = undefined, + head: u32 = 0, + inflight: u64 = 0, }; /// Parses `name[:gid],name[:gid]` (TF_TP_ROCE_HCA). @@ -137,6 +148,8 @@ pub const VerbsWire = struct { slot_bytes: usize, region: []u8, posted: u64 = 0, + /// The most stripe bytes in flight at once (posted, not completed), for tests and reports. + peak: std.atomic.Value(u64) = .init(0), completed: u64 = 0, /// Opens the HCAs, registers `region` on each and creates the queue pairs (INIT); `record` then goes to the peer. @@ -247,9 +260,10 @@ pub const VerbsWire = struct { pub fn post(w: *VerbsWire, slot: u32, seq: u32, padded: u32) !void { var plan: [max_hcas]guard.Stripe = undefined; try guard.stripes(w.regionLayout(), w.n_hca, slot, padded, &plan); + while (w.inflight() + padded > max_inflight) try w.reapWait(); var seq_copy = seq; for (w.hcas[0..w.n_hca], plan[0..w.n_hca], 0..) |*h, s, i| { - while (h.outstanding >= send_depth / 4) try w.reapHca(h); + while (h.outstanding >= ring_len) try w.reapHca(h); var flag_sge: abi.Sge = .{ .addr = @intFromPtr(&seq_copy), .length = 4, .lkey = 0 }; var flag_wr: abi.SendWr = .{ .wr_id = seq, .sg_list = @ptrCast(&flag_sge), .num_sge = 1, .opcode = abi.wr_rdma_write, .send_flags = abi.send_signaled | abi.send_inline, .remote_addr = w.peer_region + s.flag, .rkey = w.peer_rkey[i] }; var data_sge: abi.Sge = .{ .addr = @intFromPtr(w.region.ptr) + s.local, .length = s.len, .lkey = h.mr.lkey }; @@ -257,8 +271,11 @@ pub const VerbsWire = struct { var bad: ?*abi.SendWr = null; const first: *abi.SendWr = if (s.len > 0) &data_wr else &flag_wr; if (h.ctx.ops.post_send(h.qp, first, &bad) != 0) return error.PostFailed; + h.ring[(h.head + h.outstanding) % ring_len] = s.len; h.outstanding += 1; + h.inflight += s.len; } + if (w.inflight() > w.peak.load(.monotonic)) w.peak.store(w.inflight(), .monotonic); w.posted += 1; } @@ -275,10 +292,32 @@ pub const VerbsWire = struct { log.err("tp roce: RDMA write of seq {d} failed: {s} (vendor 0x{x})", .{ c.wr_id, w.v.statusText(c.status), c.vendor_err }); return error.Completion; }; + for (0..@intCast(n)) |_| { + h.inflight -= h.ring[h.head]; + h.head = (h.head + 1) % ring_len; + } h.outstanding -= @intCast(n); w.completed += @intCast(n); } + /// Bytes posted and not yet completed over every HCA. + pub fn inflight(w: *const VerbsWire) u64 { + var sum: u64 = 0; + for (w.hcas[0..w.n_hca]) |h| sum += h.inflight; + return sum; + } + + /// Reaps until some post completes (error.Completion when none does within `inflight_wait_ns`). + fn reapWait(w: *VerbsWire) !void { + const before = w.completed; + const t0 = sock.nowNs(); + while (w.completed == before) { + try w.reap(); + if (w.completed == before and sock.nowNs() - t0 > inflight_wait_ns) return error.Completion; + std.atomic.spinLoopHint(); + } + } + /// Reap every HCA's completions (after each op, and in the proxy's idle loop). pub fn reap(w: *VerbsWire) !void { for (w.hcas[0..w.n_hca]) |*h| try w.reapHca(h); diff --git a/zig/src/fabric/tp/tp.zig b/zig/src/fabric/tp/tp.zig index 8e5d5b0d34..901544b58b 100644 --- a/zig/src/fabric/tp/tp.zig +++ b/zig/src/fabric/tp/tp.zig @@ -41,7 +41,7 @@ pub const Session = struct { pub fn open(s: *Session, d: *const cuda.Driver, cfg: Config, kernel_image: []const u8) !void { var b = try Bootstrap.open(cfg); defer b.close(); - try b.agree(&cfg.digest(), "settings (TF_TP_WORLD, TF_COMM_BACKEND, TF_TP_MAILBOX_MAX_KB, TF_TP_FAILFAST, TF_TP_ROCE_WIRE / _FALLBACK)"); + try b.agree(&cfg.digest(), "settings (TF_TP_WORLD, TF_COMM_BACKEND, TF_TP_MAILBOX_MAX_KB, TF_TP_ROCE_CHUNKED_KB, TF_TP_FAILFAST, TF_TP_ROCE_WIRE / _FALLBACK)"); var lib = try cuda.nccl.Library.openPath(cfg.nccl_lib); errdefer lib.close(); const v = try lib.version(); @@ -84,7 +84,13 @@ pub const Session = struct { const opts: roce.Options = .{ .wire = if (cfg.roce_wire == .shm) .shm else .verbs, .hcas = specs[0..n_specs], .traffic_class = cfg.roce_tc, .cpu = cfg.roce_cpu, .max_bytes = cfg.mailbox_max_bytes, .timeout_ns = cfg.mailbox_timeout_ns, .fast = cfg.roce_fast, .blocks = cfg.roce_blocks }; const r = roce.Roce.init(d, b, image, opts) catch |e| return s.roceFallback(e); s.roce = r; - s.route = .{ .small = .{ .roce = r }, .nccl = &s.nccl, .max_bytes = cfg.mailbox_max_bytes }; + if (cfg.roce_chunked_bytes > 0 and !r.chunk_abi) { + log.warn("TF_TP_ROCE_CHUNKED_KB={d}: the RoCE kernel image has no chunked ABI (tp_roce_chunk_abi); refusing to start", .{cfg.roce_chunked_bytes / 1024}); + r.deinit(); + s.roce = null; + return error.ChunkAbiMissing; + } + s.route = .{ .small = .{ .roce = r }, .nccl = &s.nccl, .max_bytes = cfg.mailbox_max_bytes, .chunked_bytes = cfg.roce_chunked_bytes }; const good = probe(s, d) catch false; var blobs: [bootstrap.max_ranks][]const u8 = undefined; var buf: [bootstrap.max_ranks]u8 = undefined; @@ -203,6 +209,7 @@ test { _ = @import("roce_guard.zig"); _ = @import("roce_test.zig"); _ = @import("roce_post_test.zig"); + _ = @import("roce_inflight_test.zig"); _ = @import("host_test.zig"); _ = @import("fate_test.zig"); } From e65583cab6b6eda7789847bf4ab0d2168d2e6b7c Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 13:57:31 +0700 Subject: [PATCH 15/43] Server hooks for a model family's own chat rules. family.zig is an optional vtable a family installs in place of the chat template: prompt rendering, thinking and effort fields, a streaming reply reader that parses its own call markup, and the turn it remembers. Without a family the server runs as before. chat.zig's reply collection moves to chat_generation.zig and openai.zig's run to openai_run.zig (pure moves), so neither file passes 600 lines with the hooks in. --- zig/src/server/anthropic.zig | 3 +- zig/src/server/chat.zig | 350 +++-------------------- zig/src/server/chat_generation.zig | 439 +++++++++++++++++++++++++++++ zig/src/server/family.zig | 135 +++++++++ zig/src/server/openai.zig | 341 ++-------------------- zig/src/server/openai_run.zig | 343 ++++++++++++++++++++++ zig/src/server/prompt.zig | 19 +- zig/src/server/root.zig | 1 + zig/src/server/serve.zig | 16 ++ zig/src/server/server.zig | 13 +- zig/src/server/tokens.zig | 22 +- 11 files changed, 1055 insertions(+), 627 deletions(-) create mode 100644 zig/src/server/chat_generation.zig create mode 100644 zig/src/server/family.zig create mode 100644 zig/src/server/openai_run.zig diff --git a/zig/src/server/anthropic.zig b/zig/src/server/anthropic.zig index bf465e7ac2..a8ec8132f5 100644 --- a/zig/src/server/anthropic.zig +++ b/zig/src/server/anthropic.zig @@ -207,7 +207,8 @@ const Reply = struct { fn text(r: *Reply, kind: []const u8, value: []const u8) E!void { if (value.len == 0) return; if (r.active != null) { - r.pending.append(r.a, .{ .kind = kind, .text = value }) catch return error.Closed; + // held until the open tool_use block closes: copied, ``value`` is a slice of a text that grows by then + r.pending.append(r.a, .{ .kind = kind, .text = r.a.dupe(u8, value) catch return error.Closed }) catch return error.Closed; return; } const key = if (std.mem.eql(u8, kind, "thinking")) "thinking" else "text"; diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 7c702ef432..3183ea6df8 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -12,8 +12,12 @@ const log = @import("log.zig"); const ids = @import("ids.zig"); const clock = @import("clock.zig"); const Server = @import("server.zig").Server; +const family_mod = @import("family.zig"); const chunk_plan = @import("chunk_plan.zig"); const Mailbox = @import("mailbox.zig").Mailbox; +const chat_generation = @import("chat_generation.zig"); +const Shown = chat_generation.Shown; +const Generation = chat_generation.Generation; const Value = json.Value; const Cx = errors.Cx; const Allocator = std.mem.Allocator; @@ -35,6 +39,10 @@ pub const Input = struct { logprobs: ?u8 = null, /// ``/v1/completions``' ``prompt_logprobs``: rows for the prompt with this many alternatives; null: not asked. prompt_logprobs: ?u8 = null, + /// The request body (a family reads fields of its own from it, such as ``response_format``). + body: Value = .null, + /// ``parallel_tool_calls: false``: the reply keeps its first call (a family applies it while streaming). + single_call: bool = false, }; /// A streamed piece: content text (a string) or a delta object (reasoning or tool calls). @@ -62,6 +70,8 @@ pub const Reply = struct { logprobs: ?Value = null, /// ``choices[0].prompt_logprobs`` (vLLM's form) for a request that asked; null otherwise. prompt_logprobs: ?Value = null, + /// A family's parsed calls (OpenAI tool-call objects); null: the route parses the content itself. + calls: ?[]Value = null, }; pub const Failure = error{ Refused, Cancelled, Failed, OutOfMemory }; @@ -88,6 +98,9 @@ pub const Prepared = struct { drafts: bool, stops: fields_mod.Stops, preparing: bool, + images_omitted: u32 = 0, + /// the prompt's prepared images (a family's), held until the reply ends + images: ?family_mod.Images = null, }; /// Render ``input`` and run every check that can refuse it, before anything reaches the client or the engine. @@ -107,8 +120,10 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa if (gone.check()) return error.Cancelled; var thinking = flag(f, "enable_thinking") orelse srv.config.enable_thinking; if (input.prompt != null) thinking = false; - const effort = srv.effortFor(if (f.get("reasoning_effort")) |e| (if (e == .string) e.string else null) else null); + const asked: ?[]const u8 = if (f.get("reasoning_effort")) |e| (if (e == .string) e.string else null) else null; + const effort: ?[]const u8 = if (srv.family) |fam| asked orelse srv.config.reasoning_effort orelse fam.defaultEffort() else srv.effortFor(asked); const rendered = try prompt_mod.prepare(srv, cx, input, thinking, effort); + errdefer if (rendered.images) |im| im.release(); if (gone.check()) return error.Cancelled; if (rendered.ids.len == 0) return cx.refuse("rendered prompt is empty"); const window: i64 = srv.info.context_window; @@ -157,7 +172,7 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa } if (input.logprobs) |count| request.logprobs = try @import("logprobs.zig").admit(srv, cx, thinking, request.think_budget, count); request.prompt_logprobs = input.prompt_logprobs; - return .{ .input = input, .request = request, .prompt_len = rendered.ids.len, .received = received, .thinking = thinking, .effort = effort, .sampled = sampling != null, .drafts = drafts, .stops = stops_opt, .preparing = preparing }; + return .{ .input = input, .request = request, .prompt_len = rendered.ids.len, .received = received, .thinking = thinking, .effort = effort, .sampled = sampling != null, .drafts = drafts, .stops = stops_opt, .preparing = preparing, .images_omitted = rendered.images_omitted, .images = rendered.images }; } /// Submit a prepared request and collect its reply; ``sink`` hears the stream (null: not streamed). @@ -172,6 +187,7 @@ pub fn generate(srv: *Server, cx: *Cx, prepared: Prepared, sink: ?Sink, gone: an const stops_opt = prepared.stops; var preparing = prepared.preparing; defer release(srv, preparing); + defer if (prepared.images) |im| im.release(); // after the engine's finished event (the loop below waits for it) if (request.background) { while (true) { const now = nowNs(io); @@ -183,16 +199,19 @@ pub fn generate(srv: *Server, cx: *Cx, prepared: Prepared, sink: ?Sink, gone: an } if (gone.check()) return error.Cancelled; var stop_hook: StopHook = .{ .srv = srv, .stops = .{ .strings = stops_opt.strings } }; - if (stops_opt.strings.len > 0) request.stop = .{ .ctx = &stop_hook, .check = StopHook.check }; + // a family matches stop strings in the answer (``push``); the generic path in the raw text on the engine's thread + if (stops_opt.strings.len > 0 and srv.family == null) request.stop = .{ .ctx = &stop_hook, .check = StopHook.check }; var box: Mailbox = .{ .io = io, .gpa = srv.gpa }; defer box.deinit(); const id = srv.next_id.fetchAdd(1, .monotonic); if (srv.keepalive) |k| k.begin(); // the GPU is busy: the idle ticker holds its commits defer if (srv.keepalive) |k| k.end(); - srv.engine.submit(id, &request, .{ .ctx = &box, .event = Mailbox.onEvent }) catch |e| return switch (e) { - error.Busy => cx.fail(.capacity, "the engine is busy; retry shortly", .{}), - error.Closed => cx.fail(.other, "the scheduler is closed", .{}), - error.InvalidSpans => cx.fail(.request, "the prompt span layout is invalid", .{}), + srv.engine.submit(id, &request, .{ .ctx = &box, .event = Mailbox.onEvent }) catch |e| { + return switch (e) { + error.Busy => cx.fail(.capacity, "the engine is busy; retry shortly", .{}), + error.Closed => cx.fail(.other, "the scheduler is closed", .{}), + error.InvalidSpans => cx.fail(.request, "the prompt span layout is invalid", .{}), + }; }; release(srv, preparing); // a background request waits only while a foreground one prepares preparing = false; @@ -200,9 +219,11 @@ pub fn generate(srv: *Server, cx: *Cx, prepared: Prepared, sink: ?Sink, gone: an 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); errdefer if (!gen.engine_done) gen.cancel(); // the engine writes to the mailbox until it says finished const result: Failure!Reply = blk: { - if (sink != null and input.tools.len > 0) gen.calls = tool_stream.Streamer.init(a, input.tools) catch |e| break :blk e; + if (srv.family) |fam| { + gen.reader = fam.reader(a, thinking, input.tools, input.single_call) catch |e| break :blk e; + } else if (sink != null and input.tools.len > 0) gen.calls = tool_stream.Streamer.init(a, input.tools) catch |e| break :blk e; gen.loop(gone) catch |e| break :blk e; - break :blk gen.finish(cx, prepared.prompt_len, received, thinking, effort, prepared.sampled, prepared.drafts); + break :blk gen.finish(cx, prepared.prompt_len, received, thinking, effort, prepared.sampled, prepared.drafts, prepared.images_omitted); }; if (result) |_| {} else |e| logEnded(input.id, e, cx.message, prepared.prompt_len, gen.collected.items.len, seconds(nowNs(io) - received)); return result; @@ -230,301 +251,6 @@ pub fn release(srv: *Server, preparing: bool) void { if (preparing) _ = srv.preparing.fetchSub(1, .acq_rel); } -/// The text a stream has sent: Python's ``streamed = visible``, grown by its delta while it only grows. -const Shown = struct { - text: std.ArrayList(u8) = .empty, - chars: usize = 0, - at: ?[*]const u8 = null, // where the last shown text lay; the decode buffer only grows, so that prefix holds - extends: bool = false, - - /// Python's ``now[len(shown):]``; ``stable``: ``now`` lies in the append-only decode buffer. - fn after(s: *Shown, now: []const u8, stable: bool) []const u8 { - const sent = s.text.items; - s.extends = (stable and s.at == now.ptr and now.len >= sent.len) or std.mem.startsWith(u8, now, sent); - if (!s.extends) return reply_text.afterChars(now, s.chars); - if (stable) s.at = now.ptr; // the buffer moved but kept its bytes: the next check is free again - return now[sent.len..]; - } - - fn set(s: *Shown, a: Allocator, now: []const u8, delta: []const u8) Allocator.Error!void { - if (s.extends) { - try s.text.appendSlice(a, delta); - s.chars += reply_text.charCount(delta); - } else { - s.text = .empty; - try s.text.appendSlice(a, now); - s.chars = reply_text.charCount(now); - } - s.at = now.ptr; - } -}; - -/// The token loop and the text it streams. -const Generation = struct { - srv: *Server, - a: Allocator, - box: *Mailbox, - id: api.Id, - reply_id: []const u8, // the id the client gets, which the done line prints - sink: ?Sink, - thinking: bool, - stops: reply_text.Stops, - ignore_eos: bool, - max_tokens: u32, - tools: []const Value, - calls: ?tool_stream.Streamer = null, - collected: std.ArrayList(u32) = .empty, - visible: reply_text.Incremental = .{}, - hidden: std.ArrayList(u8) = .empty, // reused for the answer without its call blocks - streamed: Shown = .{}, // what content streamed, kept apart from the decode buffer that grows under slices - streamed_reasoning: Shown = .{}, - streaming_done: bool = false, - first_ns: ?i96 = null, - last_ns: ?i96 = null, - reason: ?[]const u8 = null, // the server ended the reply (a stop string, the length) before the engine said so - engine_done: bool = false, - consumed: usize = 0, - chunk_index: usize = 0, - lp: ?@import("logprobs.zig").Stream = null, // the request asked for logprobs: its rows, copied with each round - - fn eos(g: *const Generation, t: u32) bool { - return !g.ignore_eos and std.mem.indexOfScalar(u32, g.srv.eos, t) != null; - } - - /// What closes a call the model's end token left open (``tool_parse.closeCall``); nothing when a stop string, the length or ``ignore_eos`` ended the reply instead. - fn closeCall(g: *const Generation, text: []const u8) Allocator.Error![]const u8 { - const t = g.collected.items; - if (g.tools.len == 0 or t.len == 0 or !g.eos(t[t.len - 1])) return ""; - return tool_parse.closeCall(g.a, text, g.tools); - } - - /// The engine's stop check, here: the newest tokens' text holds a stop string. - fn stopHit(g: *Generation) Allocator.Error!bool { - if (g.stops.strings.len == 0) return false; - const t = g.collected.items; - const tail = t[t.len -| g.stops.tail()..]; - const text = try g.srv.text.decode(g.a, tail); - for (g.stops.strings) |s| if (std.mem.indexOf(u8, text, s) != null) return true; - return false; - } - - /// Commits a round's tokens as ``LaneStream.commit`` would, then streams what they add. - fn commit(g: *Generation, chunk: []const u32) Failure!void { - var landed: usize = 0; - for (chunk) |t| { - if (g.reason != null) break; - try g.collected.append(g.a, t); - landed += 1; - if (g.eos(t)) { - g.reason = "stop"; - } else if (try g.stopHit()) { - g.reason = "stop"; - g.srv.engine.cancel(g.id); // the engine ends EOS and length itself; a stop string is the server's - } else if (g.collected.items.len >= g.max_tokens) g.reason = "length"; - } - if (landed == 0) return; - const arrived = nowNs(g.srv.io); - if (g.first_ns == null) g.first_ns = arrived; - g.last_ns = arrived; - const sink = g.sink orelse return; - if (g.streaming_done) return; - var fresh: std.ArrayList(u32) = .empty; - for (chunk[0..landed]) |t| { - if (g.eos(t)) { - g.streaming_done = true; - break; - } - try fresh.append(g.a, t); - } - const all = try g.visible.extend(g.a, g.srv.text, decodeText, fresh.items); - const text = g.stops.visible(all, true); - if (g.visible.pending()) return; // a character still split across tokens - var answer = text; - if (g.thinking) { - const split = try reply_text.splitThinking(g.a, text, false, g.srv.markers); - const piece = g.streamed_reasoning.after(split.reasoning, true); - if (piece.len > 0) { - try g.streamed_reasoning.set(g.a, split.reasoning, piece); - try g.emit(sink, try deltaOf(g.a, "reasoning_content", piece)); - } - answer = split.answer; - } - const shown = if (g.calls != null) try reply_text.hideInto(g.a, &g.hidden, answer, false) else answer; - const vis = reply_text.heldBack(reply_text.streamingVisible(shown), g.calls != null); - const copied = @intFromPtr(vis.ptr) >= @intFromPtr(g.hidden.items.ptr) and @intFromPtr(vis.ptr) <= @intFromPtr(g.hidden.items.ptr) + g.hidden.items.len; - const delta = g.streamed.after(vis, !copied); - if (delta.len > 0) { - try g.streamed.set(g.a, vis, delta); - try g.emit(sink, if (g.lp) |*l| (try l.delta(g.srv, g.a, g.thinking, g.ignore_eos, g.collected.items, delta, g.streamed.text.items.len, false)).? else .{ .string = delta }); - } - if (g.calls) |*c| { - var out: std.ArrayList(Value) = .empty; - try c.feed(answer, &out); // never the reasoning: a call it mentions is not made - for (out.items) |d| try g.emit(sink, d); - } - } - - fn emit(g: *Generation, sink: Sink, delta: Value) Failure!void { - sink.call(sink.ctx, delta) catch { - g.cancel(); - return error.Cancelled; - }; - } - - /// Ends the request and waits for the engine; one it ended unfinished counts as a disconnect, as the Mac scheduler counts it. - fn cancel(g: *Generation) void { - if (!g.engine_done) g.srv.engine.cancel(g.id); - g.drain(); - if (g.box.reason == .cancelled and g.reason == null) g.srv.metrics.disconnected(g.srv.io); - } - - /// Waits for the engine's own end, so the request it holds may be freed. - fn drain(g: *Generation) void { - const m = g.box; - m.mutex.lockUncancelable(m.io); - defer m.mutex.unlock(m.io); - while (!m.finished) m.cond.waitUncancelable(m.io, &m.mutex); - g.engine_done = true; - } - - /// Takes rounds until the reply ends; a client that leaves cancels it. - fn loop(g: *Generation, gone: anytype) Failure!void { - const m = g.box; - while (true) { - m.mutex.lockUncancelable(m.io); - var chunk: ?[]u32 = null; - var ended = false; - if (g.chunk_index < m.chunks.items.len) { - const end = m.chunks.items[g.chunk_index]; - g.chunk_index += 1; - chunk = g.a.dupe(u32, m.tokens.items[g.consumed..end]) catch null; - if (g.lp) |*l| l.rows.appendSlice(g.a, m.rows.items[g.consumed..end]) catch {}; - g.consumed = end; - } else if (m.finished) { - ended = true; - } else { - m.cond.waitTimeout(m.io, &m.mutex, .{ .duration = .{ .raw = .fromMilliseconds(50), .clock = .awake } }) catch {}; - } - m.mutex.unlock(m.io); - if (ended) { - g.engine_done = true; - return; - } - if (gone.check()) { - g.cancel(); - return error.Cancelled; - } - const tokens = chunk orelse continue; - if (g.reason == null) try g.commit(tokens); - } - } - - fn finish(g: *Generation, cx: *Cx, prompt_len: usize, received: i96, thinking: bool, effort: ?[]const u8, exact: bool, drafts: bool) Failure!Reply { - const a = g.a; - const m = g.box; - if (!g.engine_done) g.drain(); - const reason: []const u8 = g.reason orelse switch (m.reason) { - .stop => "stop", - .length => "length", - .cancelled => "cancelled", - .failed => return cx.other(if (m.message.len > 0) try a.dupe(u8, m.message) else "the reply failed"), - }; - const finished_ns = nowNs(g.srv.io); - const content_tokens = if (g.ignore_eos) g.collected.items else reply_text.stripTrailing(g.collected.items, g.srv.eos); - const raw = try g.srv.text.decode(a, content_tokens); - const visible = g.stops.visible(raw, false); - // closed before the think split, so a call ending an unclosed think block is the answer - const close = try g.closeCall(visible); - const text = if (close.len > 0) try std.mem.concat(a, u8, &.{ visible, close }) else visible; - var content: []const u8 = text; - var reasoning: ?[]const u8 = null; - if (g.thinking) { - const split = try reply_text.splitThinking(a, text, true, g.srv.markers); - const r = reply_text.pyStrip(split.reasoning); - reasoning = if (r.len > 0) r else null; - content = split.answer; - } else { - const h = reply_text.parseHarmony(text); - content = h.content; - reasoning = h.reasoning; - } - if (g.sink) |sink| { - const thought = g.streamed_reasoning.text.items; - if (reasoning) |r| if (std.mem.startsWith(u8, r, thought) and r.len > thought.len) - try g.emit(sink, try deltaOf(a, "reasoning_content", r[thought.len..])); - const shown = if (g.calls != null) try reply_text.finishedProse(a, content) else content; - const sent = g.streamed.text.items; - const rest = if (std.mem.startsWith(u8, shown, sent) and shown.len > sent.len) shown[sent.len..] else ""; - if (g.lp) |*l| (if (try l.delta(g.srv, a, g.thinking, g.ignore_eos, g.collected.items, rest, shown.len, true)) |d| try g.emit(sink, d)) else if (rest.len > 0) try g.emit(sink, .{ .string = rest }); - if (g.calls) |*c| if (close.len > 0) { - // the streamer reads the closers as markup the model wrote, so the call ends with the same deltas - var out: std.ArrayList(Value) = .empty; - try c.feed(if (g.thinking) content else text, &out); - for (out.items) |d| try g.emit(sink, d); - }; - } - const prefilled = m.prefilled_ns; - const decode_s: f64 = if (prefilled) |p| @max(0, seconds(finished_ns - p)) else 0; - const decode_tokens = g.collected.items.len -| 1; - const runtime = try json.newObject(a); - try runtime.put(a, "enable_thinking", .{ .bool = thinking }); - try runtime.put(a, "reasoning_effort", if (!thinking) .{ .string = "none" } else if (effort) |e| .{ .string = e } else .null); - try runtime.put(a, "engine", .{ .string = g.srv.info.name }); - try runtime.put(a, "tokens_per_second", .{ .float = if (decode_s > 0) @as(f64, @floatFromInt(decode_tokens)) / decode_s else 0 }); - try runtime.put(a, "seconds", .{ .float = @max(0, decode_s) }); // the reply's decode seconds, as tokens_per_second divides them - try runtime.put(a, "prefill_seconds", if (m.stats.prefill_seconds) |p| .{ .float = p } else .null); - const widths = try a.alloc(Value, m.widths.len); - for (m.widths, widths) |w, *slot| slot.* = try json.intValue(a, w); - try runtime.put(a, "prefill_widths", .{ .array = widths }); - const raised = try a.alloc(Value, m.raised.len); - for (m.raised, raised) |r, *slot| slot.* = .{ .bool = r }; - try runtime.put(a, "prefill_raised", .{ .array = raised }); - try runtime.put(a, "time_to_first_token", if (g.first_ns) |t| .{ .float = seconds(t - received) } else .null); - try runtime.put(a, "sampling", .{ .string = if (exact) "exact" else "greedy" }); - try runtime.put(a, "drafts", .{ .bool = drafts }); - const sha = @import("tokens.zig").tokenSha(g.collected.items); - try runtime.put(a, "token_sha", .{ .string = try a.dupe(u8, &sha) }); - try runtime.put(a, "min_rows", try json.intValue(a, m.stats.min_rows)); - if (m.stats.loop_period) |period| { - const loop_field = try json.newObject(a); - try loop_field.put(a, "period", try json.intValue(a, period)); - try runtime.put(a, "loop", .{ .object = loop_field }); - } - const s = m.stats; - const spec = try json.newObject(a); - try spec.put(a, "rounds", try json.intValue(a, s.rounds)); - try spec.put(a, "drafted", try json.intValue(a, s.drafted)); - try spec.put(a, "accepted", try json.intValue(a, s.accepted)); - try spec.put(a, "acceptance_rate", .{ .float = if (s.drafted > 0) @as(f64, @floatFromInt(s.accepted)) / @as(f64, @floatFromInt(s.drafted)) else 0 }); - try spec.put(a, "tokens_per_round", .{ .float = if (s.rounds > 0) @as(f64, @floatFromInt(g.collected.items.len)) / @as(f64, @floatFromInt(s.rounds)) else 0 }); - if (m.telemetry.len > 0) if ((try json.parse(a, m.telemetry)) == .ok) try spec.put(a, "proposer", (try json.parse(a, m.telemetry)).ok); - const think_end: ?u32 = if (thinking) g.srv.text.tokenId(g.srv.markers.close) else null; - const reply: Reply = .{ - .content = content, - .stop_sequence = g.stops.matched(raw), - .reasoning = reasoning, - .tool_calls_streamed = g.calls != null and g.calls.?.streamed, - .finish_reason = reason, - .prompt_tokens = prompt_len, - .cached_tokens = m.cached orelse 0, - .completion_tokens = g.collected.items.len, - .reasoning_tokens = reply_text.reasoningCount(g.collected.items, think_end), - .runtime = .{ .object = runtime }, - .speculative = .{ .object = spec }, - .thinking = thinking, - .effort = effort, - .logprobs = if (g.lp) |l| try @import("logprobs.zig").value(g.srv, cx, a, g.thinking, content_tokens, l.rows.items) else null, - .prompt_logprobs = if (m.prompt_rows) |r| try @import("logprobs.zig").promptValue(g.srv, a, r) else null, - }; - if (std.mem.eql(u8, reason, "length") and thinking and reply_text.pyStrip(content).len == 0) - log.line("warning: a reply reached max_tokens while still thinking, so its content is empty and its text is all in reasoning_content; raise max_tokens, or send chat_template_kwargs {{\"enable_thinking\": false}} (server: --no-thinking)", .{}); - var cycle_text: [32]u8 = undefined; - const cycle = if (s.loop_period) |period| std.fmt.bufPrint(&cycle_text, " loop=period:{d}", .{period}) catch "" else ""; - log.line("done {s} prompt={d} cached={d} thinking={s} effort={s} tokens={d} sha={s} finish={s}{s} rounds={d} accepted={d}/{d}", .{ g.reply_id, prompt_len, reply.cached_tokens, if (thinking) "True" else "False", if (thinking) effort orelse "none" else "none", g.collected.items.len, sha, reason, cycle, s.rounds, s.accepted, s.drafted }); - return reply; - } -}; - /// The request's stop strings as the engine checks them after each token: the newest tokens' text holds one. const StopHook = struct { srv: *Server, @@ -541,7 +267,21 @@ const StopHook = struct { } }; -fn decodeText(t: anytype, a: Allocator, tokens: []const u32) ![]u8 { +/// What closes a call the model's end token left open (``tool_parse.closeCall``); nothing when a stop string, the length or ``ignore_eos`` ended the reply instead. +pub fn closeCall(g: *const Generation, text: []const u8) Allocator.Error![]const u8 { + const t = g.collected.items; + if (g.tools.len == 0 or t.len == 0 or !g.eos(t[t.len - 1])) return ""; + return tool_parse.closeCall(g.a, text, g.tools); +} + +/// Ends the request and waits for the engine; one it ended unfinished counts as a disconnect, as the Mac scheduler counts it. +pub fn cancel(g: *Generation) void { + if (!g.engine_done) g.srv.engine.cancel(g.id); + g.drain(); + if (g.box.reason == .cancelled and g.reason == null) g.srv.metrics.disconnected(g.srv.io); +} + +pub fn decodeText(t: anytype, a: Allocator, tokens: []const u32) ![]u8 { return t.decode(a, tokens); } diff --git a/zig/src/server/chat_generation.zig b/zig/src/server/chat_generation.zig new file mode 100644 index 0000000000..9c49e3f14a --- /dev/null +++ b/zig/src/server/chat_generation.zig @@ -0,0 +1,439 @@ +//! A submitted reply's collection: the engine's events in, streamed deltas and the finished ``Reply`` out. +const std = @import("std"); +const api = @import("engine_api"); +const json = @import("json.zig"); +const errors = @import("errors.zig"); +const reply_text = @import("reply_text.zig"); +const tool_stream = @import("tool_stream.zig"); +const tool_parse = @import("tool_parse.zig"); +const log = @import("log.zig"); +const clock = @import("clock.zig"); +const Server = @import("server.zig").Server; +const family_mod = @import("family.zig"); +const chunk_plan = @import("chunk_plan.zig"); +const Value = json.Value; +const Cx = errors.Cx; +const Allocator = std.mem.Allocator; +const chat = @import("chat.zig"); +const Sink = chat.Sink; +const Reply = chat.Reply; +const Failure = chat.Failure; +const Mailbox = @import("mailbox.zig").Mailbox; +const decodeText = chat.decodeText; +const deltaOf = chat.deltaOf; +const nowNs = clock.nowNs; +const seconds = clock.seconds; + +/// The text a stream has sent: Python's ``streamed = visible``, grown by its delta while it only grows. +pub const Shown = struct { + text: std.ArrayList(u8) = .empty, + chars: usize = 0, + at: ?[*]const u8 = null, // where the last shown text lay; the decode buffer only grows, so that prefix holds + extends: bool = false, + + /// Python's ``now[len(shown):]``; ``stable``: ``now`` lies in the append-only decode buffer. + fn after(s: *Shown, now: []const u8, stable: bool) []const u8 { + const sent = s.text.items; + s.extends = (stable and s.at == now.ptr and now.len >= sent.len) or std.mem.startsWith(u8, now, sent); + if (!s.extends) return reply_text.afterChars(now, s.chars); + if (stable) s.at = now.ptr; // the buffer moved but kept its bytes: the next check is free again + return now[sent.len..]; + } + + fn set(s: *Shown, a: Allocator, now: []const u8, delta: []const u8) Allocator.Error!void { + if (s.extends) { + try s.text.appendSlice(a, delta); + s.chars += reply_text.charCount(delta); + } else { + s.text = .empty; + try s.text.appendSlice(a, now); + s.chars = reply_text.charCount(now); + } + s.at = now.ptr; + } +}; + +/// The token loop and the text it streams. +pub const Generation = struct { + srv: *Server, + a: Allocator, + box: *Mailbox, + id: api.Id, + reply_id: []const u8, // the id the client gets, which the done line prints + sink: ?Sink, + thinking: bool, + stops: reply_text.Stops, + ignore_eos: bool, + max_tokens: u32, + tools: []const Value, + calls: ?tool_stream.Streamer = null, + collected: std.ArrayList(u32) = .empty, + visible: reply_text.Incremental = .{}, + hidden: std.ArrayList(u8) = .empty, // reused for the answer without its call blocks + streamed: Shown = .{}, // what content streamed, kept apart from the decode buffer that grows under slices + streamed_reasoning: Shown = .{}, + streaming_done: bool = false, + first_ns: ?i96 = null, + last_ns: ?i96 = null, + reason: ?[]const u8 = null, // the server ended the reply (a stop string, the length) before the engine said so + engine_done: bool = false, + consumed: usize = 0, + chunk_index: usize = 0, + lp: ?@import("logprobs.zig").Stream = null, // the request asked for logprobs: its rows, copied with each round + // a family's reply: its reader, the answer so far, how much of it streamed, and a stop string's hit + reader: ?family_mod.Reader = null, + answer: std.ArrayList(u8) = .empty, + answer_sent: usize = 0, + deltas: std.ArrayList(Value) = .empty, + stop_hit: bool = false, + stop_at: usize = 0, + + pub fn eos(g: *const Generation, t: u32) bool { + return !g.ignore_eos and std.mem.indexOfScalar(u32, g.srv.eos, t) != null; + } + + const closeCall = chat.closeCall; // in chat.zig with its doc + + /// The engine's stop check, here: the newest tokens' text holds a stop string. + fn stopHit(g: *Generation) Allocator.Error!bool { + if (g.stops.strings.len == 0) return false; + const t = g.collected.items; + const tail = t[t.len -| g.stops.tail()..]; + const text = try g.srv.text.decode(g.a, tail); + for (g.stops.strings) |s| if (std.mem.indexOf(u8, text, s) != null) return true; + return false; + } + + /// Commits a round's tokens as ``LaneStream.commit`` would, then streams what they add. + fn commit(g: *Generation, chunk: []const u32) Failure!void { + if (g.reader) |r| return g.commitFamily(r, chunk); + var landed: usize = 0; + for (chunk) |t| { + if (g.reason != null) break; + try g.collected.append(g.a, t); + landed += 1; + if (g.eos(t)) { + g.reason = "stop"; + } else if (try g.stopHit()) { + g.reason = "stop"; + g.srv.engine.cancel(g.id); // the engine ends EOS and length itself; a stop string is the server's + } else if (g.collected.items.len >= g.max_tokens) g.reason = "length"; + } + if (landed == 0) return; + const arrived = nowNs(g.srv.io); + if (g.first_ns == null) g.first_ns = arrived; + g.last_ns = arrived; + const sink = g.sink orelse return; + if (g.streaming_done) return; + var fresh: std.ArrayList(u32) = .empty; + for (chunk[0..landed]) |t| { + if (g.eos(t)) { + g.streaming_done = true; + break; + } + try fresh.append(g.a, t); + } + const all = try g.visible.extend(g.a, g.srv.text, decodeText, fresh.items); + const text = g.stops.visible(all, true); + if (g.visible.pending()) return; // a character still split across tokens + var answer = text; + if (g.thinking) { + const split = try reply_text.splitThinking(g.a, text, false, g.srv.markers); + const piece = g.streamed_reasoning.after(split.reasoning, true); + if (piece.len > 0) { + try g.streamed_reasoning.set(g.a, split.reasoning, piece); + try g.emit(sink, try deltaOf(g.a, "reasoning_content", piece)); + } + answer = split.answer; + } + const shown = if (g.calls != null) try reply_text.hideInto(g.a, &g.hidden, answer, false) else answer; + const vis = reply_text.heldBack(reply_text.streamingVisible(shown), g.calls != null); + const copied = @intFromPtr(vis.ptr) >= @intFromPtr(g.hidden.items.ptr) and @intFromPtr(vis.ptr) <= @intFromPtr(g.hidden.items.ptr) + g.hidden.items.len; + const delta = g.streamed.after(vis, !copied); + if (delta.len > 0) { + try g.streamed.set(g.a, vis, delta); + try g.emit(sink, if (g.lp) |*l| (try l.delta(g.srv, g.a, g.thinking, g.ignore_eos, g.collected.items, delta, g.streamed.text.items.len, false)).? else .{ .string = delta }); + } + if (g.calls) |*c| { + var out: std.ArrayList(Value) = .empty; + try c.feed(answer, &out); // never the reasoning: a call it mentions is not made + for (out.items) |d| try g.emit(sink, d); + } + } + + fn emit(g: *Generation, sink: Sink, delta: Value) Failure!void { + sink.call(sink.ctx, delta) catch { + g.cancel(); + return error.Cancelled; + }; + } + + pub const cancel = chat.cancel; // in chat.zig with its doc + + /// Waits for the engine's own end, so the request it holds may be freed. + pub fn drain(g: *Generation) void { + const m = g.box; + m.mutex.lockUncancelable(m.io); + defer m.mutex.unlock(m.io); + while (!m.finished) m.cond.waitUncancelable(m.io, &m.mutex); + g.engine_done = true; + } + + /// Takes rounds until the reply ends; a client that leaves cancels it. + pub fn loop(g: *Generation, gone: anytype) Failure!void { + const m = g.box; + while (true) { + m.mutex.lockUncancelable(m.io); + var chunk: ?[]u32 = null; + var ended = false; + if (g.chunk_index < m.chunks.items.len) { + const end = m.chunks.items[g.chunk_index]; + g.chunk_index += 1; + chunk = g.a.dupe(u32, m.tokens.items[g.consumed..end]) catch null; + if (g.lp) |*l| l.rows.appendSlice(g.a, m.rows.items[g.consumed..end]) catch {}; + g.consumed = end; + } else if (m.finished) { + ended = true; + } else { + m.cond.waitTimeout(m.io, &m.mutex, .{ .duration = .{ .raw = .fromMilliseconds(50), .clock = .awake } }) catch {}; + } + m.mutex.unlock(m.io); + if (ended) { + g.engine_done = true; + return; + } + if (gone.check()) { + g.cancel(); + return error.Cancelled; + } + const tokens = chunk orelse continue; + if (g.reason == null) try g.commit(tokens); + } + } + + /// A family's round: committed tokens through the family's reader, its deltas pushed (an EOS id ends, unshown). + fn commitFamily(g: *Generation, r: family_mod.Reader, chunk: []const u32) Failure!void { + var landed: usize = 0; + const start = g.collected.items.len; + for (chunk) |t| { + if (g.reason != null) break; + try g.collected.append(g.a, t); + landed += 1; + if (g.eos(t)) { + g.reason = "stop"; + } else if (g.collected.items.len >= g.max_tokens) g.reason = "length"; + } + if (landed == 0) return; + const arrived = nowNs(g.srv.io); + if (g.first_ns == null) g.first_ns = arrived; + g.last_ns = arrived; + for (g.collected.items[start..]) |t| if (!g.isEos(t)) try r.push(&.{t}); + if (r.pending()) return; // a character still split across tokens + g.deltas.clearRetainingCapacity(); + try r.feed(false, &g.deltas); + try g.pushFamily(g.deltas.items, false); + } + + /// Any end-of-sentence id, ``ignore_eos`` or not: never text. + fn isEos(g: *const Generation, t: u32) bool { + return std.mem.indexOfScalar(u32, g.srv.eos, t) != null; + } + + /// ``push``: reasoning and calls go out as they come; the answer is held where a stop may begin, cut where it does. + fn pushFamily(g: *Generation, deltas: []const Value, finished: bool) Failure!void { + for (deltas) |d| switch (d) { + .string => |t| try g.answer.appendSlice(g.a, t), + else => if (g.sink) |sink| if (!g.streaming_done or finished) try g.emit(sink, d), + }; + var full = g.answer.items; + const stops = g.stops.strings; + if (stops.len > 0 and !g.stop_hit) { + var longest: usize = 0; + for (stops) |st| longest = @max(longest, st.len); + const from = (g.answer_sent + 1) -| longest; + var cut: ?usize = null; + for (stops) |st| if (std.mem.indexOfPos(u8, full, @min(from, full.len), st)) |at| { + if (cut == null or at < cut.?) cut = at; + }; + if (cut) |c| { + g.answer.shrinkRetainingCapacity(c); + full = g.answer.items; + g.stop_hit = true; + g.stop_at = g.collected.items.len; + if (g.reason == null) { + g.reason = "stop"; + g.srv.engine.cancel(g.id); // the engine ends EOS and length itself; a stop string is the server's + } + } + } + var upto = full.len; + if (!finished and !g.stop_hit) { + var hold: usize = 0; + for (stops) |st| { + var k = @min(st.len - 1, full.len); + while (k > 0) : (k -= 1) if (std.mem.endsWith(u8, full, st[0..k])) break; + hold = @max(hold, k); + } + upto -= hold; + } + if (upto > g.answer_sent) { + if (g.sink) |sink| try g.emit(sink, .{ .string = full[g.answer_sent..upto] }); + g.answer_sent = upto; + } + } + + /// A family's finished reply: the stream's last deltas, then the parts (``app._run``'s end). + fn finishFamily(g: *Generation, r: family_mod.Reader) Failure!struct { content: []const u8, reasoning: ?[]const u8, calls: []Value, reason: []const u8, used: usize } { + const streaming = g.sink != null; + g.deltas.clearRetainingCapacity(); + try r.feed(true, &g.deltas); + try g.pushFamily(g.deltas.items, true); + const parsed = try r.parse(); + var content: []const u8 = undefined; + var calls: []Value = undefined; + if (streaming) { + calls = try r.calls(); + content = g.answer.items[0..g.answer_sent]; + } else { + calls = parsed.calls; + content = parsed.content; + for (g.stops.strings) |st| if (std.mem.indexOf(u8, content, st)) |at| { + content = content[0..at]; + if (!g.stop_hit) { + g.stop_hit = true; + g.stop_at = g.collected.items.len; + } + }; + } + const t = g.collected.items; + const reason: []const u8 = if (calls.len > 0) "tool_calls" else if (g.stop_hit or (!g.ignore_eos and t.len > 0 and g.isEos(t[t.len - 1]))) "stop" else "length"; + if (g.srv.family) |fam| fam.remember(parsed.reasoning, content, calls); + // tools offered, no call, but tool markup: log up to 300 bytes of what the parser saw (UTF-8 cut, escaped) + if (calls.len == 0 and g.tools.len > 0) if (r.markup() catch null) |m| + std.log.warn("tool markup with no call: request {s}, finish {s}: {s}; reply bytes {d}..: {f}", .{ g.reply_id, reason, m.reason, m.at, std.zig.fmtString(m.window) }); + return .{ .content = content, .reasoning = if (parsed.reasoning.len > 0) parsed.reasoning else null, .calls = calls, .reason = reason, .used = if (g.stop_hit and g.stop_at > 0) g.stop_at else t.len }; + } + + pub fn finish(g: *Generation, cx: *Cx, prompt_len: usize, received: i96, thinking: bool, effort: ?[]const u8, exact: bool, drafts: bool, images_omitted: u32) Failure!Reply { + const a = g.a; + const m = g.box; + if (!g.engine_done) g.drain(); + const reason_generic: []const u8 = g.reason orelse switch (m.reason) { + .stop => "stop", + .length => "length", + .cancelled => "cancelled", + .failed => return cx.other(if (m.message.len > 0) try a.dupe(u8, m.message) else "the reply failed"), + }; + var reason = reason_generic; + var raw: []const u8 = ""; + var content: []const u8 = undefined; + var reasoning: ?[]const u8 = null; + var calls: ?[]Value = null; + var used = g.collected.items.len; + var logprobs: ?Value = null; + if (g.reader) |r| { + const f = try g.finishFamily(r); + content = f.content; + reasoning = f.reasoning; + calls = f.calls; + used = f.used; + if (!(m.reason == .cancelled and g.reason == null)) reason = f.reason; + } else { + const content_tokens = if (g.ignore_eos) g.collected.items else reply_text.stripTrailing(g.collected.items, g.srv.eos); + raw = try g.srv.text.decode(a, content_tokens); + if (g.lp) |l| logprobs = try @import("logprobs.zig").value(g.srv, cx, a, g.thinking, content_tokens, l.rows.items); + const visible = g.stops.visible(raw, false); + // closed before the think split, so a call ending an unclosed think block is the answer + const close = try g.closeCall(visible); + const text = if (close.len > 0) try std.mem.concat(a, u8, &.{ visible, close }) else visible; + content = text; + if (g.thinking) { + const split = try reply_text.splitThinking(a, text, true, g.srv.markers); + const r = reply_text.pyStrip(split.reasoning); + reasoning = if (r.len > 0) r else null; + content = split.answer; + } else { + const h = reply_text.parseHarmony(text); + content = h.content; + reasoning = h.reasoning; + } + if (g.sink) |sink| { + const thought = g.streamed_reasoning.text.items; + if (reasoning) |r| if (std.mem.startsWith(u8, r, thought) and r.len > thought.len) + try g.emit(sink, try deltaOf(a, "reasoning_content", r[thought.len..])); + const shown = if (g.calls != null) try reply_text.finishedProse(a, content) else content; + const sent = g.streamed.text.items; + const rest = if (std.mem.startsWith(u8, shown, sent) and shown.len > sent.len) shown[sent.len..] else ""; + if (g.lp) |*l| (if (try l.delta(g.srv, a, g.thinking, g.ignore_eos, g.collected.items, rest, shown.len, true)) |d| try g.emit(sink, d)) else if (rest.len > 0) try g.emit(sink, .{ .string = rest }); + if (g.calls) |*c| if (close.len > 0) { + // the streamer reads the closers as markup the model wrote, so the call ends with the same deltas + var out: std.ArrayList(Value) = .empty; + try c.feed(if (g.thinking) content else text, &out); + for (out.items) |d| try g.emit(sink, d); + }; + } + } + const finished_ns = nowNs(g.srv.io); + const prefilled = m.prefilled_ns; + const decode_s: f64 = if (prefilled) |p| @max(0, seconds(finished_ns - p)) else 0; + const decode_tokens = g.collected.items.len -| 1; + const runtime = try json.newObject(a); + try runtime.put(a, "enable_thinking", .{ .bool = thinking }); + try runtime.put(a, "reasoning_effort", if (!thinking) .{ .string = "none" } else if (effort) |e| .{ .string = e } else .null); + try runtime.put(a, "engine", .{ .string = g.srv.info.name }); + try runtime.put(a, "tokens_per_second", .{ .float = if (decode_s > 0) @as(f64, @floatFromInt(decode_tokens)) / decode_s else 0 }); + try runtime.put(a, "seconds", .{ .float = @max(0, decode_s) }); // the reply's decode seconds, as tokens_per_second divides them + try runtime.put(a, "prefill_seconds", if (m.stats.prefill_seconds) |p| .{ .float = p } else .null); + const widths = try a.alloc(Value, m.widths.len); + for (m.widths, widths) |w, *slot| slot.* = try json.intValue(a, w); + try runtime.put(a, "prefill_widths", .{ .array = widths }); + const raised = try a.alloc(Value, m.raised.len); + for (m.raised, raised) |r, *slot| slot.* = .{ .bool = r }; + try runtime.put(a, "prefill_raised", .{ .array = raised }); + try runtime.put(a, "time_to_first_token", if (g.first_ns) |t| .{ .float = seconds(t - received) } else .null); + try runtime.put(a, "sampling", .{ .string = if (exact) "exact" else "greedy" }); + try runtime.put(a, "drafts", .{ .bool = drafts }); + const sha = @import("tokens.zig").tokenSha(g.collected.items); + try runtime.put(a, "token_sha", .{ .string = try a.dupe(u8, &sha) }); + if (images_omitted > 0) try runtime.put(a, "images_omitted", try json.intValue(a, images_omitted)); + try runtime.put(a, "min_rows", try json.intValue(a, m.stats.min_rows)); + if (m.stats.loop_period) |period| { + const loop_field = try json.newObject(a); + try loop_field.put(a, "period", try json.intValue(a, period)); + try runtime.put(a, "loop", .{ .object = loop_field }); + } + const s = m.stats; + const spec = try json.newObject(a); + try spec.put(a, "rounds", try json.intValue(a, s.rounds)); + try spec.put(a, "drafted", try json.intValue(a, s.drafted)); + try spec.put(a, "accepted", try json.intValue(a, s.accepted)); + try spec.put(a, "acceptance_rate", .{ .float = if (s.drafted > 0) @as(f64, @floatFromInt(s.accepted)) / @as(f64, @floatFromInt(s.drafted)) else 0 }); + try spec.put(a, "tokens_per_round", .{ .float = if (s.rounds > 0) @as(f64, @floatFromInt(g.collected.items.len)) / @as(f64, @floatFromInt(s.rounds)) else 0 }); + if (m.telemetry.len > 0) if ((try json.parse(a, m.telemetry)) == .ok) try spec.put(a, "proposer", (try json.parse(a, m.telemetry)).ok); + const think_end: ?u32 = if (thinking) g.srv.text.tokenId(g.srv.markers.close) else null; + const reply: Reply = .{ + .content = content, + .stop_sequence = g.stops.matched(raw), + .reasoning = reasoning, + .tool_calls_streamed = g.calls != null and g.calls.?.streamed, + .finish_reason = reason, + .prompt_tokens = prompt_len, + .cached_tokens = m.cached orelse 0, + .completion_tokens = used, + .reasoning_tokens = reply_text.reasoningCount(g.collected.items, think_end), + .runtime = .{ .object = runtime }, + .speculative = .{ .object = spec }, + .thinking = thinking, + .effort = effort, + .logprobs = logprobs, + .prompt_logprobs = if (m.prompt_rows) |r| try @import("logprobs.zig").promptValue(g.srv, a, r) else null, + .calls = calls, + }; + if (std.mem.eql(u8, reason, "length") and thinking and reply_text.pyStrip(content).len == 0) + log.line("warning: a reply reached max_tokens while still thinking, so its content is empty and its text is all in reasoning_content; raise max_tokens, or send chat_template_kwargs {{\"enable_thinking\": false}} (server: --no-thinking)", .{}); + var cycle_text: [32]u8 = undefined; + const cycle = if (s.loop_period) |period| std.fmt.bufPrint(&cycle_text, " loop=period:{d}", .{period}) catch "" else ""; + log.line("done {s} prompt={d} cached={d} thinking={s} effort={s} tokens={d} sha={s} finish={s}{s} rounds={d} accepted={d}/{d}", .{ g.reply_id, prompt_len, reply.cached_tokens, if (thinking) "True" else "False", if (thinking) effort orelse "none" else "none", g.collected.items.len, sha, reason, cycle, s.rounds, s.accepted, s.drafted }); + return reply; + } +}; diff --git a/zig/src/server/family.zig b/zig/src/server/family.zig new file mode 100644 index 0000000000..9e24c8dab4 --- /dev/null +++ b/zig/src/server/family.zig @@ -0,0 +1,135 @@ +//! A model family's own chat rules (prompt encoding, effort budgets, call markup) in place of the chat template. +const std = @import("std"); +const json = @import("json.zig"); +const errors = @import("errors.zig"); +const Value = json.Value; +const Cx = errors.Cx; +const Allocator = std.mem.Allocator; + +/// A request's thinking switch and effort as it asked for them (null: the server's default). +pub const Thinking = struct { enable: ?bool = null, effort: ?[]const u8 = null }; + +/// What a chat prompt is rendered from: the request's own messages and tools, as the client sent them. +pub const Prompt = struct { + messages: Value, + tools: []const Value, + thinking: bool, + /// The effort label (``Family.thinking``'s, else the default). + effort: []const u8, + generation: bool = true, + /// The request body, for fields the family reads itself (``response_format``). + body: Value = .null, +}; + +pub const Rendered = struct { ids: []const u32, images_omitted: u32 = 0, images: ?Images = null }; + +/// A prompt's prepared images: ``held`` lives until ``release``; ``token`` for /tokenize. +pub const Images = struct { + held: *const anyopaque, + token: u32, + ctx: *anyopaque, + release_fn: *const fn (ctx: *anyopaque) void, + + pub fn release(im: Images) void { + im.release_fn(im.ctx); + } +}; + +/// A finished reply's parts, calls as OpenAI tool-call objects. +pub const Parsed = struct { reasoning: []const u8, content: []const u8, calls: []Value }; + +/// One reply as its tokens arrive: the text so far, and the deltas it adds. +pub const Reader = struct { + ctx: *anyopaque, + vt: *const VTable, + + pub const VTable = struct { + /// Tokens committed (end-of-sentence ids already left out). + push: *const fn (ctx: *anyopaque, ids: []const u32) Allocator.Error!void, + /// A character is split across tokens: nothing new is readable yet. + pending: *const fn (ctx: *anyopaque) bool, + /// Deltas since the last feed (content a string, reasoning and calls delta objects); ``finished`` flushes. + feed: *const fn (ctx: *anyopaque, finished: bool, out: *std.ArrayList(Value)) Allocator.Error!void, + /// The whole reply parsed (after the last ``push``; a split character at the end reads as U+FFFD). + parse: *const fn (ctx: *anyopaque) Allocator.Error!Parsed, + /// The calls the stream sent, in order, as OpenAI tool-call objects. + calls: *const fn (ctx: *anyopaque) Allocator.Error![]Value, + /// The finished reply's tool markup that yielded no call (a reason, at most 300 bytes) for the log; null: none + markup: ?*const fn (ctx: *anyopaque) Allocator.Error!?Markup = null, + }; + + pub fn push(r: Reader, ids: []const u32) Allocator.Error!void { + return r.vt.push(r.ctx, ids); + } + pub fn pending(r: Reader) bool { + return r.vt.pending(r.ctx); + } + pub fn feed(r: Reader, finished: bool, out: *std.ArrayList(Value)) Allocator.Error!void { + return r.vt.feed(r.ctx, finished, out); + } + pub fn parse(r: Reader) Allocator.Error!Parsed { + return r.vt.parse(r.ctx); + } + pub fn calls(r: Reader) Allocator.Error![]Value { + return r.vt.calls(r.ctx); + } + pub fn markup(r: Reader) Allocator.Error!?Markup { + const f = r.vt.markup orelse return null; + return f(r.ctx); + } +}; + +/// A reply's tool markup that the parser turned into no call (``Reader.markup``). +pub const Markup = struct { at: usize, reason: []const u8, window: []const u8 }; + +/// Where a reply's reasoning goes: ``reasoning_content`` (OpenAI's), ``reasoning`` (vLLM's), or both. +pub const ReasoningFields = enum { both, reasoning_content, reasoning }; + +pub const Family = struct { + ctx: *anyopaque, + vt: *const VTable, + + pub const VTable = struct { + name: []const u8, + /// ``chat_template_kwargs`` and ``reasoning_effort`` read; a bad value refuses the request. + thinking: *const fn (ctx: *anyopaque, cx: *Cx, body: Value) errors.Refused!Thinking, + /// The effort a request that names none gets. + default_effort: *const fn (ctx: *anyopaque) []const u8, + /// The tools the prompt offers: ``tool_choice`` none offers none; otherwise all of them. + tools: *const fn (ctx: *anyopaque, cx: *Cx, body: Value) errors.Refused![]const Value, + render: *const fn (ctx: *anyopaque, cx: *Cx, p: Prompt) errors.Refused!Rendered, + reader: *const fn (ctx: *anyopaque, a: Allocator, thinking: bool, tools: []const Value, single: bool) Allocator.Error!Reader, + /// A finished reply's reasoning, kept for a later request whose history comes back without it. + remember: *const fn (ctx: *anyopaque, reasoning: []const u8, content: []const u8, calls: []const Value) void, + reasoning_fields: ReasoningFields = .both, + /// Whether the request's ``min_p`` reaches the sampler (some families ignore it). + reads_min_p: bool = true, + }; + + pub fn thinking(f: Family, cx: *Cx, body: Value) errors.Refused!Thinking { + return f.vt.thinking(f.ctx, cx, body); + } + pub fn defaultEffort(f: Family) []const u8 { + return f.vt.default_effort(f.ctx); + } + pub fn tools(f: Family, cx: *Cx, body: Value) errors.Refused![]const Value { + return f.vt.tools(f.ctx, cx, body); + } + pub fn render(f: Family, cx: *Cx, p: Prompt) errors.Refused!Rendered { + return f.vt.render(f.ctx, cx, p); + } + pub fn reader(f: Family, a: Allocator, think: bool, offered: []const Value, single: bool) Allocator.Error!Reader { + return f.vt.reader(f.ctx, a, think, offered, single); + } + pub fn remember(f: Family, reasoning: []const u8, content: []const u8, calls: []const Value) void { + f.vt.remember(f.ctx, reasoning, content, calls); + } + + /// A reasoning delta under the family's field names. + pub fn reasoningDelta(f: Family, a: Allocator, text: []const u8) Allocator.Error!Value { + const o = try json.newObject(a); + if (f.vt.reasoning_fields != .reasoning) try o.put(a, "reasoning_content", .{ .string = text }); + if (f.vt.reasoning_fields != .reasoning_content) try o.put(a, "reasoning", .{ .string = text }); + return .{ .object = o }; + } +}; diff --git a/zig/src/server/openai.zig b/zig/src/server/openai.zig index b8782aaf96..41b0ce7fa1 100644 --- a/zig/src/server/openai.zig +++ b/zig/src/server/openai.zig @@ -20,6 +20,8 @@ const warm = @import("warm.zig"); const cortheon = @import("cortheon/cortheon.zig"); const Server = @import("server.zig").Server; const Conn = @import("http_conn.zig").Conn; +const openai_run = @import("openai_run.zig"); +const Run = openai_run.Run; const Value = json.Value; const Cx = errors.Cx; const Allocator = std.mem.Allocator; @@ -108,7 +110,7 @@ pub fn replyModel(srv: *const Server, body: Value) []const u8 { return srv.config.served_name; } -const Plan = struct { +pub const Plan = struct { input: chat.Input, stream: bool, separate_usage: bool, @@ -125,11 +127,18 @@ fn plan(srv: *Server, cx: *Cx, is_chat: bool, raw: Value) errors.Refused!Plan { const exact_bytes = asked and srv.info.logprobs and (srv.text.tokenBytes(cx.a, 0) catch return error.OutOfMemory) != null; const logprobs = if (is_chat) try fields.probabilityOptions(cx, body, exact_bytes) else try fields.legacyLogprobs(cx, body, exact_bytes); if (srv.config.request_log) |path| request_log.append(cx.a, path, body); - var input: chat.Input = .{ .fields = undefined }; - if (is_chat) { + var input: chat.Input = .{ .fields = undefined, .body = body }; + if (srv.family) |fam| if (is_chat) { + // a family renders the client's messages and tools itself (its encoding reads OpenAI's shapes) + const m = body.get("messages") orelse Value.null; + if (m != .array or m.array.len == 0) return cx.refuse("messages must be a non-empty list"); + input.messages = m; + input.tools = try fam.tools(cx, body); + }; + if (is_chat and srv.family == null) { input.messages = try messages.normalize(cx, body.get("messages"), "system", srv.needs_user_after_tool); input.tools = try tool_specs.active(cx, body.get("tools"), body.get("tool_choice")); - } else if (body.get("messages")) |m| if (m == .array and m.array.len > 0) { + } else if (!is_chat) if (body.get("messages")) |m| if (m == .array and m.array.len > 0) { input.messages = try messages.normalize(cx, m, "system", srv.needs_user_after_tool); }; if (!is_chat and input.messages.array.len == 0) input.prompt = try legacyPrompt(srv, cx, body.get("prompt")); @@ -147,9 +156,12 @@ fn plan(srv: *Server, cx: *Cx, is_chat: bool, raw: Value) errors.Refused!Plan { const sampling = try json.newObject(cx.a); for ([_][]const u8{ "temperature", "top_p", "top_k", "min_p", "seed", "priority", "draft", "thinking_budget", "ignore_eos", "stop" } ++ grammar.fields) |k| if (body.get(k)) |v| try sampling.put(cx.a, k, v); try grammar.refusal(cx, body); - if (input.tools.len > 0 and try tool_specs.choiceRequiresCall(cx.a, body.get("tool_choice"))) try sampling.put(cx.a, "tool_call_required", .{ .bool = true }); - if (body.field("parallel_tool_calls")) |p| if (p == .bool and !p.bool) try sampling.put(cx.a, "tool_call_single", .{ .bool = true }); - const thinking = try fields.thinkingFields(cx, body, srv.effort_levels); + if (srv.family == null and input.tools.len > 0 and try tool_specs.choiceRequiresCall(cx.a, body.get("tool_choice"))) try sampling.put(cx.a, "tool_call_required", .{ .bool = true }); + if (srv.family == null) if (body.field("parallel_tool_calls")) |p| if (p == .bool and !p.bool) try sampling.put(cx.a, "tool_call_single", .{ .bool = true }); + const thinking: fields.Thinking = if (srv.family) |fam| blk: { + const t = try fam.thinking(cx, body); + break :blk .{ .enable = t.enable, .effort = t.effort }; + } else try fields.thinkingFields(cx, body, srv.effort_levels); if (thinking.effort) |e| try sampling.put(cx.a, "reasoning_effort", .{ .string = e }); if (thinking.enable) |on| try sampling.put(cx.a, "enable_thinking", .{ .bool = on }); input.fields = .{ .object = sampling }; @@ -167,6 +179,7 @@ fn plan(srv: *Server, cx: *Cx, is_chat: bool, raw: Value) errors.Refused!Plan { if (body.field("parallel_tool_calls")) |p| { if (p != .bool) return cx.refuse("parallel_tool_calls must be a boolean"); policy.single = !p.bool; + input.single_call = policy.single; } const options = body.get("stream_options"); return .{ @@ -234,326 +247,20 @@ pub fn run(srv: *Server, a: Allocator, out: Out, gone: Gone, is_chat: bool, raw: return; }; p.input.id = id; - const turn = cortheon.begin(srv, a, raw, id, is_chat, &p.input); + // Cortheon reads calls in the template path's format: a family with its own chat rules runs without it + const turn = if (srv.family == null) cortheon.begin(srv, a, raw, id, is_chat, &p.input) else null; var r: Run = .{ .srv = srv, .a = a, .out = out, .is_chat = is_chat, .plan = p, .id = id, .created = std.Io.Clock.real.now(srv.io).toSeconds(), .turn = turn }; if (p.stream) r.stream(gone, field) else r.whole(gone, field); } /// Why a request got an error reply, logged before it: its access line shows only the status. -fn logRefused(id: []const u8, message: []const u8) void { +pub fn logRefused(id: []const u8, message: []const u8) void { var buf: [1400]u8 = undefined; log.line("{s}", .{log.refused(&buf, id, message)}); } -fn errorOther(a: Allocator, message: []const u8) Allocator.Error!Value { +pub fn errorOther(a: Allocator, message: []const u8) Allocator.Error!Value { const o = try json.newObject(a); try o.put(a, "message", .{ .string = message }); return .{ .object = o }; } - -const Run = struct { - srv: *Server, - a: Allocator, - out: Out, - is_chat: bool, - plan: Plan, - id: []const u8, - created: i64, - streamed_prose: bool = false, - prose_sent: std.ArrayList(u8) = .empty, // the content a tool stream sent - text_offset: usize = 0, // characters a completion stream's logprob rows have covered - turn: ?*cortheon.Turn = null, // --cortheon's work on this request - - fn chunk(r: *Run, delta: ?Value, finish: ?[]const u8) Allocator.Error!Value { - const a = r.a; - const o = try json.newObject(a); - try o.put(a, "id", .{ .string = r.id }); - try o.put(a, "object", .{ .string = if (r.is_chat) "chat.completion.chunk" else "text_completion" }); - try o.put(a, "created", try json.intValue(a, r.created)); - try o.put(a, "model", .{ .string = r.plan.named }); - const choice = try json.newObject(a); - try choice.put(a, "index", .{ .int = "0" }); - if (r.is_chat) { - const d: Value = if (delta) |v| (if (v == .object) v else if (v.string.len > 0) try chat.deltaOf(a, "content", v.string) else .{ .object = try json.newObject(a) }) else .{ .object = try json.newObject(a) }; - try choice.put(a, "delta", d); - // a content delta's rows (logprobs.Stream) are the choice's, as OpenAI streams them - if (d == .object) if (d.object.fetchOrderedRemove("logprobs")) |kv| try choice.put(a, "logprobs", kv.value); - try choice.put(a, "finish_reason", if (finish) |f| .{ .string = f } else .null); - } else { - const d = delta orelse Value{ .string = "" }; - const rows = if (d == .object) d.object.get("logprobs") else null; - try choice.put(a, "text", if (d == .object) d.object.get("content") orelse .{ .string = "" } else d); - try choice.put(a, "finish_reason", if (finish) |f| .{ .string = f } else .null); - try choice.put(a, "logprobs", if (rows) |v| try @import("logprobs.zig").legacy(a, v, &r.text_offset) else .null); - } - const list = try a.alloc(Value, 1); - list[0] = .{ .object = choice }; - try o.put(a, "choices", .{ .array = list }); - return .{ .object = o }; - } - - fn emit(r: *Run, payload: Value) error{Closed}!void { - return r.out.vt.event(r.out.ctx, payload); - } - - fn usage(r: *Run, reply: *const chat.Reply) Allocator.Error!Value { - const a = r.a; - const o = try json.newObject(a); - try o.put(a, "prompt_tokens", try json.intValue(a, reply.prompt_tokens)); - try o.put(a, "completion_tokens", try json.intValue(a, reply.completion_tokens)); - try o.put(a, "total_tokens", try json.intValue(a, reply.prompt_tokens + reply.completion_tokens)); - const cached = try json.newObject(a); - try cached.put(a, "cached_tokens", try json.intValue(a, reply.cached_tokens)); - try o.put(a, "prompt_tokens_details", .{ .object = cached }); - const thought = try json.newObject(a); - try thought.put(a, "reasoning_tokens", try json.intValue(a, reply.reasoning_tokens)); - try o.put(a, "completion_tokens_details", .{ .object = thought }); - return .{ .object = o }; - } - - /// ``response_extras``: exact_mode, the matched stop, the runtime and the draft counters. - fn extras(r: *Run, o: *json.Object, reply: *const chat.Reply) Allocator.Error!void { - try o.put(r.a, "exact_mode", .{ .string = "exact" }); - if (reply.stop_sequence) |s| try o.put(r.a, "stop_sequence", .{ .string = s }); - try o.put(r.a, "tensorfold", reply.runtime); - try o.put(r.a, "speculative", reply.speculative); - if (r.turn) |t| try o.put(r.a, "cortheon", try t.summary(r.a)); - } - - /// ``ToolCallPolicy.finish``: the reply's calls parsed out of its content. - fn attachCalls(r: *Run, reply: *chat.Reply) Allocator.Error!?[]Value { - const tools = r.plan.input.tools; - if (tools.len == 0) return null; - const parsed = try tool_parse.parse(r.a, reply.content, tools, r.plan.policy.maxCalls()); - const calls = (if (r.turn != null) cortheon.clientCalls(parsed.calls) else parsed.calls) orelse { - if (r.plan.policy.single or parsed.calls != null) reply.content = try r.plan.policy.parsedContent(r.a, parsed.content); - return null; - }; - reply.content = try r.plan.policy.parsedContent(r.a, parsed.content); - reply.finish_reason = "tool_calls"; - if (r.plan.policy.single) reply.tool_calls_streamed = false; - return if (r.plan.policy.single and calls.len > 1) calls[0..1] else calls; - } - - fn fail(r: *Run, cx: *const Cx, field: []const u8) void { - const a = r.a; - const kind_ok = cx.kind != .other and cx.kind != .server; - const body = (if (kind_ok) errorBody(a, cx, field) else errorOther(a, cx.message)) catch return; - r.out.vt.reply(r.out.ctx, if (kind_ok) 400 else 500, wrapError(a, body) catch return); - } - - /// A reply that ended before anything was sent: a refusal's 400, a failure's 500, nothing when the client left. - fn unsent(r: *Run, cx: *Cx, e: chat.Failure, field: []const u8) void { - switch (e) { - error.Cancelled => {}, - error.OutOfMemory => r.fail(&.{ .a = r.a, .kind = .server, .message = "out of memory" }, field), - error.Failed => r.fail(&.{ .a = r.a, .kind = .server, .message = "the reply failed" }, field), - error.Refused => { - if (cx.kind == .other) cx.kind = .server; - if (cx.kind != .server) logRefused(r.id, cx.message); // a 500 is a failure, not a refusal - r.fail(cx, field); - }, - } - } - - /// A context-window refusal on a stream: the role chunk when there are no tools, then the error event and [DONE]. - fn streamedContext(r: *Run, cx: *const Cx, tools: bool, field: []const u8) void { - const a = r.a; - // a context-window refusal: the head opens with no prompt behind it, so it carries no count - r.out.vt.open(r.out.ctx, 0) catch return; - if (!tools and r.is_chat) r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; - logRefused(r.id, cx.message); - const body = errorBody(a, cx, field) catch return; - r.emit(wrapError(a, body) catch return) catch return; - r.out.vt.event(r.out.ctx, null) catch {}; - } - - fn whole(r: *Run, gone: Gone, field: []const u8) void { - const a = r.a; - var cx: Cx = .{ .a = a }; - var reply = (if (r.turn) |t| cortheon.loop.whole(t, r.srv, &cx, &r.plan.input, gone) else chat.run(r.srv, &cx, r.plan.input, null, gone)) catch |e| return r.unsent(&cx, e, field); - const calls = r.attachCalls(&reply) catch return; - const o = json.newObject(a) catch return; - r.wholeBody(o, &reply, calls) catch return; - r.out.vt.reply(r.out.ctx, 200, .{ .object = o }); - r.warmNext(&reply, calls); - } - - /// The next turn's prompt prefilled in the background, from the reply as its client got it (warm.zig). - fn warmNext(r: *Run, reply: *const chat.Reply, calls: ?[]Value) void { - if (!r.is_chat) return; - const message = json.newObject(r.a) catch return; - message.put(r.a, "role", .{ .string = "assistant" }) catch return; - message.put(r.a, "content", if (calls != null and reply.content.len == 0) .null else .{ .string = reply.content }) catch return; - if (calls) |c| message.put(r.a, "tool_calls", .{ .array = c }) catch return; - if (r.turn) |t| t.remember(r.a, r.plan.input, .{ .object = message }); - warm.reply(r.srv, r.plan.input, .{ .object = message }, reply.thinking, reply.effort, reply.prompt_tokens); - } - - fn wholeBody(r: *Run, o: *json.Object, reply: *const chat.Reply, calls: ?[]Value) Allocator.Error!void { - const a = r.a; - try o.put(a, "id", .{ .string = r.id }); - try o.put(a, "object", .{ .string = if (r.is_chat) "chat.completion" else "text_completion" }); - try o.put(a, "created", try json.intValue(a, r.created)); - try o.put(a, "model", .{ .string = r.plan.named }); - const choice = try json.newObject(a); - try choice.put(a, "index", .{ .int = "0" }); - if (r.is_chat) { - const message = try json.newObject(a); - try message.put(a, "role", .{ .string = "assistant" }); - try message.put(a, "content", if (calls != null and reply.content.len == 0) .null else .{ .string = reply.content }); - if (reply.reasoning) |t| if (t.len > 0) try message.put(a, "reasoning_content", .{ .string = t }); - if (calls) |c| try message.put(a, "tool_calls", .{ .array = c }); - try choice.put(a, "message", .{ .object = message }); - try choice.put(a, "finish_reason", .{ .string = reply.finish_reason }); - if (reply.logprobs) |lp| try choice.put(a, "logprobs", lp); - } else { - try choice.put(a, "text", .{ .string = reply.content }); - try choice.put(a, "finish_reason", .{ .string = reply.finish_reason }); - var from: usize = 0; - try choice.put(a, "logprobs", if (reply.logprobs) |lp| try @import("logprobs.zig").legacy(a, lp, &from) else .null); - if (reply.prompt_logprobs) |p| try choice.put(a, "prompt_logprobs", p); - } - const list = try a.alloc(Value, 1); - list[0] = .{ .object = choice }; - try o.put(a, "choices", .{ .array = list }); - try o.put(a, "usage", try r.usage(reply)); - try r.extras(o, reply); - } - - const StreamSink = struct { - run: *Run, - tools: bool, - - fn call(ctx: *anyopaque, delta: Value) error{Closed}!void { - const s: *StreamSink = @ptrCast(@alignCast(ctx)); - const r = s.run; - if (!s.tools) { - // completions stream text only, with its logprob rows when asked - if (!r.is_chat and delta != .string and (delta != .object or delta.object.get("content") == null)) return; - return r.emit(r.chunk(delta, null) catch return error.Closed); - } - const filtered = r.policyDelta(delta) catch return error.Closed; - const d = filtered orelse return; - try r.prose(); - r.prose_sent.appendSlice(r.a, if (d == .string) d.string else d.strField("content") orelse "") catch return error.Closed; - return r.emit(r.chunk(d, null) catch return error.Closed); - } - }; - - /// The one-call policy on a streamed delta: text filtered, tool_calls dropped; null when nothing is left. - fn policyDelta(r: *Run, delta: Value) Allocator.Error!?Value { - if (!r.plan.policy.single) return if (delta.truthy()) delta else null; - if (delta == .string) { - const t = try r.plan.policy.text(r.a, delta.string); - return if (t.len > 0) Value{ .string = t } else null; - } - const o = try json.newObject(r.a); - var content: ?[]const u8 = null; - for (delta.object.keys(), delta.object.values()) |k, v| { - if (std.mem.eql(u8, k, "tool_calls")) continue; - if (std.mem.eql(u8, k, "content") and v == .string) { - content = v.string; - continue; - } - try o.put(r.a, k, v); - } - if (content) |c| { - const t = try r.plan.policy.text(r.a, c); - if (t.len > 0) try o.put(r.a, "content", .{ .string = t }); - } - return if (o.count() > 0) Value{ .object = o } else null; - } - - /// The assistant role before the first prose of a tool-using stream. - fn prose(r: *Run) error{Closed}!void { - if (r.streamed_prose) return; - r.streamed_prose = true; - const role = roleDelta(r.a) catch return error.Closed; - return r.emit(r.chunk(role, null) catch return error.Closed); - } - - fn stream(r: *Run, gone: Gone, field: []const u8) void { - const a = r.a; - const tools = r.plan.input.tools.len > 0; - var cx: Cx = .{ .a = a }; - // a context-window refusal is reported inside the stream; every other refusal is still a 400 before it opens - var compaction: ?compact.Stamp = null; - const prepared = if (r.srv.config.compact_at == null) chat.prepare(r.srv, &cx, r.plan.input, gone) catch |e| { - if (e == error.Refused and cx.kind == .context_length) return r.streamedContext(&cx, tools, field); - return r.unsent(&cx, e, field); - } else blk: { - const ready = compact.prepare(r.srv, &cx, r.plan.input, gone) catch |e| { - if (e == error.Refused and cx.kind == .context_length) return r.streamedContext(&cx, tools, field); - return r.unsent(&cx, e, field); - }; - compaction = ready.stamp; - 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); - r.out.vt.open(r.out.ctx, prepared.prompt_len) catch return; - if (!tools and r.is_chat) r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; - var sink_state: StreamSink = .{ .run = r, .tools = tools }; - handed = true; - const sink: chat.Sink = .{ .ctx = &sink_state, .call = StreamSink.call }; - var reply = (if (r.turn) |t| cortheon.loop.stream(t, r.srv, &cx, &r.plan.input, prepared, sink, gone) else chat.generate(r.srv, &cx, prepared, sink, gone)) catch |e| { - switch (e) { - error.Cancelled => return, - error.Refused => if (cx.kind != .other and cx.kind != .server) { - logRefused(r.id, cx.message); - const body = errorBody(a, &cx, field) catch return; - r.emit(wrapError(a, body) catch return) catch return; - r.out.vt.event(r.out.ctx, null) catch {}; - return; - }, - else => {}, - } - const message = if (e == error.Refused) cx.message else if (e == error.OutOfMemory) "out of memory" else "the reply failed"; - log.line("stream error: {s}", .{message}); - const err = json.newObject(a) catch return; - err.put(a, "message", .{ .string = message }) catch return; - err.put(a, "type", .{ .string = "server_error" }) catch return; - r.emit(wrapError(a, .{ .object = err }) catch return) catch return; - r.out.vt.event(r.out.ctx, null) catch {}; - return; - }; - var calls: ?[]Value = null; - if (tools) { - calls = r.attachCalls(&reply) catch return; - const tail = r.plan.policy.flush(); - if (tail.len > 0) { - r.prose() catch return; - r.prose_sent.appendSlice(a, tail) catch return; - r.emit(r.chunk(.{ .string = tail }, null) catch return) catch return; - } - if (calls != null and !reply.tool_calls_streamed) { - r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; - const deltas = tool_parse.deltas(a, calls.?) catch return; - for (deltas) |d| r.emit(r.chunk(d, null) catch return) catch return; - } else if (reply.content.len > 0 and !r.streamed_prose) { - r.emit(r.chunk(.{ .string = reply.content }, null) catch return) catch return; - } else if (r.plan.policy.kept(reply.content, r.prose_sent.items)) |rest| { - r.emit(r.chunk(.{ .string = rest }, null) catch return) catch return; - } - } - const last = r.chunk(.{ .string = "" }, if (reply.finish_reason.len > 0) reply.finish_reason else "length") catch return; - if (compaction) |s| compact.stamp(&cx, &reply, s) catch return; - r.extras(last.object, &reply) catch return; - const use = r.usage(&reply) catch return; - if (!r.plan.separate_usage) last.object.put(a, "usage", use) catch return; - r.emit(last) catch return; - if (r.plan.separate_usage) { - const u = r.chunk(null, null) catch return; - u.object.put(a, "choices", .{ .array = &.{} }) catch return; - u.object.put(a, "usage", use) catch return; - r.emit(u) catch return; - } - r.out.vt.event(r.out.ctx, null) catch {}; - r.warmNext(&reply, calls); - } -}; - -fn roleDelta(a: Allocator) Allocator.Error!Value { - return chat.deltaOf(a, "role", "assistant"); -} diff --git a/zig/src/server/openai_run.zig b/zig/src/server/openai_run.zig new file mode 100644 index 0000000000..9ea9f740ab --- /dev/null +++ b/zig/src/server/openai_run.zig @@ -0,0 +1,343 @@ +//! One OpenAI request's run: the reply generated, then sent whole or streamed as chunks. +const std = @import("std"); +const json = @import("json.zig"); +const errors = @import("errors.zig"); +const chat = @import("chat.zig"); +const compact = @import("compact.zig"); +const tool_parse = @import("tool_parse.zig"); +const tool_stream = @import("tool_stream.zig"); +const ids = @import("ids.zig"); +const log = @import("log.zig"); +const request_log = @import("request_log.zig"); +const sse = @import("sse.zig"); +const warm = @import("warm.zig"); +const cortheon = @import("cortheon/cortheon.zig"); +const Server = @import("server.zig").Server; +const Value = json.Value; +const Cx = errors.Cx; +const Allocator = std.mem.Allocator; +const openai = @import("openai.zig"); +const Out = openai.Out; +const Gone = openai.Gone; +const errorBody = openai.errorBody; +const wrapError = openai.wrapError; +const replyModel = openai.replyModel; +const Plan = openai.Plan; +const logRefused = openai.logRefused; +const errorOther = openai.errorOther; + +pub const Run = struct { + srv: *Server, + a: Allocator, + out: Out, + is_chat: bool, + plan: Plan, + id: []const u8, + created: i64, + streamed_prose: bool = false, + prose_sent: std.ArrayList(u8) = .empty, // the content a tool stream sent + text_offset: usize = 0, // characters a completion stream's logprob rows have covered + turn: ?*cortheon.Turn = null, // --cortheon's work on this request + + fn chunk(r: *Run, delta: ?Value, finish: ?[]const u8) Allocator.Error!Value { + const a = r.a; + const o = try json.newObject(a); + try o.put(a, "id", .{ .string = r.id }); + try o.put(a, "object", .{ .string = if (r.is_chat) "chat.completion.chunk" else "text_completion" }); + try o.put(a, "created", try json.intValue(a, r.created)); + try o.put(a, "model", .{ .string = r.plan.named }); + const choice = try json.newObject(a); + try choice.put(a, "index", .{ .int = "0" }); + if (r.is_chat) { + const d: Value = if (delta) |v| (if (v == .object) v else if (v.string.len > 0) try chat.deltaOf(a, "content", v.string) else .{ .object = try json.newObject(a) }) else .{ .object = try json.newObject(a) }; + try choice.put(a, "delta", d); + // a content delta's rows (logprobs.Stream) are the choice's, as OpenAI streams them + if (d == .object) if (d.object.fetchOrderedRemove("logprobs")) |kv| try choice.put(a, "logprobs", kv.value); + try choice.put(a, "finish_reason", if (finish) |f| .{ .string = f } else .null); + } else { + const d = delta orelse Value{ .string = "" }; + const rows = if (d == .object) d.object.get("logprobs") else null; + try choice.put(a, "text", if (d == .object) d.object.get("content") orelse .{ .string = "" } else d); + try choice.put(a, "finish_reason", if (finish) |f| .{ .string = f } else .null); + try choice.put(a, "logprobs", if (rows) |v| try @import("logprobs.zig").legacy(a, v, &r.text_offset) else .null); + } + const list = try a.alloc(Value, 1); + list[0] = .{ .object = choice }; + try o.put(a, "choices", .{ .array = list }); + return .{ .object = o }; + } + + fn emit(r: *Run, payload: Value) error{Closed}!void { + return r.out.vt.event(r.out.ctx, payload); + } + + fn usage(r: *Run, reply: *const chat.Reply) Allocator.Error!Value { + const a = r.a; + const o = try json.newObject(a); + try o.put(a, "prompt_tokens", try json.intValue(a, reply.prompt_tokens)); + try o.put(a, "completion_tokens", try json.intValue(a, reply.completion_tokens)); + try o.put(a, "total_tokens", try json.intValue(a, reply.prompt_tokens + reply.completion_tokens)); + const cached = try json.newObject(a); + try cached.put(a, "cached_tokens", try json.intValue(a, reply.cached_tokens)); + try o.put(a, "prompt_tokens_details", .{ .object = cached }); + const thought = try json.newObject(a); + try thought.put(a, "reasoning_tokens", try json.intValue(a, reply.reasoning_tokens)); + try o.put(a, "completion_tokens_details", .{ .object = thought }); + return .{ .object = o }; + } + + /// ``response_extras``: exact_mode, the matched stop, the runtime and the draft counters. + fn extras(r: *Run, o: *json.Object, reply: *const chat.Reply) Allocator.Error!void { + try o.put(r.a, "exact_mode", .{ .string = "exact" }); + if (reply.stop_sequence) |s| try o.put(r.a, "stop_sequence", .{ .string = s }); + try o.put(r.a, "tensorfold", reply.runtime); + try o.put(r.a, "speculative", reply.speculative); + if (r.turn) |t| try o.put(r.a, "cortheon", try t.summary(r.a)); + } + + /// ``ToolCallPolicy.finish``: the reply's calls parsed out of its content. + fn attachCalls(r: *Run, reply: *chat.Reply) Allocator.Error!?[]Value { + if (reply.calls) |calls| return if (calls.len > 0) calls else null; // the family parsed them (and applied the one-call policy) + const tools = r.plan.input.tools; + if (tools.len == 0) return null; + const parsed = try tool_parse.parse(r.a, reply.content, tools, r.plan.policy.maxCalls()); + const calls = (if (r.turn != null) cortheon.clientCalls(parsed.calls) else parsed.calls) orelse { + if (r.plan.policy.single or parsed.calls != null) reply.content = try r.plan.policy.parsedContent(r.a, parsed.content); + return null; + }; + reply.content = try r.plan.policy.parsedContent(r.a, parsed.content); + reply.finish_reason = "tool_calls"; + if (r.plan.policy.single) reply.tool_calls_streamed = false; + return if (r.plan.policy.single and calls.len > 1) calls[0..1] else calls; + } + + fn fail(r: *Run, cx: *const Cx, field: []const u8) void { + const a = r.a; + const kind_ok = cx.kind != .other and cx.kind != .server; + const body = (if (kind_ok) errorBody(a, cx, field) else errorOther(a, cx.message)) catch return; + r.out.vt.reply(r.out.ctx, if (kind_ok) 400 else 500, wrapError(a, body) catch return); + } + + /// A reply that ended before anything was sent: a refusal's 400, a failure's 500, nothing when the client left. + fn unsent(r: *Run, cx: *Cx, e: chat.Failure, field: []const u8) void { + switch (e) { + error.Cancelled => {}, + error.OutOfMemory => r.fail(&.{ .a = r.a, .kind = .server, .message = "out of memory" }, field), + error.Failed => r.fail(&.{ .a = r.a, .kind = .server, .message = "the reply failed" }, field), + error.Refused => { + if (cx.kind == .other) cx.kind = .server; + if (cx.kind != .server) logRefused(r.id, cx.message); // a 500 is a failure, not a refusal + r.fail(cx, field); + }, + } + } + + /// A context-window refusal on a stream: the role chunk when there are no tools, then the error event and [DONE]. + fn streamedContext(r: *Run, cx: *const Cx, tools: bool, field: []const u8) void { + const a = r.a; + // a context-window refusal: the head opens with no prompt behind it, so it carries no count + r.out.vt.open(r.out.ctx, 0) catch return; + if (!tools and r.is_chat) r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; + logRefused(r.id, cx.message); + const body = errorBody(a, cx, field) catch return; + r.emit(wrapError(a, body) catch return) catch return; + r.out.vt.event(r.out.ctx, null) catch {}; + } + + pub fn whole(r: *Run, gone: Gone, field: []const u8) void { + const a = r.a; + var cx: Cx = .{ .a = a }; + var reply = (if (r.turn) |t| cortheon.loop.whole(t, r.srv, &cx, &r.plan.input, gone) else chat.run(r.srv, &cx, r.plan.input, null, gone)) catch |e| return r.unsent(&cx, e, field); + const calls = r.attachCalls(&reply) catch return; + const o = json.newObject(a) catch return; + r.wholeBody(o, &reply, calls) catch return; + r.out.vt.reply(r.out.ctx, 200, .{ .object = o }); + r.warmNext(&reply, calls); + } + + /// The next turn's prompt prefilled in the background, from the reply as its client got it (warm.zig). + fn warmNext(r: *Run, reply: *const chat.Reply, calls: ?[]Value) void { + if (!r.is_chat) return; + const message = json.newObject(r.a) catch return; + message.put(r.a, "role", .{ .string = "assistant" }) catch return; + message.put(r.a, "content", if (calls != null and reply.content.len == 0) .null else .{ .string = reply.content }) catch return; + if (calls) |c| message.put(r.a, "tool_calls", .{ .array = c }) catch return; + if (r.turn) |t| t.remember(r.a, r.plan.input, .{ .object = message }); + warm.reply(r.srv, r.plan.input, .{ .object = message }, reply.thinking, reply.effort, reply.prompt_tokens); + } + + fn wholeBody(r: *Run, o: *json.Object, reply: *const chat.Reply, calls: ?[]Value) Allocator.Error!void { + const a = r.a; + try o.put(a, "id", .{ .string = r.id }); + try o.put(a, "object", .{ .string = if (r.is_chat) "chat.completion" else "text_completion" }); + try o.put(a, "created", try json.intValue(a, r.created)); + try o.put(a, "model", .{ .string = r.plan.named }); + const choice = try json.newObject(a); + try choice.put(a, "index", .{ .int = "0" }); + if (r.is_chat) { + const message = try json.newObject(a); + try message.put(a, "role", .{ .string = "assistant" }); + // a family's reply: ``content or None`` and no prose beside calls; the template path keeps a call's prose + const no_content = if (r.srv.family != null) calls != null or reply.content.len == 0 else calls != null and reply.content.len == 0; + try message.put(a, "content", if (no_content) .null else .{ .string = reply.content }); + if (reply.reasoning) |t| if (t.len > 0) { + if (r.srv.family) |fam| { + const d = try fam.reasoningDelta(a, t); + for (d.object.keys(), d.object.values()) |k, v| try message.put(a, k, v); + } else try message.put(a, "reasoning_content", .{ .string = t }); + }; + if (calls) |c| try message.put(a, "tool_calls", .{ .array = c }); + try choice.put(a, "message", .{ .object = message }); + try choice.put(a, "finish_reason", .{ .string = reply.finish_reason }); + if (reply.logprobs) |lp| try choice.put(a, "logprobs", lp); + } else { + try choice.put(a, "text", .{ .string = reply.content }); + try choice.put(a, "finish_reason", .{ .string = reply.finish_reason }); + var from: usize = 0; + try choice.put(a, "logprobs", if (reply.logprobs) |lp| try @import("logprobs.zig").legacy(a, lp, &from) else .null); + if (reply.prompt_logprobs) |p| try choice.put(a, "prompt_logprobs", p); + } + const list = try a.alloc(Value, 1); + list[0] = .{ .object = choice }; + try o.put(a, "choices", .{ .array = list }); + try o.put(a, "usage", try r.usage(reply)); + try r.extras(o, reply); + } + + const StreamSink = struct { + run: *Run, + tools: bool, + + fn call(ctx: *anyopaque, delta: Value) error{Closed}!void { + const s: *StreamSink = @ptrCast(@alignCast(ctx)); + const r = s.run; + if (!s.tools) { + // completions stream text only, with its logprob rows when asked + if (!r.is_chat and delta != .string and (delta != .object or delta.object.get("content") == null)) return; + return r.emit(r.chunk(delta, null) catch return error.Closed); + } + const filtered = r.policyDelta(delta) catch return error.Closed; + const d = filtered orelse return; + try r.prose(); + r.prose_sent.appendSlice(r.a, if (d == .string) d.string else d.strField("content") orelse "") catch return error.Closed; + return r.emit(r.chunk(d, null) catch return error.Closed); + } + }; + + /// The one-call policy on a streamed delta: text filtered, tool_calls dropped; null when nothing is left. + fn policyDelta(r: *Run, delta: Value) Allocator.Error!?Value { + if (!r.plan.policy.single) return if (delta.truthy()) delta else null; + if (delta == .string) { + const t = try r.plan.policy.text(r.a, delta.string); + return if (t.len > 0) Value{ .string = t } else null; + } + const o = try json.newObject(r.a); + var content: ?[]const u8 = null; + for (delta.object.keys(), delta.object.values()) |k, v| { + if (std.mem.eql(u8, k, "tool_calls")) continue; + if (std.mem.eql(u8, k, "content") and v == .string) { + content = v.string; + continue; + } + try o.put(r.a, k, v); + } + if (content) |c| { + const t = try r.plan.policy.text(r.a, c); + if (t.len > 0) try o.put(r.a, "content", .{ .string = t }); + } + return if (o.count() > 0) Value{ .object = o } else null; + } + + /// The assistant role before the first prose of a tool-using stream. + fn prose(r: *Run) error{Closed}!void { + if (r.streamed_prose) return; + r.streamed_prose = true; + const role = roleDelta(r.a) catch return error.Closed; + return r.emit(r.chunk(role, null) catch return error.Closed); + } + + pub fn stream(r: *Run, gone: Gone, field: []const u8) void { + const a = r.a; + // a family streams calls and prose itself, so its stream is the plain one + const tools = r.plan.input.tools.len > 0 and r.srv.family == null; + var cx: Cx = .{ .a = a }; + // a context-window refusal is reported inside the stream + var compaction: ?compact.Stamp = null; + const prepared = if (r.srv.config.compact_at == null) chat.prepare(r.srv, &cx, r.plan.input, gone) catch |e| { + if (e == error.Refused and cx.kind == .context_length) return r.streamedContext(&cx, tools, field); + return r.unsent(&cx, e, field); + } else blk: { + const ready = compact.prepare(r.srv, &cx, r.plan.input, gone) catch |e| { + if (e == error.Refused and cx.kind == .context_length) return r.streamedContext(&cx, tools, field); + return r.unsent(&cx, e, field); + }; + compaction = ready.stamp; + 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); + r.out.vt.open(r.out.ctx, prepared.prompt_len) catch return; + if (!tools and r.is_chat) r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; + var sink_state: StreamSink = .{ .run = r, .tools = tools }; + handed = true; + const sink: chat.Sink = .{ .ctx = &sink_state, .call = StreamSink.call }; + var reply = (if (r.turn) |t| cortheon.loop.stream(t, r.srv, &cx, &r.plan.input, prepared, sink, gone) else chat.generate(r.srv, &cx, prepared, sink, gone)) catch |e| { + switch (e) { + error.Cancelled => return, + error.Refused => if (cx.kind != .other and cx.kind != .server) { + logRefused(r.id, cx.message); + const body = errorBody(a, &cx, field) catch return; + r.emit(wrapError(a, body) catch return) catch return; + r.out.vt.event(r.out.ctx, null) catch {}; + return; + }, + else => {}, + } + const message = if (e == error.Refused) cx.message else if (e == error.OutOfMemory) "out of memory" else "the reply failed"; + log.line("stream error: {s}", .{message}); + const err = json.newObject(a) catch return; + err.put(a, "message", .{ .string = message }) catch return; + err.put(a, "type", .{ .string = "server_error" }) catch return; + r.emit(wrapError(a, .{ .object = err }) catch return) catch return; + r.out.vt.event(r.out.ctx, null) catch {}; + return; + }; + var calls: ?[]Value = null; + if (tools) { + calls = r.attachCalls(&reply) catch return; + const tail = r.plan.policy.flush(); + if (tail.len > 0) { + r.prose() catch return; + r.prose_sent.appendSlice(a, tail) catch return; + r.emit(r.chunk(.{ .string = tail }, null) catch return) catch return; + } + if (calls != null and !reply.tool_calls_streamed) { + r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; + const deltas = tool_parse.deltas(a, calls.?) catch return; + for (deltas) |d| r.emit(r.chunk(d, null) catch return) catch return; + } else if (reply.content.len > 0 and !r.streamed_prose) { + r.emit(r.chunk(.{ .string = reply.content }, null) catch return) catch return; + } else if (r.plan.policy.kept(reply.content, r.prose_sent.items)) |rest| { + r.emit(r.chunk(.{ .string = rest }, null) catch return) catch return; + } + } + const last = r.chunk(.{ .string = "" }, if (reply.finish_reason.len > 0) reply.finish_reason else "length") catch return; + if (compaction) |s| compact.stamp(&cx, &reply, s) catch return; + r.extras(last.object, &reply) catch return; + const use = r.usage(&reply) catch return; + if (!r.plan.separate_usage) last.object.put(a, "usage", use) catch return; + r.emit(last) catch return; + if (r.plan.separate_usage) { + const u = r.chunk(null, null) catch return; + u.object.put(a, "choices", .{ .array = &.{} }) catch return; + u.object.put(a, "usage", use) catch return; + r.emit(u) catch return; + } + r.out.vt.event(r.out.ctx, null) catch {}; + r.warmNext(&reply, calls); + } +}; + +fn roleDelta(a: Allocator) Allocator.Error!Value { + return chat.deltaOf(a, "role", "assistant"); +} diff --git a/zig/src/server/prompt.zig b/zig/src/server/prompt.zig index e68363bfe7..f3fa2c9931 100644 --- a/zig/src/server/prompt.zig +++ b/zig/src/server/prompt.zig @@ -6,10 +6,11 @@ const messages_mod = @import("messages.zig"); const model_text = @import("model_text.zig"); const chat = @import("chat.zig"); const Server = @import("server.zig").Server; +const family_mod = @import("family.zig"); const Value = json.Value; const Cx = errors.Cx; -pub const Rendered = struct { ids: []const u32, history_len: usize = 0, rewind_len: usize = 0 }; +pub const Rendered = struct { ids: []const u32, history_len: usize = 0, rewind_len: usize = 0, images_omitted: u32 = 0, images: ?family_mod.Images = null }; /// Where lessons decide to act on a reply: the row before the generation prompt's ``, else before its last. pub fn decideOf(srv: *const Server, ids: []const u32) u32 { @@ -27,6 +28,11 @@ pub const isTitle = messages_mod.isTitleRequest; /// ``render_prompt_ids``: messages normalized for this template, rendered, and a dangling ```` closed. pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, thinking: bool, effort: ?[]const u8, generation: bool) errors.Refused![]const u32 { + if (srv.family) |fam| { + const got = try fam.render(cx, .{ .messages = messages, .tools = tools, .thinking = thinking, .effort = effort orelse fam.defaultEffort(), .generation = generation }); + if (got.images) |im| im.release(); // ids only: the images are not held past the render + return got.ids; + } const normalized = try messages_mod.toolArguments(cx, try messages_mod.normalize(cx, messages, srv.late_system, srv.needs_user_after_tool)); var problem: []const u8 = ""; const options: model_text.RenderOptions = .{ @@ -96,6 +102,17 @@ pub fn prepare(srv: *Server, cx: *Cx, input: chat.Input, thinking: bool, effort: } }, .ids => |ids| return .{ .ids = ids }, }; + if (srv.family) |fam| { + const p: family_mod.Prompt = .{ .messages = input.messages, .tools = input.tools, .thinking = thinking, .effort = effort orelse fam.defaultEffort(), .body = input.body }; + const got = try fam.render(cx, p); + var history = p; + history.generation = false; + errdefer if (got.images) |im| im.release(); + const before = try fam.render(cx, history); + if (before.images) |im| im.release(); + const n = before.ids.len; + return .{ .ids = got.ids, .history_len = if (n > 0 and n < got.ids.len and std.mem.eql(u32, got.ids[0..n], before.ids)) n else 0, .images_omitted = got.images_omitted, .images = got.images }; + } const prompt = try renderIds(srv, cx, input.messages, input.tools, thinking, effort, true); const history = try renderIds(srv, cx, input.messages, input.tools, thinking, effort, false); var history_len: usize = 0; diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 7630f5217a..0c0f60224d 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -12,6 +12,7 @@ pub const server = @import("server.zig"); pub const serve = @import("serve.zig"); pub const hub = @import("hub.zig"); pub const tool_specs = @import("tool_specs.zig"); +pub const family = @import("family.zig"); pub const Server = server.Server; pub const Config = server.Config; diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index da27cb6b0e..6ad9cfd36f 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -73,6 +73,8 @@ pub const Setup = struct { started: i96 = 0, /// Called once the socket listens, with its port (tests read it). on_listen: ?*const fn (port: u16) void = null, + /// The checkpoint family's own chat rules (family.zig), or null. + family: ?@import("family.zig").Family = null, }; fn env(s: Setup, name: []const u8) ?[]const u8 { @@ -119,6 +121,7 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { .seed_salt = salt, .request_log = env(s, "TENSORFOLD_REQUEST_LOG"), .dashboard = args.dashboard, + .family = s.family, .compact_at = if (args.compact_auto) .{ .auto = {} } else if (args.compact_fraction) |f| .{ .fraction = f } else null, .compact_keep = args.compact_keep, .compact_memory = args.compact_memory, @@ -156,6 +159,8 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { s.served, args.host, port, shownSampling(a, s.sampling), if (args.no_drafts or srv.info.plain_only) "off" else "on", if (window > 0) std.fmt.bufPrint(&window_text, "{d}", .{window}) catch "?" else "unlimited", loaded, }); + if (greedy(s.sampling)) + log.line("WARNING: default sampling is greedy (temperature <= 0) - a client that omits `temperature` decodes greedily, which on a long agentic turn can loop in reasoning and return empty content", .{}); if (args.thinking and std.mem.indexOf(u8, s.text.templateSource(), "enable_thinking") != null) log.line("thinking on (the chat template's default): replies reason in reasoning_content before the answer in content, and max_tokens counts both. --no-thinking turns it off; a request can send chat_template_kwargs {{\"enable_thinking\": false}}", .{}); if (srv.cortheon) |c| log.line("cortheon on: chat requests get Cortheon's reasoning program, web research before building, and cortheon_search/cortheon_fetch run here ({s}); a request can send \"cortheon\": false", .{c.engine}); @@ -180,6 +185,17 @@ fn resolve(io: std.Io, host: []const u8, port: u16) ?std.Io.net.IpAddress { return std.Io.net.IpAddress.resolve(io, host, port) catch null; } +/// Whether requests that name no temperature decode greedily (the startup warning). +fn greedy(sampling: ?json.Value) bool { + const v = sampling orelse return true; + const t = v.get("temperature") orelse return true; + return switch (t) { + .float => |f| f <= 0, + .int => |i| (std.fmt.parseFloat(f64, i) catch 0) <= 0, + else => true, + }; +} + /// The serving line's sampling: greedy, or each default as ``name value``. fn shownSampling(a: Allocator, sampling: ?json.Value) []const u8 { const v = sampling orelse return "greedy"; diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 011734ee9a..fc90467612 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -23,6 +23,7 @@ const live = @import("live.zig"); const slide_graph = @import("slide_graph.zig"); const slide_lesson = @import("slide_lesson.zig"); const cortheon = @import("cortheon/cortheon.zig"); +const family_mod = @import("family.zig"); const Value = json.Value; const Cx = errors.Cx; const Allocator = std.mem.Allocator; @@ -62,6 +63,8 @@ pub const Config = struct { /// BRAVE_SEARCH_API_KEY: Cortheon searches through Brave instead of DuckDuckGo. cortheon_search_key: ?[]const u8 = null, timeouts: http_conn.Timeouts = .{}, + /// The checkpoint family's own chat rules (family.zig); null: the chat template path. + family: ?family_mod.Family = null, }; pub const Server = struct { @@ -82,6 +85,7 @@ pub const Server = struct { think_close_end: ?u32 = null, metrics: metrics.Metrics, store: responses.Store, + family: ?family_mod.Family = null, next_id: std.atomic.Value(u64) = .init(1), /// Connections open now; a stop waits for them before freeing what they read. open_connections: std.atomic.Value(u32) = .init(0), @@ -102,7 +106,7 @@ pub const Server = struct { /// Reads what the template and tokenizer decide once: late system role, think markers, efforts, the forced close. pub fn init(gpa: Allocator, io: std.Io, engine: api.Engine, text: model_text.Text, config: Config, keys: ?*auth.Store) !*Server { const srv = try gpa.create(Server); - srv.* = .{ .gpa = gpa, .io = io, .engine = engine, .info = engine.info(), .text = text, .config = config, .keys = keys, .metrics = .{ .gpa = gpa }, .store = .{ .gpa = gpa }, .arena = .init(gpa), .slide = .init(gpa, io, config.slide_graph), .teacher = .init(gpa) }; + srv.* = .{ .gpa = gpa, .io = io, .engine = engine, .info = engine.info(), .text = text, .config = config, .keys = keys, .metrics = .{ .gpa = gpa }, .store = .{ .gpa = gpa }, .arena = .init(gpa), .slide = .init(gpa, io, config.slide_graph), .teacher = .init(gpa), .family = config.family }; errdefer srv.deinit(); // a later step's failure stops the keepalive before the caller closes the engine if (config.cortheon) srv.cortheon = try cortheon.State.init(gpa, io, "TensorFold-Cortheon (+https://github.com/ashhart/TensorFold)", config.cortheon_search_key); if (config.keep_warm_s > 0) if (engine.keepaliveTarget()) |target| { @@ -155,7 +159,12 @@ pub const Server = struct { const options = try json.newObject(cx.a); const defaults = if (srv.config.default_sampling) |d| try fields.parseNumbers(cx, d) else null; if (defaults) |d| try merge(options, cx, d); - try merge(options, cx, try fields.parseNumbers(cx, f)); + const asked = try fields.parseNumbers(cx, f); + // a family whose Python server ignores the request's min_p keeps the default's + if (srv.family) |fam| if (!fam.vt.reads_min_p) { + _ = asked.object.orderedRemove("min_p"); + }; + try merge(options, cx, asked); const temp = number(options.get("temperature")) orelse 0; if (temp <= 0) return null; if (srv.info.greedy_only) return cx.refuse(if ((if (defaults) |d| number(d.get("temperature")) else null) orelse 0 > 0) diff --git a/zig/src/server/tokens.zig b/zig/src/server/tokens.zig index 3a2e351921..52bcb87092 100644 --- a/zig/src/server/tokens.zig +++ b/zig/src/server/tokens.zig @@ -27,7 +27,27 @@ fn tokenize(srv: *Server, cx: *Cx, body: Value) errors.Refused!Value { if (body != .object) return cx.refuse("the request body must be a JSON object"); const strings = try flag(cx, body, "return_token_strs", false); var ids: []const u32 = undefined; - if (body.has("messages")) { + if (body.has("messages") and srv.family != null) { + // rendered exactly as a chat request with the same body + const fam = srv.family.?; + const generation = try flag(cx, body, "add_generation_prompt", true); + const m = body.get("messages").?; + if (m != .array or m.array.len == 0) return cx.refuse("messages must be a non-empty list"); + const t = try fam.thinking(cx, body); + const rendered = fam.render(cx, .{ .messages = m, .tools = try fam.tools(cx, body), .thinking = t.enable orelse srv.config.enable_thinking, .effort = t.effort orelse srv.config.reasoning_effort orelse fam.defaultEffort(), .generation = generation, .body = body }) catch |e| switch (e) { + error.Refused => return cx.refuse(cx.message), + else => |x| return x, + }; + ids = rendered.ids; + if (rendered.images) |im| { // /tokenize shows the image token at each image position (expand(virtual=False)) + defer im.release(); + const out = try cx.a.dupe(u32, rendered.ids); + for (out) |*id| if (id.* >= 1 << 24) { + id.* = im.token; + }; + ids = out; + } + } else if (body.has("messages")) { const generation = try flag(cx, body, "add_generation_prompt", true); const msgs = try messages.normalize(cx, body.get("messages"), "system", srv.needs_user_after_tool); const tools = tool_specs.active(cx, body.get("tools"), body.get("tool_choice")) catch |e| switch (e) { From bf19238c00a700755fe5c99702465f9dfe146ce5 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:05:28 +0700 Subject: [PATCH 16/43] TP RoCE: a failed transport halts its QPs, drains them and keeps its region; the region is DONTFORK and NOHUGEPAGE A local protection error can arrive while the IOMMU invalidates its mappings lazily. On such an error the old path left the other HCA's QP ready to send (still taking the peer's writes), hung up the fate channel without telling the peer to stop, and freed the region at once: unregistered from CUDA with no GPU sync, then QP destroy, MR deregistration and munmap. A device write after the unmap could then reach a reused page through a stale IOTLB entry. - Roce.fail and Roce.abort (the fail-fast hook) move every QP to the error state: no send goes out and no peer write lands. - Roce.deinit quiesces first: QPs to error, every posted request reaped (flushed or not), cuCtxSynchronize. A failed transport, or one that does not quiesce, keeps its region registered and mapped; the process exit tears it down. - Session.close on a failed transport goes fail-fast before closing, so the peer aborts and halts its HCAs first. - The region is MADV_DONTFORK and MADV_NOHUGEPAGE before ibv_reg_mr; a refusal fails the start. - A QP whose actual max_inline_data is under 4 B is refused (the flag post uses lkey 0 on a stack word). - One log line a rank: region base and length, each HCA's lkey and rkey, the peer's region and rkeys. Tests: on fake verbs Roce.abort halts both QPs; quiesce halts both before the first flush and reaps every posted write; an aborted transport reports itself failed, and over the shm wire the aborted rank keeps its region while the other frees it; smaps shows dc and nh on a prepared region; the inline bound refuses 3 B and takes 4. --- zig/src/fabric/tp/roce.zig | 42 ++++++++++++++ zig/src/fabric/tp/roce_fake.zig | 37 ++++++++++++ zig/src/fabric/tp/roce_inflight_test.zig | 74 ++++++++++++++++++++++++ zig/src/fabric/tp/roce_post_test.zig | 16 +++++ zig/src/fabric/tp/roce_test.zig | 3 + zig/src/fabric/tp/roce_verbs.zig | 66 ++++++++++++++++++++- zig/src/fabric/tp/tp.zig | 2 + 7 files changed, 237 insertions(+), 3 deletions(-) diff --git a/zig/src/fabric/tp/roce.zig b/zig/src/fabric/tp/roce.zig index 3e7c1e644a..c8e8010810 100644 --- a/zig/src/fabric/tp/roce.zig +++ b/zig/src/fabric/tp/roce.zig @@ -198,6 +198,7 @@ pub const Roce = struct { r.mapping = try posix.mmap(null, len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .PRIVATE, .ANONYMOUS = true }, -1, 0); r.region = r.mapping; errdefer posix.munmap(r.mapping); + try rv.prepareRegion(r.mapping); r.wire = .{ .verbs = try rv.VerbsWire.open(o.hcas, r.region, .{ .send = send_off, .recv = send_off + 2 * r.slot_bytes, .flags = flag_off, .slot = r.slot_bytes }, o.traffic_class) }; }, .shm => { @@ -222,6 +223,8 @@ pub const Roce = struct { if (all[1 - r.me].len != @sizeOf(rv.Record)) return error.BadRecord; @memcpy(std.mem.asBytes(&peer), all[1 - r.me]); try v.connect(peer); + var line: [512]u8 = undefined; + log.info("[tensorfold] tp roce: rank {d} registered {s}", .{ r.me, v.describe(&line) }); return true; }, } @@ -261,14 +264,37 @@ pub const Roce = struct { posix.munmap(r.mapping); } + /// Regions a failed transport left registered and mapped (the process exit is their only teardown). + pub var kept: std.atomic.Value(u32) = .init(0); + + /// HCAs and GPU quiet, then unregister and unmap; a failed transport's region is kept until exit. pub fn deinit(r: *Roce) void { + const was_failed = r.failed(); + r.word(w_abort).store(1, .release); r.running.store(false, .release); if (r.thread) |t| t.join(); + const quiet = r.quiesce(); + if (was_failed or !quiet) { + const why = if (was_failed) "failed" else "did not quiesce"; + log.warn("[tensorfold] tp roce: rank {d}: the transport {s}; region 0x{x} ({d} B) kept registered until exit", .{ r.me, why, @intFromPtr(r.region.ptr), r.region.len }); + _ = kept.fetchAdd(1, .monotonic); + return; + } r.closeDevice(); r.closeWire(); std.heap.page_allocator.destroy(r); } + /// The QPs to error with every posted request reaped, and the GPU's work on the region done. + fn quiesce(r: *Roce) bool { + const wire = switch (r.wire) { + .verbs => |*v| v.quiesce(rv.quiesce_ns), + .shm => true, + }; + const gpu = if (r.d) |d| d.api.cuCtxSynchronize() == cuda.abi.success else true; + return wire and gpu; + } + fn word(r: *const Roce, i: usize) *std.atomic.Value(u32) { return @ptrCast(@alignCast(r.region[ctrl_off + 4 * i ..].ptr)); } @@ -279,6 +305,21 @@ pub const Roce = struct { log.err("[tensorfold] tp roce proxy: {s}", .{s}); r.word(w_proxy_failed).store(1, .release); r.word(w_abort).store(1, .release); // waits on this rank end now instead of at the timeout + r.halt(); + } + + /// The HCAs stop: nothing of this rank's goes out and none of the peer's writes lands while the ranks go down. + fn halt(r: *Roce) void { + switch (r.wire) { + .verbs => |*v| _ = v.halt(), + .shm => {}, + } + } + + /// The transport failed or was aborted: a device may still write to its region. + pub fn failed(r: *const Roce) bool { + if (r.aborted.load(.acquire) or r.word(w_abort).load(.acquire) != 0) return true; + return r.word(w_proxy_failed).load(.acquire) != 0 or r.word(w_failed).load(.acquire) != 0; } /// The proxy loop: spin on the doorbell, post each new sequence (at most two pending), reap while idle, then nap. @@ -382,6 +423,7 @@ pub const Roce = struct { pub fn abort(r: *Roce) void { r.aborted.store(true, .release); r.word(w_abort).store(1, .release); + r.halt(); } pub fn check(r: *Roce) Error!void { diff --git a/zig/src/fabric/tp/roce_fake.zig b/zig/src/fabric/tp/roce_fake.zig index a6b02db77d..fbb55891f6 100644 --- a/zig/src/fabric/tp/roce_fake.zig +++ b/zig/src/fabric/tp/roce_fake.zig @@ -76,6 +76,43 @@ pub const Fake = struct { } }; +/// QPs that complete nothing on their own: a QP moved to error flushes what it holds, as an HCA reports it. +pub const Halts = struct { + pub var halted: [2]bool = .{ false, false }; + pub var flushed_before_halt: u32 = 0; + pub var flushes: u32 = 0; + + pub fn reset() void { + halted = .{ false, false }; + flushed_before_halt = 0; + flushes = 0; + } + + pub fn modifyQp(qp: *abi.Qp, attr: *abi.QpAttr, mask: c_int) callconv(.c) c_int { + if (attr.qp_state == abi.qps_err and mask == abi.mask.state) halted[qp.qp_num] = true; + return 0; + } + + pub fn pollCq(cq: *abi.Cq, _: c_int, wc: [*]abi.Wc) callconv(.c) c_int { + const q = cq.handle; + if (Fake.len[q] == 0 or !halted[q]) return 0; + if (!halted[0] or !halted[1]) flushed_before_halt += 1; + Fake.inflight -= Fake.queue[q][Fake.head[q]]; + Fake.head[q] = (Fake.head[q] + 1) % 512; + Fake.len[q] -= 1; + flushes += 1; + wc[0] = std.mem.zeroes(abi.Wc); + wc[0].status = 5; // IBV_WC_WR_FLUSH_ERR + return 1; + } + + /// Points a wire's QPs at these handlers. + pub fn install(hw: *Hw, w: *rv.VerbsWire) void { + w.v.api.ibv_modify_qp = modifyQp; + for (&hw.ctx) |*c| c.ops.poll_cq = pollCq; + } +}; + /// Two fake HCAs behind one VerbsWire, connected to a peer whose region has this rank's layout. pub const Hw = struct { ctx: [2]abi.Context = undefined, diff --git a/zig/src/fabric/tp/roce_inflight_test.zig b/zig/src/fabric/tp/roce_inflight_test.zig index a8b83a1564..130f1e1eab 100644 --- a/zig/src/fabric/tp/roce_inflight_test.zig +++ b/zig/src/fabric/tp/roce_inflight_test.zig @@ -4,6 +4,7 @@ const rv = @import("roce_verbs.zig"); const fake = @import("roce_fake.zig"); const Roce = @import("roce.zig").Roce; const Fake = fake.Fake; +const Halts = fake.Halts; const slot = fake.slot; test "chunked posts: the wire's own in-flight count and peak match the HCAs', and reach the 2 MiB cap" { @@ -24,3 +25,76 @@ test "chunked posts: the wire's own in-flight count and peak match the HCAs', an try std.testing.expectEqual(Fake.max_inflight, w.peak.load(.monotonic)); try std.testing.expectEqual(@as(u64, 2 << 20), w.peak.load(.monotonic)); } + +test "a failed transport: both QPs go to error before the first flush and every posted write is reaped" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(slot)); + defer std.testing.allocator.free(region); + @memset(region, 0); + Fake.region = region; + Fake.slot_bytes = slot; + Fake.reset(1); + Halts.reset(); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + Halts.install(&hw, &w); + // two 1 MiB ops posted, neither completed: 2 MiB in flight at the HCAs + try w.post(1, 1, slot); + try w.post(0, 2, slot); + try std.testing.expectEqual(@as(u64, 2 * slot), w.inflight()); + try std.testing.expect(w.quiesce(std.time.ns_per_s)); + try std.testing.expect(Halts.halted[0] and Halts.halted[1]); + try std.testing.expectEqual(@as(u32, 0), Halts.flushed_before_halt); + try std.testing.expectEqual(@as(u32, 4), Halts.flushes); + try std.testing.expectEqual(@as(u64, 0), w.inflight()); + try std.testing.expectEqual(@as(u64, 0), Fake.inflight); + // an aborted transport reports itself failed: its deinit keeps the region + Halts.reset(); + var r: Roce = .{ .me = 0, .slot_bytes = slot, .timeout_ns = 1, .region = region, .mapping = region, .wire = .{ .verbs = w }, .cpu = null }; + try std.testing.expect(!r.failed()); + r.abort(); + try std.testing.expect(Halts.halted[0] and Halts.halted[1]); + try std.testing.expect(r.failed()); +} + +test "a QP whose inline capacity cannot take the 4-byte flag post is refused" { + try std.testing.expect(!rv.inlineFits(.{ .max_send_wr = 1, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 3 })); + try std.testing.expect(rv.inlineFits(.{ .max_send_wr = 1, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 4 })); +} + +/// /proc/self/smaps' VmFlags of the mapping that starts at `at`. +fn vmFlags(at: usize, buf: []u8) ![]const u8 { + const fd = std.c.open("/proc/self/smaps", .{}); + if (fd < 0) return error.NoSmaps; + defer _ = std.c.close(fd); + var n: usize = 0; + while (n < buf.len) { + const got = std.c.read(fd, buf[n..].ptr, buf.len - n); + if (got <= 0) break; + n += @intCast(got); + } + var head: [32]u8 = undefined; + const want = try std.fmt.bufPrint(&head, "{x}-", .{at}); + var lines = std.mem.splitScalar(u8, buf[0..n], '\n'); + var inside = false; + while (lines.next()) |l| { + if (std.mem.startsWith(u8, l, want)) inside = true; + if (inside and std.mem.startsWith(u8, l, "VmFlags:")) return l; + } + return error.NotFound; +} + +test "the region is DONTFORK and NOHUGEPAGE before registration" { + if (@import("builtin").os.tag != .linux) return error.SkipZigTest; + const len = Roce.regionBytes(slot); + const m = try std.posix.mmap(null, len, .{ .READ = true, .WRITE = true }, .{ .TYPE = .PRIVATE, .ANONYMOUS = true }, -1, 0); + defer std.posix.munmap(m); + const buf = try std.testing.allocator.alloc(u8, 4 << 20); + defer std.testing.allocator.free(buf); + // smaps: dc = MADV_DONTFORK, nh = MADV_NOHUGEPAGE + const before = try vmFlags(@intFromPtr(m.ptr), buf); + try std.testing.expect(std.mem.indexOf(u8, before, " dc") == null and std.mem.indexOf(u8, before, " nh") == null); + try rv.prepareRegion(m); + const flags = try vmFlags(@intFromPtr(m.ptr), buf); + try std.testing.expect(std.mem.indexOf(u8, flags, " dc") != null); + try std.testing.expect(std.mem.indexOf(u8, flags, " nh") != null); +} diff --git a/zig/src/fabric/tp/roce_post_test.zig b/zig/src/fabric/tp/roce_post_test.zig index a4f7971fed..a284293000 100644 --- a/zig/src/fabric/tp/roce_post_test.zig +++ b/zig/src/fabric/tp/roce_post_test.zig @@ -102,3 +102,19 @@ test "chunked strides posted back to back: at most 2 MiB in flight over both HCA try std.testing.expect(Fake.max_inflight > 1 << 20); } } + +test "Roce.abort (the fail-fast hook) moves both QPs to error: no send goes out and no peer write lands" { + const region = try std.testing.allocator.alignedAlloc(u8, .fromByteUnits(4096), Roce.regionBytes(fake.slot)); + defer std.testing.allocator.free(region); + @memset(region, 0); + Fake.region = region; + Fake.slot_bytes = fake.slot; + Fake.reset(1); + fake.Halts.reset(); + var hw: fake.Hw = .{}; + var w = hw.wire(region); + fake.Halts.install(&hw, &w); + var r: Roce = .{ .me = 0, .slot_bytes = fake.slot, .timeout_ns = 1, .region = region, .mapping = region, .wire = .{ .verbs = w }, .cpu = null }; + r.abort(); + try std.testing.expect(fake.Halts.halted[0] and fake.Halts.halted[1]); +} diff --git a/zig/src/fabric/tp/roce_test.zig b/zig/src/fabric/tp/roce_test.zig index dff1ab7eb9..ec6bc4d65e 100644 --- a/zig/src/fabric/tp/roce_test.zig +++ b/zig/src/fabric/tp/roce_test.zig @@ -154,7 +154,10 @@ test "a peer that stops makes the waiting rank time out, and abort ends a wait a try std.testing.expectEqual(@as(?anyerror, error.Timeout), run.errs[0]); try std.testing.expectEqual(@as(u32, 20), run.done[0]); var run2: Run = .{ .port = 0, .wire = .shm, .ops = 50, .stop_after = 20, .abort_after_ms = 50 }; + const kept = roce.Roce.kept.load(.monotonic); try twoRanks(&run2); + // the aborted rank keeps its region until exit, the other frees it + try std.testing.expectEqual(kept + 1, roce.Roce.kept.load(.monotonic)); try std.testing.expectEqual(@as(?anyerror, error.Aborted), run2.errs[0]); try std.testing.expectEqual(@as(?anyerror, null), run2.errs[1]); } diff --git a/zig/src/fabric/tp/roce_verbs.zig b/zig/src/fabric/tp/roce_verbs.zig index cfcb2d143d..8ed5e3eca4 100644 --- a/zig/src/fabric/tp/roce_verbs.zig +++ b/zig/src/fabric/tp/roce_verbs.zig @@ -20,6 +20,25 @@ pub const max_inflight: u64 = 2 * @as(u64, guard.max_write); /// A post that waits on completions this long fails the proxy (a peer or HCA that stopped answering). const inflight_wait_ns: u64 = 10 * std.time.ns_per_s; const ring_len = send_depth / 4; +/// The flag post's bytes: a u32 on the stack under lkey 0, legal only while the QP sends it inline. +pub const flag_bytes: u32 = @sizeOf(u32); +/// The longest `quiesce` waits for the HCAs to flush what was posted. +pub const quiesce_ns: u64 = 5 * std.time.ns_per_s; + +/// Before ibv_reg_mr: no fork shares or copies the region (DONTFORK), its pages stay 4 KiB (NOHUGEPAGE); else no start. +pub fn prepareRegion(region: []align(std.heap.page_size_min) u8) Error!void { + if (comptime @import("builtin").os.tag != .linux) return; // the verbs wire runs on Linux only + const advice = [_]u32{ std.os.linux.MADV.DONTFORK, std.os.linux.MADV.NOHUGEPAGE }; + for (advice) |a| std.posix.madvise(region.ptr, region.len, a) catch |e| { + log.warn("tp roce: madvise({d}) on the {d} B region failed: {t}; refusing to start", .{ a, region.len, e }); + return error.VerbsFailed; + }; +} + +/// Whether a QP's actual inline capacity takes the flag post. +pub fn inlineFits(cap: abi.QpCap) bool { + return cap.max_inline_data >= flag_bytes; +} /// One HCA as configured: its device name and RoCE v2 GID index (null: detected). pub const HcaSpec = struct { name: [64]u8 = @splat(0), name_len: u8 = 0, gid: ?u8 = null }; @@ -190,6 +209,10 @@ pub const VerbsWire = struct { var init: abi.QpInitAttr = .{ .send_cq = cq, .recv_cq = cq, .cap = .{ .max_send_wr = send_depth, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 16 }, .qp_type = abi.qpt_rc }; const qp = api.ibv_create_qp(pd, &init) orelse return error.VerbsFailed; errdefer _ = api.ibv_destroy_qp(qp); + if (!inlineFits(init.cap)) { + log.warn("tp roce: {s}'s QP sends {d} B inline, the flag post needs {d}; refusing to start", .{ name, init.cap.max_inline_data, flag_bytes }); + return error.VerbsFailed; + } var a: abi.QpAttr = .{ .qp_state = abi.qps_init, .port_num = 1, .qp_access_flags = abi.access_remote_write }; if (api.ibv_modify_qp(qp, &a, abi.mask.state | abi.mask.pkey_index | abi.mask.port | abi.mask.access_flags) != 0) return error.VerbsFailed; w.hcas[w.n_hca] = .{ .ctx = ctx, .pd = pd, .mr = mr, .cq = cq, .qp = qp, .gid_index = gi, .gid = gid.raw, .mtu = @intCast(port.active_mtu), .lid = port.lid }; @@ -210,6 +233,39 @@ pub const VerbsWire = struct { w.v.close(); } + /// The region and each HCA's keys, this side's and the peer's (a protection error is checked against them). + pub fn describe(w: *const VerbsWire, buf: []u8) []const u8 { + var out: std.Io.Writer = .fixed(buf); + out.print("region 0x{x} len {d}", .{ @intFromPtr(w.region.ptr), w.region.len }) catch return out.buffered(); + for (w.hcas[0..w.n_hca], 0..) |h, i| { + out.print("; HCA {d} lkey 0x{x} rkey 0x{x}, peer region 0x{x} rkey 0x{x}", .{ i, h.mr.lkey, h.mr.rkey, w.peer_region, w.peer_rkey[i] }) catch break; + } + return out.buffered(); + } + + /// The QPs to error: the HCAs stop sending and refuse the peer's writes, and what was posted completes as flushed. + pub fn halt(w: *VerbsWire) bool { + var ok = true; + for (w.hcas[0..w.n_hca]) |h| { + var a: abi.QpAttr = .{ .qp_state = abi.qps_err }; + if (w.v.api.ibv_modify_qp(h.qp, &a, abi.mask.state) != 0) ok = false; + } + return ok; + } + + /// `halt`, then every posted request's completion reaped, flushed or not: no request of the region left at an HCA. + pub fn quiesce(w: *VerbsWire, timeout_ns: u64) bool { + if (!w.halt()) return false; + const t0 = sock.nowNs(); + for (w.hcas[0..w.n_hca]) |*h| while (h.outstanding > 0) { + var wc: [32]abi.Wc = undefined; + const n = h.ctx.ops.poll_cq(h.cq, 32, &wc); + if (n < 0) return false; + if (n > 0) w.popped(h, @intCast(n)) else if (sock.nowNs() - t0 > timeout_ns) return false; + }; + return true; + } + pub fn record(w: *const VerbsWire) Record { var r: Record = .{ .n_hca = w.n_hca, .region = @intFromPtr(w.region.ptr), .slot_bytes = w.slot_bytes, .region_len = w.region.len }; for (w.hcas[0..w.n_hca], 0..) |h, i| { @@ -292,12 +348,16 @@ pub const VerbsWire = struct { log.err("tp roce: RDMA write of seq {d} failed: {s} (vendor 0x{x})", .{ c.wr_id, w.v.statusText(c.status), c.vendor_err }); return error.Completion; }; - for (0..@intCast(n)) |_| { + w.popped(h, @intCast(n)); + } + + fn popped(w: *VerbsWire, h: *Hca, n: u32) void { + for (0..n) |_| { h.inflight -= h.ring[h.head]; h.head = (h.head + 1) % ring_len; } - h.outstanding -= @intCast(n); - w.completed += @intCast(n); + h.outstanding -= n; + w.completed += n; } /// Bytes posted and not yet completed over every HCA. diff --git a/zig/src/fabric/tp/tp.zig b/zig/src/fabric/tp/tp.zig index 901544b58b..6cc5ceb6ec 100644 --- a/zig/src/fabric/tp/tp.zig +++ b/zig/src/fabric/tp/tp.zig @@ -187,6 +187,8 @@ pub const Session = struct { /// A normal stop: every rank calls it (rank 0 first marks the stop as expected with `fate.expectStop`). pub fn close(s: *Session) void { + // a failed transport: every rank goes down (the peers stop posting) before this one lets go of anything + if (s.roce) |r| if (r.failed()) if (s.fate) |*f| f.fatal("transport", error.BackendFailed); if (s.fate) |*f| f.close(); s.plan.close(); if (s.roce) |r| r.deinit(); From 9a486bbd493ffe70a032d5743dd13534f996cebc Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:25:57 +0700 Subject: [PATCH 17/43] Server: a stop logs its signal and the open connections when the accept loop ends, and again if it exits with connections open. --- zig/src/server/serve.zig | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index 6ad9cfd36f..6974118c28 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -13,11 +13,20 @@ const Allocator = std.mem.Allocator; const posix = std.posix; var stop_requested: std.atomic.Value(bool) = .init(false); +/// The signal that asked for the stop (0: none yet), for the shutdown line +var stop_signal: std.atomic.Value(i32) = .init(0); -fn onStop(_: posix.SIG) callconv(.c) void { +fn onStop(sig: posix.SIG) callconv(.c) void { + stop_signal.store(@intCast(@backingInt(sig)), .release); stop_requested.store(true, .release); } +fn signalName(sig: i32) []const u8 { + if (sig == @backingInt(posix.SIG.TERM)) return "SIGTERM"; + if (sig == @backingInt(posix.SIG.INT)) return "SIGINT"; + return "a stop signal"; +} + /// SIGTERM and SIGINT stop the server cleanly; SIGHUP rereads the key file; SIGPIPE is ignored, as Python ignores it. pub fn installSignals(keys: bool) void { const stop: posix.Sigaction = .{ .handler = .{ .handler = onStop }, .mask = posix.sigemptyset(), .flags = 0 }; @@ -171,10 +180,15 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { const accept = std.Thread.spawn(.{}, server_mod.Server.serve, .{ srv, lis, &stop_requested }) catch return 1; accept.join(); lis.close(); + // before the replies still running fail (their rounds' CUDA calls fail once the exit tears the driver down) + log.line("shutting down ({s}): no new requests; {d} connection(s) open, waiting up to 2 s before exiting", .{ signalName(stop_signal.load(.acquire)), srv.open_connections.load(.acquire) }); if (drawing) ticker.finish(); var waited: u32 = 0; while (srv.open_connections.load(.acquire) > 0 and waited < 40) : (waited += 1) std.Io.sleep(io, .fromMilliseconds(50), .awake) catch {}; - if (srv.open_connections.load(.acquire) > 0) std.process.exit(0); // replies still open: end without freeing what they read + if (srv.open_connections.load(.acquire) > 0) { + log.line("shutting down: exiting with {d} connection(s) still open; their requests fail", .{srv.open_connections.load(.acquire)}); + std.process.exit(0); // replies still open: end without freeing what they read + } return 0; } From beae7e870456400b0b892128a8917a2886a063f4 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:06:03 +0700 Subject: [PATCH 18/43] TP RoCE: tf-tp-test failfast V transport checks the failed-transport path on real HCAs, only on an idle link - `tf-tp-test failfast V transport [N [HOLD_S]]`: N clean 4 KiB sums, the ranks meet over NCCL, each waits for its link to go idle (every HCA's posts completed), both agree over NCCL, then the victim takes the proxy's failure path (`Roce.inject`, which is `Roce.fail`) and closes its session: both ranks halt their QPs and exit 70, and the victim's region stays registered. When either rank's link is not idle there is no inject, a clean stop, REFUSED and rc 2. - Its host-only forms (TF_TP_FAILFAST=0, or a hold) are refused unless TF_TP_ALLOW_HOST_ONLY=1 and no RDMA wire is in use (the shm wire, or no RoCE), so they never run on an HCA. - `run_transport.sh MASTER PEER_SSH OUT`, driven from rank 0's host (rank 1 over ssh): prechecks on both hosts (no GPU process, memlock unlimited), then T0 (suite), T1 and T3 (each rank as the victim). After each step it reads both hosts' kernel log since the step and every HCA *err* counter; a hit, a counter that grew, an unexpected rc or a failure line prints STOP: and exits 1. `run_transport.sh t2 OUT` runs the host-only form on one host. - The `halted` and `quiesced` lines carry each QP's state read back with ibv_query_qp (an optional symbol), each HCA's posts and completions, the completions the drain reaped and how many of them were flushes. - `Roce.allGatherV` refuses a stride past the slot (one op; the route never sends one). - The README states the RoCE limits and how to run the check. Tests: on fake verbs the link reads busy with two posts out and idle after the drain, and the report line names both QPs in the error state with every completion a flush. --- zig/src/fabric/tp/README.md | 11 +++ zig/src/fabric/tp/gpu_main.zig | 21 +++- zig/src/fabric/tp/gpu_transport.zig | 71 ++++++++++++++ zig/src/fabric/tp/roce.zig | 30 +++++- zig/src/fabric/tp/roce_fake.zig | 7 ++ zig/src/fabric/tp/roce_inflight_test.zig | 8 +- zig/src/fabric/tp/roce_verbs.zig | 40 ++++++++ zig/src/fabric/tp/run_transport.sh | 120 +++++++++++++++++++++++ zig/src/fabric/verbs.zig | 4 + 9 files changed, 307 insertions(+), 5 deletions(-) create mode 100644 zig/src/fabric/tp/gpu_transport.zig create mode 100755 zig/src/fabric/tp/run_transport.sh diff --git a/zig/src/fabric/tp/README.md b/zig/src/fabric/tp/README.md index 5d5c57358f..a0b4af5342 100644 --- a/zig/src/fabric/tp/README.md +++ b/zig/src/fabric/tp/README.md @@ -51,3 +51,14 @@ zig/src/fabric/tp/run_hosts.sh 1 192.0.2.1 ./zig-out/bin/tf-tp-test out-rank1 `ibv_devices` lists the RDMA devices; `ibv_devinfo -d NAME` shows a port's state and link layer. `/sys/class/infiniband/NAME/ports/1/gids/N` and `gid_attrs/ndevs/N` show each GID and its netdev. `ulimit -l` must allow the RoCE region to be pinned (`unlimited` is simplest). + +## RoCE limits and the failed-transport check + +No RDMA write is larger than 1 MiB, so on `TF_COMM_BACKEND=roce` a `TF_TP_MAILBOX_MAX_KB` above 1024 fails at boot. +`TF_TP_ROCE_CHUNKED_KB` (default 0, off) sends variable-length gathers past the slot over RoCE in slot-sized ops. +It needs a kernel built from this tree's `mailbox.cu`. +At most 2 MiB is posted and not yet completed at any time. +A rank whose transport fails moves its queue pairs to the error state and keeps its region until the process exits. + +`run_transport.sh MASTER_ADDR PEER_SSH OUTDIR` checks that failure path on two hosts with `tf-tp-test failfast V +transport` and stops at the first kernel log line or HCA error counter that grows; run it with no arguments for help. diff --git a/zig/src/fabric/tp/gpu_main.zig b/zig/src/fabric/tp/gpu_main.zig index e3b9a420e7..07d7c7afb9 100644 --- a/zig/src/fabric/tp/gpu_main.zig +++ b/zig/src/fabric/tp/gpu_main.zig @@ -1,4 +1,4 @@ -//! `tf-tp-test suite | failfast RANK fatal|kill | roce-host | kernel`: one rank of a two-process TP=2 test. +//! `tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | kernel`: one rank of a TP=2 test. const std = @import("std"); const cuda = @import("cuda"); const tp = @import("tp"); @@ -11,7 +11,7 @@ const max_bytes = (4 << 20) + 4096; pub fn main(init: std.process.Init) !u8 { const args = try init.minimal.args.toSlice(init.arena.allocator()); if (args.len < 2) { - std.debug.print("usage: tf-tp-test suite | failfast RANK fatal|kill | roce-host | kernel\n", .{}); + std.debug.print("usage: tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | kernel\n", .{}); return 2; } // `kernel`: whether the build embeds the mailbox kernel (-Dtp-nvcc or -Dtp-fatbin), for scripts to check first @@ -21,6 +21,10 @@ pub fn main(init: std.process.Init) !u8 { } const cfg = try tp.Config.fromEnv(); if (std.mem.eql(u8, args[1], "roce-host")) return @import("roce_host.zig").run(cfg, init.gpa); + if (std.mem.eql(u8, args[1], "failfast") and (!cfg.failfast or args.len > 5) and !hostOnlyAllowed(cfg)) { + std.debug.print("REFUSED: failfast with TF_TP_FAILFAST=0 or a hold is host-only: TF_TP_ALLOW_HOST_ONLY=1 and no RDMA wire\n", .{}); + return 2; + } if (cfg.backend != .nccl and tp_kernel.mailbox.len == 0) { std.debug.print("FAIL TF_COMM_BACKEND={t} needs the mailbox kernel: rebuild with -Dtp-nvcc or -Dtp-fatbin\n", .{cfg.backend}); return 2; @@ -36,6 +40,11 @@ pub fn main(init: std.process.Init) !u8 { const up_ms = @as(f64, @floatFromInt(tp.sock.nowNs() - t0)) / 1e6; std.debug.print("INFO rank {d}/{d} on {s} (sm_{d}), backend {t}, mailbox kernel {d} B, bootstrap + NCCL init {d:.0} ms\n", .{ cfg.rank, cfg.world, try ctx.name(&name_buf), try ctx.capability(), s.comm().kind(), tp_kernel.mailbox.len, up_ms }); if (std.mem.eql(u8, args[1], "suite")) return suite(&s, &driver, init.gpa, cfg); + if (std.mem.eql(u8, args[1], "failfast") and args.len >= 4 and std.mem.eql(u8, args[3], "transport")) { + const n = if (args.len > 4) try std.fmt.parseInt(u32, args[4], 10) else 50; + const hold = if (args.len > 5) try std.fmt.parseInt(u32, args[5], 10) else 0; + return @import("gpu_transport.zig").run(&s, &driver, init.gpa, cfg, try std.fmt.parseInt(u32, args[2], 10), n, hold); + } if (std.mem.eql(u8, args[1], "failfast") and args.len == 4) return failfast(&s, &driver, init.gpa, cfg, try std.fmt.parseInt(u32, args[2], 10), args[3]); std.debug.print("unknown command {s}\n", .{args[1]}); return 2; @@ -70,7 +79,13 @@ fn suite(s: *tp.Session, d: *const cuda.Driver, gpa: std.mem.Allocator, cfg: tp. return if (c.failures == 0) 0 else 1; } -fn realNow() f64 { +/// The failed-transport check's host-only forms (fail-fast off, a hold): TF_TP_ALLOW_HOST_ONLY=1 and no RDMA wire. +fn hostOnlyAllowed(cfg: tp.Config) bool { + const flag = std.mem.eql(u8, std.mem.span(std.c.getenv("TF_TP_ALLOW_HOST_ONLY") orelse ""), "1"); + return flag and (cfg.backend != .roce or cfg.roce_wire == .shm); +} + +pub fn realNow() f64 { var ts: std.c.timespec = undefined; _ = std.c.clock_gettime(.REALTIME, &ts); return @as(f64, @floatFromInt(ts.sec)) + @as(f64, @floatFromInt(ts.nsec)) / 1e9; diff --git a/zig/src/fabric/tp/gpu_transport.zig b/zig/src/fabric/tp/gpu_transport.zig new file mode 100644 index 0000000000..49cd625844 --- /dev/null +++ b/zig/src/fabric/tp/gpu_transport.zig @@ -0,0 +1,71 @@ +//! `tf-tp-test failfast V transport [N [HOLD_S]]`: the failed-transport path on real HCAs, injected on an idle link. +const std = @import("std"); +const cuda = @import("cuda"); +const tp = @import("tp"); +const checks = @import("gpu_checks.zig"); +const realNow = @import("gpu_main.zig").realNow; + +/// N clean sums, an NCCL meet, the victim's `Roce.inject` once every link is idle, then its close (fail-fast: exit 70). +pub fn run(s: *tp.Session, d: *const cuda.Driver, gpa: std.mem.Allocator, cfg: tp.Config, victim: u32, n: u32, hold_s: u32) !u8 { + const r = s.roce orelse { + std.debug.print("failfast transport needs TF_COMM_BACKEND=roce (the session runs {t})\n", .{s.comm().kind()}); + return 2; + }; + var c = try checks.Ctx.init(d, gpa, cfg.rank, 4096); + defer c.deinit(); + const comm = s.comm(); + for (0..n) |_| { + try comm.allReduce(c.in.ptr, c.out.ptr, 2048, .bf16, .sum, c.stream.handle); + try c.stream.synchronize(); + try comm.check(); + } + try s.nccl.iface().barrier(); + std.debug.print("rank {d}: {d} clean exchanges; both ranks met over NCCL\n", .{ cfg.rank, n }); + // the inject only on an idle link: every rank's posts completed, else no rank injects + const t_idle = tp.sock.nowNs(); + while (!r.linkIdle() and tp.sock.nowNs() - t_idle < 2 * std.time.ns_per_s) tp.sock.sleepNs(std.time.ns_per_ms); + var line: [384]u8 = undefined; + std.debug.print("rank {d}: link {s} before the inject: {s}\n", .{ cfg.rank, if (r.linkIdle()) "idle" else "BUSY", r.describeLink(&line) }); + if (try idleRanks(s, &c, r.linkIdle()) != cfg.world) { + std.debug.print("REFUSED rank {d}: a link was not idle; no inject, a clean stop\n", .{cfg.rank}); + if (s.fate) |*f| f.expectStop(); + try s.nccl.iface().barrier(); + s.close(); + return 2; + } + if (cfg.rank == victim) { + std.debug.print("INJECT transport rank {d} at {d:.6}\n", .{ cfg.rank, realNow() }); + r.inject(); + comm.check() catch |e| std.debug.print("rank {d}: transport reports {t}\n", .{ cfg.rank, e }); + } else if (s.fate != null) { + while (true) tp.sock.sleepNs(std.time.ns_per_s); // the fate channel ends this rank (exit 70) + } + const kept0 = tp.roce.Roce.kept.load(.monotonic); + s.close(); // fail-fast on: exit 70 from here + const kept = tp.roce.Roce.kept.load(.monotonic) - kept0; + const want: u32 = @intFromBool(cfg.rank == victim); + const live = cfg.rank != victim or mapped(r.region); + std.debug.print("rank {d}: closed; regions kept {d} (want {d}){s}\n", .{ cfg.rank, kept, want, if (cfg.rank != victim) "" else if (live) ", region still mapped" else ", region NOT mapped" }); + if (hold_s > 0) { + std.debug.print("HOLD rank {d} pid {d} for {d} s\n", .{ cfg.rank, std.c.getpid(), hold_s }); + tp.sock.sleepNs(@as(u64, hold_s) * std.time.ns_per_s); + } + const ok = kept == want and live; + std.debug.print("{s} rank {d}\n", .{ if (ok) "PASS" else "FAIL", cfg.rank }); + return if (ok) 0 else 1; +} + +fn mapped(region: []const u8) bool { + var vec: [1]u8 = undefined; + return std.os.linux.errno(std.os.linux.mincore(@constCast(region.ptr), 4096, &vec)) == .SUCCESS; +} + +/// The ranks whose link was idle, summed over NCCL (one f32 a rank). +fn idleRanks(s: *tp.Session, c: *checks.Ctx, idle: bool) !u32 { + var v: [1]f32 = .{@floatFromInt(@intFromBool(idle))}; + try c.in.uploadAsync(0, std.mem.asBytes(&v), c.stream.handle); + try s.nccl.iface().allReduce(c.in.ptr, c.out.ptr, 1, .f32, .sum, c.stream.handle); + try c.stream.synchronize(); + try c.out.download(0, std.mem.asBytes(&v)); + return @intFromFloat(v[0]); +} diff --git a/zig/src/fabric/tp/roce.zig b/zig/src/fabric/tp/roce.zig index c8e8010810..c9bb72a2c3 100644 --- a/zig/src/fabric/tp/roce.zig +++ b/zig/src/fabric/tp/roce.zig @@ -292,9 +292,19 @@ pub const Roce = struct { .shm => true, }; const gpu = if (r.d) |d| d.api.cuCtxSynchronize() == cuda.abi.success else true; + var line: [384]u8 = undefined; + log.info("[tensorfold] tp roce: rank {d} quiesced ({s}): {s}; GPU sync {s}", .{ r.me, if (wire) "drained" else "NOT drained", r.describeLink(&line), if (gpu) "done" else "FAILED" }); return wire and gpu; } + /// Each HCA's QP state, posts and completions, and what a drain reaped (tf-tp-test's idle check; shm: none). + pub fn describeLink(r: *const Roce, buf: []u8) []const u8 { + return switch (r.wire) { + .verbs => |*v| v.report(buf), + .shm => "shm wire", + }; + } + fn word(r: *const Roce, i: usize) *std.atomic.Value(u32) { return @ptrCast(@alignCast(r.region[ctrl_off + 4 * i ..].ptr)); } @@ -311,11 +321,28 @@ pub const Roce = struct { /// The HCAs stop: nothing of this rank's goes out and none of the peer's writes lands while the ranks go down. fn halt(r: *Roce) void { switch (r.wire) { - .verbs => |*v| _ = v.halt(), + .verbs => |*v| { + const ok = v.halt(); + var line: [384]u8 = undefined; + log.warn("[tensorfold] tp roce: rank {d} halted{s}: {s}", .{ r.me, if (ok) "" else " (a QP refused ERR)", v.report(&line) }); + }, .shm => {}, } } + /// Every post of this rank's completed (shm: nothing is ever in flight). + pub fn linkIdle(r: *const Roce) bool { + return switch (r.wire) { + .verbs => |*v| v.idle(), + .shm => true, + }; + } + + /// tf-tp-test's `failfast RANK transport`: the proxy's failure path (`fail`), here with no op in flight. + pub fn inject(r: *Roce) void { + r.fail("failure injected by tf-tp-test (no RDMA fault)", .{}); + } + /// The transport failed or was aborted: a device may still write to its region. pub fn failed(r: *const Roce) bool { if (r.aborted.load(.acquire) or r.word(w_abort).load(.acquire) != 0) return true; @@ -394,6 +421,7 @@ pub const Roce = struct { /// `allGatherBytes` where rank r posts only lens[r] bytes (device int32 [2], padded to 16, at least 16). pub fn allGatherV(r: *Roce, s: DevicePtr, o: DevicePtr, n: usize, lens: DevicePtr, stream: collective.Stream) Error!void { + if (n > r.slot_bytes) return error.Unsupported; // one op: past a slot is allGatherVChunked's if (lens == 0) return error.Invalid; return r.launch(.gather, s, o, n, lens, stream); } diff --git a/zig/src/fabric/tp/roce_fake.zig b/zig/src/fabric/tp/roce_fake.zig index fbb55891f6..c316448ef8 100644 --- a/zig/src/fabric/tp/roce_fake.zig +++ b/zig/src/fabric/tp/roce_fake.zig @@ -93,6 +93,12 @@ pub const Halts = struct { return 0; } + /// ibv_query_qp as the device answers it: ERR once halted, else RTS. + pub fn queryQp(qp: *abi.Qp, attr: *abi.QpAttr, _: c_int, _: *abi.QpInitAttr) callconv(.c) c_int { + attr.qp_state = if (halted[qp.qp_num]) abi.qps_err else abi.qps_rts; + return 0; + } + pub fn pollCq(cq: *abi.Cq, _: c_int, wc: [*]abi.Wc) callconv(.c) c_int { const q = cq.handle; if (Fake.len[q] == 0 or !halted[q]) return 0; @@ -109,6 +115,7 @@ pub const Halts = struct { /// Points a wire's QPs at these handlers. pub fn install(hw: *Hw, w: *rv.VerbsWire) void { w.v.api.ibv_modify_qp = modifyQp; + w.v.query_qp = queryQp; for (&hw.ctx) |*c| c.ops.poll_cq = pollCq; } }; diff --git a/zig/src/fabric/tp/roce_inflight_test.zig b/zig/src/fabric/tp/roce_inflight_test.zig index 130f1e1eab..6441431794 100644 --- a/zig/src/fabric/tp/roce_inflight_test.zig +++ b/zig/src/fabric/tp/roce_inflight_test.zig @@ -37,16 +37,22 @@ test "a failed transport: both QPs go to error before the first flush and every var hw: fake.Hw = .{}; var w = hw.wire(region); Halts.install(&hw, &w); - // two 1 MiB ops posted, neither completed: 2 MiB in flight at the HCAs + // two 1 MiB ops posted, neither completed: 2 MiB in flight at the HCAs, the link busy + try std.testing.expect(w.idle()); try w.post(1, 1, slot); try w.post(0, 2, slot); try std.testing.expectEqual(@as(u64, 2 * slot), w.inflight()); + try std.testing.expect(!w.idle()); try std.testing.expect(w.quiesce(std.time.ns_per_s)); try std.testing.expect(Halts.halted[0] and Halts.halted[1]); try std.testing.expectEqual(@as(u32, 0), Halts.flushed_before_halt); try std.testing.expectEqual(@as(u32, 4), Halts.flushes); try std.testing.expectEqual(@as(u64, 0), w.inflight()); try std.testing.expectEqual(@as(u64, 0), Fake.inflight); + // the line tf-tp-test's transport run reads: the device's states and every reap, all of them flushes + var line: [256]u8 = undefined; + try std.testing.expect(w.idle()); + try std.testing.expectEqualStrings("HCA 0 QP ERR; posted 2, completed 2; HCA 1 QP ERR; posted 2, completed 2; reaped 4 (flushed 4), still posted 0", w.report(&line)); // an aborted transport reports itself failed: its deinit keeps the region Halts.reset(); var r: Roce = .{ .me = 0, .slot_bytes = slot, .timeout_ns = 1, .region = region, .mapping = region, .wire = .{ .verbs = w }, .cpu = null }; diff --git a/zig/src/fabric/tp/roce_verbs.zig b/zig/src/fabric/tp/roce_verbs.zig index 8ed5e3eca4..6917d292be 100644 --- a/zig/src/fabric/tp/roce_verbs.zig +++ b/zig/src/fabric/tp/roce_verbs.zig @@ -73,6 +73,9 @@ const Hca = struct { ring: [ring_len]u32 = undefined, head: u32 = 0, inflight: u64 = 0, + /// Posts made and their completions reaped (the proxy writes them, others read): equal = nothing in flight. + posts: std.atomic.Value(u64) = .init(0), + done: std.atomic.Value(u64) = .init(0), }; /// Parses `name[:gid],name[:gid]` (TF_TP_ROCE_HCA). @@ -170,6 +173,9 @@ pub const VerbsWire = struct { /// The most stripe bytes in flight at once (posted, not completed), for tests and reports. peak: std.atomic.Value(u64) = .init(0), completed: u64 = 0, + /// Completions `quiesce` reaped, and those of them flushed (not sent: the QP was in error). + reaped: u64 = 0, + flushed: u64 = 0, /// Opens the HCAs, registers `region` on each and creates the queue pairs (INIT); `record` then goes to the peer. pub fn open(specs: []const HcaSpec, region: []u8, layout: struct { send: usize, recv: usize, flags: usize, slot: usize }, traffic_class: u8) !VerbsWire { @@ -261,11 +267,43 @@ pub const VerbsWire = struct { var wc: [32]abi.Wc = undefined; const n = h.ctx.ops.poll_cq(h.cq, 32, &wc); if (n < 0) return false; + for (wc[0..@intCast(n)]) |c| w.flushed += @intFromBool(c.status != abi.wc_success); + w.reaped += @intCast(n); if (n > 0) w.popped(h, @intCast(n)) else if (sock.nowNs() - t0 > timeout_ns) return false; }; return true; } + /// Every HCA's posts completed: no request of the region at an HCA (readable beside the proxy). + pub fn idle(w: *const VerbsWire) bool { + for (w.hcas[0..w.n_hca]) |*h| if (h.posts.load(.acquire) != h.done.load(.acquire)) return false; + return true; + } + + /// Each QP's state as its device reports it (ibv_query_qp), then what `quiesce` reaped and what is still posted. + pub fn report(w: *const VerbsWire, buf: []u8) []const u8 { + const names = [_][]const u8{ "RESET", "INIT", "RTR", "RTS", "SQD", "SQE", "ERR" }; + var out: std.Io.Writer = .fixed(buf); + for (w.hcas[0..w.n_hca], 0..) |*h, i| { + var a: abi.QpAttr = .{}; + var init: abi.QpInitAttr = undefined; + const q = w.v.query_qp orelse { + out.print("HCA {d} QP state unknown (no ibv_query_qp); ", .{i}) catch break; + continue; + }; + if (q(h.qp, &a, abi.mask.state, &init) != 0) { + out.print("HCA {d} QP query failed; ", .{i}) catch break; + } else if (a.qp_state >= 0 and a.qp_state < names.len) { + out.print("HCA {d} QP {s}; ", .{ i, names[@intCast(a.qp_state)] }) catch break; + } else out.print("HCA {d} QP state {d}; ", .{ i, a.qp_state }) catch break; + out.print("posted {d}, completed {d}; ", .{ h.posts.load(.acquire), h.done.load(.acquire) }) catch break; + } + var held: u64 = 0; + for (w.hcas[0..w.n_hca]) |h| held += h.outstanding; + out.print("reaped {d} (flushed {d}), still posted {d}", .{ w.reaped, w.flushed, held }) catch {}; + return out.buffered(); + } + pub fn record(w: *const VerbsWire) Record { var r: Record = .{ .n_hca = w.n_hca, .region = @intFromPtr(w.region.ptr), .slot_bytes = w.slot_bytes, .region_len = w.region.len }; for (w.hcas[0..w.n_hca], 0..) |h, i| { @@ -330,6 +368,7 @@ pub const VerbsWire = struct { h.ring[(h.head + h.outstanding) % ring_len] = s.len; h.outstanding += 1; h.inflight += s.len; + _ = h.posts.fetchAdd(1, .release); } if (w.inflight() > w.peak.load(.monotonic)) w.peak.store(w.inflight(), .monotonic); w.posted += 1; @@ -358,6 +397,7 @@ pub const VerbsWire = struct { } h.outstanding -= n; w.completed += n; + _ = h.done.fetchAdd(n, .release); } /// Bytes posted and not yet completed over every HCA. diff --git a/zig/src/fabric/tp/run_transport.sh b/zig/src/fabric/tp/run_transport.sh new file mode 100755 index 0000000000..68ebcddd9a --- /dev/null +++ b/zig/src/fabric/tp/run_transport.sh @@ -0,0 +1,120 @@ +#!/usr/bin/env bash +# The failed-transport check on two hosts (tf-tp-test failfast V transport); `run_transport.sh` alone prints how. +set -u +usage() { + cat << 'TXT' +run_transport.sh MASTER_ADDR PEER_SSH OUTDIR (on rank 0's host, as root; rank 1 runs over ssh) + Two hosts with this script and tf-tp-test (built with the mailbox kernel) at the same path on both. Before the steps, + on both hosts: no GPU process and memlock unlimited. Steps: T0 (suite), T1 (victim 1), T3 (victim 0), fail-fast on; + tf-tp-test injects only when both links are idle. After each step, on both hosts: the kernel log since the step + began and every HCA *err* counter against the step's start. A matching line, a counter that grew, an unexpected rc + or a failure line prints STOP: and ends the run (exit 1), with no later step. +run_transport.sh t2 OUTDIR (one host: TF_TP_ALLOW_HOST_ONLY=1) + T2 (fail-fast off, a hold) runs both ranks on this host over the shm wire; tf-tp-test refuses it on an RDMA wire. +TXT +} +[ $# -ge 2 ] || { usage; exit 2; } +KIT=$(cd "$(dirname "$0")" && pwd) +BIN=$KIT/tf-tp-test +KLOG='EXT4-fs|smmu|iommu|mlx5|AER|nvme|Xid' +BAD='RDMA write of seq|local protection|NOT drained|GPU sync FAILED|REFUSED|refused ERR|panic' +"$BIN" kernel > /dev/null 2>&1 || { echo "STOP: $BIN has no mailbox kernel: rebuild with -Dtp-nvcc=PATH or -Dtp-fatbin=PATH" >&2; exit 2; } + +if [ "${1:-}" = t2 ]; then + OUT=${2:?outdir} + [ "${TF_TP_ALLOW_HOST_ONLY:-}" = 1 ] || { echo "STOP: T2 is host-only: set TF_TP_ALLOW_HOST_ONLY=1" >&2; exit 2; } + mkdir -p "$OUT" + for r in 0 1; do + TF_TP_RANK=$r TF_TP_MASTER=127.0.0.1 TF_TP_PORT=29990 TF_COMM_BACKEND=roce TF_TP_ROCE_WIRE=shm TF_TP_FAILFAST=0 \ + timeout 300 "$BIN" failfast 1 transport 200 5 > "$OUT/T2-r$r.log" 2>&1 & + done + wait + grep -h '^PASS\|^FAIL\|REFUSED\|regions kept' "$OUT"/T2-r*.log + exit 0 +fi + +MASTER=${1:?master address}; PEER=${2:?peer ssh target}; OUT=${3:?outdir} +mkdir -p "$OUT" +stop() { echo "STOP: $*" | tee -a "$OUT/verdict.txt"; exit 1; } + +precheck() { + [ -z "$(nvidia-smi --query-compute-apps=pid --format=csv,noheader 2> /dev/null)" ] || { echo "a GPU process runs"; return 1; } + [ "$(bash -c 'ulimit -l unlimited 2> /dev/null; ulimit -l')" = unlimited ] || { echo "memlock not unlimited"; return 1; } + echo "ok $(hostname) $(uname -r)" +} +counters() { + find -L /sys/class/infiniband/ -maxdepth 5 -path '*ports*counters*' -name '*err*' 2> /dev/null | sort | while read -r f; do + echo "$f $(cat "$f" 2> /dev/null)" + done +} +klog() { # kernel lines since epoch $1 matching $2 (journalctl must answer: an unreadable log is a STOP) + journalctl -k --since "@$1" --no-pager -o short-iso 2> /dev/null | grep -E "$2" + return "${PIPESTATUS[0]}" +} +probe() { # this host's tf-tp-test QPs and MRs once a second until it exits + local pid="" t + for _ in $(seq 200); do pid=$(pgrep -nx tf-tp-test) && break; sleep 0.1; done + [ -z "$pid" ] && return + for t in $(seq 180); do + kill -0 "$pid" 2> /dev/null || break + echo "== t=$t pid $pid" + rdma resource show qp pid "$pid" 2>&1 + rdma resource show mr pid "$pid" 2>&1 + sleep 1 + done +} +on() { # NODE CMD ARGS...: a command or one of the functions above, here (local) or on the peer + local node=$1 fn=$2 def + shift 2 + if [ "$node" = local ]; then "$fn" "$@"; return; fi + def=$(declare -f "$fn" 2> /dev/null) + ssh -o BatchMode=yes "$PEER" "${def:+$def +}$fn $(printf '%q ' "$@")" +} +grew() { # BEFORE AFTER: counters that went up + awk 'NR == FNR { b[$1] = $2; next } ($1 in b) && $2 > b[$1] { print $1 " " b[$1] " -> " $2 }' "$1" "$2" +} + +for n in local peer; do + v=$(on $n precheck) || stop "$n: $v" + echo "precheck $n: $v" | tee -a "$OUT/verdict.txt" +done +on peer test -x "$BIN" || stop "peer: no $BIN (the kit must sit at the same path on both hosts)" + +PORT=29900 +step() { # NAME WANT_RC -- ARGS... + local name=$1 want=$2 n rc rc0 rc1 hits g + shift 3 + PORT=$((PORT + 10)) + declare -A t + for n in local peer; do + t[$n]=$(on $n date +%s) + on $n counters > "$OUT/$name.$n.err0" + done + on local probe > "$OUT/$name.local.rdma.txt" 2>&1 & + on peer probe > "$OUT/$name.peer.rdma.txt" 2>&1 & + local envs=(TF_TP_MASTER="$MASTER" TF_TP_PORT=$PORT TF_COMM_BACKEND=roce TF_TP_ROCE_FALLBACK=error TF_TP_MAILBOX_MAX_KB=1024 + TF_TP_ROCE_FAST=1 TF_TP_FAILFAST_GRACE_S=15) + ssh -o BatchMode=yes "$PEER" "ulimit -l unlimited && env $(printf '%q ' "${envs[@]}") TF_TP_RANK=1 timeout 300 $(printf '%q ' "$BIN" "$@"); echo \"EXIT rank 1 rc \$? at \$(date +%s.%N)\"" > "$OUT/$name-r1.log" 2>&1 & + ( ulimit -l unlimited && env "${envs[@]}" TF_TP_RANK=0 timeout 300 "$BIN" "$@"; echo "EXIT rank 0 rc $? at $(date +%s.%N)" ) > "$OUT/$name-r0.log" 2>&1 + wait + rc0=$(grep -oP '^EXIT rank 0 rc \K[0-9]+' "$OUT/$name-r0.log") + rc1=$(grep -oP '^EXIT rank 1 rc \K[0-9]+' "$OUT/$name-r1.log") + echo "$name: rank 0 rc ${rc0:-?}, rank 1 rc ${rc1:-?} (want $want)" | tee -a "$OUT/verdict.txt" + for n in local peer; do + on $n counters > "$OUT/$name.$n.err1" + g=$(grew "$OUT/$name.$n.err0" "$OUT/$name.$n.err1") + [ -n "$g" ] && stop "$name $n: HCA error counters grew: $g" + hits=$(on $n klog "${t[$n]}" "$KLOG"); rc=$? + [ "$rc" -ne 0 ] && stop "$name $n: the kernel log is unreadable (journalctl rc $rc)" + [ -n "$hits" ] && { echo "$hits" > "$OUT/$name.$n.klog.txt"; stop "$name $n: kernel lines: $hits"; } + done + hits=$(grep -hE "$BAD" "$OUT/$name-r0.log" "$OUT/$name-r1.log") + [ -n "$hits" ] && stop "$name: $hits" + [ "$rc0" = "$want" ] && [ "$rc1" = "$want" ] || stop "$name: rc $rc0 / $rc1, want $want" + grep -h 'halted\|quiesced\|idle before\|regions kept' "$OUT/$name"-r*.log | sed "s/^/ /" +} +step T0 0 -- suite +step T1 70 -- failfast 1 transport 200 +step T3 70 -- failfast 0 transport 200 +echo "PASS transport check (T0, T1, T3)" | tee -a "$OUT/verdict.txt" diff --git a/zig/src/fabric/verbs.zig b/zig/src/fabric/verbs.zig index f3af36587b..8ae58532ef 100644 --- a/zig/src/fabric/verbs.zig +++ b/zig/src/fabric/verbs.zig @@ -50,12 +50,15 @@ pub const Api = struct { /// Exported by rdma-core 34+ but not in every header; used only when present. const GidEx = *const fn (ctx: Ctx, port: u32, index: u32, entry: *abi.GidEntry, flags: u32, size: usize) callconv(.c) c_int; const InOrder = *const fn (qp: *abi.Qp, op: c_int, flags: u32) callconv(.c) c_int; +/// ibv_query_qp: the QP's state as the device holds it (not every library exports it; used only when present). +pub const QueryQp = *const fn (qp: *abi.Qp, attr: *abi.QpAttr, mask: c_int, init: *abi.QpInitAttr) callconv(.c) c_int; pub const Verbs = struct { lib: std.DynLib, api: Api, gid_ex: ?GidEx = null, in_order: ?InOrder = null, + query_qp: ?QueryQp = null, pub fn open(path: ?[:0]const u8) Error!Verbs { var lib = std.DynLib.openZ((path orelse default_path).ptr) catch return error.LibraryUnavailable; @@ -68,6 +71,7 @@ pub const Verbs = struct { var v: Verbs = .{ .lib = lib, .api = api }; v.gid_ex = v.lib.lookup(GidEx, "_ibv_query_gid_ex"); v.in_order = v.lib.lookup(InOrder, "ibv_query_qp_data_in_order"); + v.query_qp = v.lib.lookup(QueryQp, "ibv_query_qp"); return v; } From 467504a1303beaf2ecbf7796bb4b607bf3eacae0 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:28:48 +0700 Subject: [PATCH 19/43] Server: SIGTERM / SIGINT drains before stopping. New generation requests get 503 with Retry-After (code server_restarting; Anthropic overloaded_error) and /health answers "draining" without asking the engine. Requests in progress get TF_DRAIN_S seconds (default 20; a second signal cuts it short), then the engine's optional halt ends its rounds at a boundary and what is left fails with "server restarting: retry shortly". The engine closes before the server is freed, once. --- zig/src/server/drain_test.zig | 302 +++++++++++++++++++++++++++++++ zig/src/server/main.zig | 25 ++- zig/src/server/root.zig | 1 + zig/src/server/routes.zig | 46 ++++- zig/src/server/serve.zig | 65 ++++++- zig/src/server/server.zig | 4 + zig/src/server/status_routes.zig | 2 + 7 files changed, 434 insertions(+), 11 deletions(-) create mode 100644 zig/src/server/drain_test.zig diff --git a/zig/src/server/drain_test.zig b/zig/src/server/drain_test.zig new file mode 100644 index 0000000000..75106abcb9 --- /dev/null +++ b/zig/src/server/drain_test.zig @@ -0,0 +1,302 @@ +//! A stop's drain over HTTP: serve.run with a real SIGTERM, a slow stub engine, bytes as tokens. +const std = @import("std"); +const api = @import("engine_api"); +const json = @import("json.zig"); +const model_text = @import("model_text.zig"); +const serve = @import("serve.zig"); +const cli = @import("cli.zig"); +const routes = @import("routes.zig"); +const Allocator = std.mem.Allocator; +const testing = std.testing; + +/// Slow enough to stop mid-reply: "@short" ends after a dozen tokens, "@long" only when halted or cancelled. +const Slow = struct { + gpa: Allocator, + io: std.Io, + halted: std.atomic.Value(bool) = .init(false), + reason: []const u8 = "", + active: std.atomic.Value(u32) = .init(0), + halts: u32 = 0, + closes: u32 = 0, + /// halts before the first close (the halt must come first) + halts_at_close: u32 = 0, + cancelled: std.atomic.Value(u64) = .init(0), + + fn engine(s: *Slow) api.Engine { + return .{ .ctx = s, .vtable = &.{ .info = info, .submit = submit, .cancel = cancel, .status = status, .memory = memory } }; + } + + fn self(ctx: *anyopaque) *Slow { + return @ptrCast(@alignCast(ctx)); + } + + fn info(_: *anyopaque) api.Info { + return .{ .context_window = 4096, .name = "slow" }; + } + + fn cancel(ctx: *anyopaque, id: api.Id) void { + self(ctx).cancelled.store(id, .release); + } + + fn status(ctx: *anyopaque, out: *api.Status, _: []u32) void { + out.* = .{ .running = self(ctx).active.load(.acquire) }; + } + + fn memory(_: *anyopaque, _: bool) ?api.Memory { + return null; + } + + fn submit(ctx: *anyopaque, id: api.Id, r: *const api.Request, sink: api.Sink) api.SubmitError!void { + const s = self(ctx); + _ = s.active.fetchAdd(1, .acq_rel); + const t = std.Thread.spawn(.{}, run, .{ s, id, r, sink }) catch { + _ = s.active.fetchSub(1, .acq_rel); + return error.Busy; + }; + t.detach(); + } + + fn run(s: *Slow, id: api.Id, r: *const api.Request, sink: api.Sink) void { + defer _ = s.active.fetchSub(1, .acq_rel); + var arena = std.heap.ArenaAllocator.init(s.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const prompt = a.alloc(u8, r.prompt.len) catch return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .failed, .message = "out of memory" } }); + for (r.prompt, prompt) |t, *b| b.* = @truncate(t); + const long = std.mem.indexOf(u8, prompt, "@long") != null; + const word: []const u32 = if (long) &.{ ' ', 'l', 'a' } else &.{ ' ', 'o', 'k' }; + sink.event(sink.ctx, id, &.{ .prefilled = 0 }); + var n: usize = 0; + while (true) : (n += 1) { + if (s.halted.load(.acquire)) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .failed, .message = s.reason } }); + if (s.cancelled.load(.acquire) == id) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .cancelled } }); + if (!long and n == 12) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .stop } }); + sink.event(sink.ctx, id, &.{ .tokens = word }); + std.Io.sleep(s.io, .fromMilliseconds(15), .awake) catch {}; + } + } + + fn halt(ctx: *anyopaque, reason: []const u8) void { + const s = self(ctx); + s.reason = reason; + s.halts += 1; + s.halted.store(true, .release); + // a round boundary: every reply has ended before halt returns + while (s.active.load(.acquire) > 0) std.Io.sleep(s.io, .fromMilliseconds(1), .awake) catch {}; + } + + fn close(ctx: *anyopaque) void { + const s = self(ctx); + if (s.closes == 0) s.halts_at_close = s.halts; + s.closes += 1; + } +}; + +/// Bytes as tokens; a prompt renders as its messages' JSON, so the engine sees the client's marker. +const TestText = struct { + fn text(t: *@This()) model_text.Text { + return .{ .ctx = t, .vtable = &.{ .encode = encode, .decode = decode, .token_id = tokenId, .token_string = tokenString, .vocab_size = vocabSize, .eos_ids = eosIds, .render = render, .template_source = templateSource } }; + } + + fn encode(_: *anyopaque, a: Allocator, input: []const u8, _: bool) model_text.Error![]u32 { + const ids = try a.alloc(u32, input.len); + for (input, ids) |byte, *id| id.* = byte; + return ids; + } + + fn decode(_: *anyopaque, a: Allocator, ids: []const u32) Allocator.Error![]u8 { + const decoded = try a.alloc(u8, ids.len); + for (ids, decoded) |id, *byte| byte.* = @intCast(id); + return decoded; + } + + fn tokenId(_: *anyopaque, _: []const u8) ?u32 { + return null; + } + + fn tokenString(_: *anyopaque, a: Allocator, id: u32) Allocator.Error![]u8 { + return std.fmt.allocPrint(a, "{d}", .{id}); + } + + fn vocabSize(_: *anyopaque) u32 { + return 256; + } + + fn eosIds(_: *anyopaque) []const u32 { + return &.{}; + } + + fn render(_: *anyopaque, a: Allocator, messages: json.Value, _: model_text.RenderOptions, _: *[]const u8) model_text.Error![]u8 { + return json.stringify(a, messages, .{}) catch error.OutOfMemory; + } + + fn templateSource(_: *anyopaque) []const u8 { + return ""; + } +}; + +var listen_port: std.atomic.Value(u16) = .init(0); + +fn onListen(port: u16) void { + listen_port.store(port, .release); +} + +const Reply = struct { status: u16, head: []const u8, body: []const u8 }; + +/// One request on its own connection (closed by us after it), the whole response read. +fn request(a: Allocator, io: std.Io, port: u16, method: []const u8, path: []const u8, body: []const u8) !Reply { + const addr: std.Io.net.IpAddress = .{ .ip4 = .loopback(port) }; + const stream = try addr.connect(io, .{ .mode = .stream }); + defer stream.close(io); + const head = try std.fmt.allocPrint(a, "{s} {s} HTTP/1.1\r\nHost: x\r\nContent-Type: application/json\r\nContent-Length: {d}\r\nConnection: close\r\n\r\n{s}", .{ method, path, body.len, body }); + var wbuf: [4096]u8 = undefined; + var w = std.Io.net.Stream.Writer.init(stream, io, &wbuf); + try w.interface.writeAll(head); + try w.interface.flush(); + var rbuf: [16384]u8 = undefined; + var r = std.Io.net.Stream.Reader.init(stream, io, &rbuf); + const text = try r.interface.allocRemaining(a, .unlimited); + const split = std.mem.indexOf(u8, text, "\r\n\r\n") orelse return error.Head; + var payload: []const u8 = text[split + 4 ..]; + if (std.mem.indexOf(u8, text[0..split], "chunked") != null) payload = try unchunk(a, payload); + return .{ .status = try std.fmt.parseInt(u16, text[9..12], 10), .head = text[0..split], .body = payload }; +} + +fn unchunk(a: Allocator, s: []const u8) ![]const u8 { + var out: std.ArrayList(u8) = .empty; + var i: usize = 0; + while (i < s.len) { + const eol = std.mem.indexOfPos(u8, s, i, "\r\n") orelse break; + const n = try std.fmt.parseInt(usize, std.mem.trim(u8, s[i..eol], " "), 16); + if (n == 0) break; + try out.appendSlice(a, s[eol + 2 .. eol + 2 + n]); + i = eol + 2 + n + 2; + } + return out.items; +} + +const Background = struct { + a: Allocator, + io: std.Io, + port: u16, + body: []const u8, + reply: ?Reply = null, + failed: bool = false, + + fn run(b: *Background) void { + b.reply = request(b.a, b.io, b.port, "POST", "/v1/chat/completions", b.body) catch { + b.failed = true; + return; + }; + } +}; + +fn chatBody(a: Allocator, marker: []const u8, stream: bool) ![]const u8 { + return std.fmt.allocPrint(a, "{{\"messages\":[{{\"role\":\"user\",\"content\":\"go {s}\"}}],\"stream\":{},\"max_tokens\":4000}}", .{ marker, stream }); +} + +const Served = struct { + code: u8 = 255, + fn run(x: *Served, gpa: Allocator, io: std.Io, setup: serve.Setup) void { + x.code = serve.run(gpa, io, .{ .host = "127.0.0.1", .port = 0 }, setup); + } +}; + +/// serve.run on its own thread with `drain_s`, the engine `slow`; the port once it listens. +fn start(gpa: Allocator, io: std.Io, env: *std.process.Environ.Map, slow: *Slow, tt: *TestText, x: *Served) !std.Thread { + listen_port.store(0, .release); + const t = try std.Thread.spawn(.{}, Served.run, .{ x, gpa, io, serve.Setup{ + .engine = slow.engine(), + .text = tt.text(), + .served = "slow", + .environ = env, + .on_listen = onListen, + .stop = .{ .ctx = slow, .halt = Slow.halt, .close = Slow.close }, + } }); + while (listen_port.load(.acquire) == 0) std.Io.sleep(io, .fromMilliseconds(2), .awake) catch {}; + return t; +} + +fn waitActive(io: std.Io, slow: *Slow, n: u32) void { + while (slow.active.load(.acquire) < n) std.Io.sleep(io, .fromMilliseconds(2), .awake) catch {}; +} + +test "SIGTERM drains: 503 + Retry-After and /health draining while a reply finishes; at the deadline the rest end \"server restarting\" at a round boundary, then the engine closes" { + const gpa = testing.allocator; + const io = testing.io; + var tt: TestText = .{}; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var env = std.process.Environ.Map.init(gpa); + defer env.deinit(); + try env.put("TENSORFOLD_NO_LIVE", "1"); + + // 1. the reply in progress finishes within the drain: no halt, a whole reply; new work refused meanwhile + { + try env.put("TF_DRAIN_S", "10"); + var slow: Slow = .{ .gpa = gpa, .io = io }; + var x: Served = .{}; + const t = try start(gpa, io, &env, &slow, &tt, &x); + const port = listen_port.load(.acquire); + var bg: Background = .{ .a = a, .io = io, .port = port, .body = try chatBody(a, "@short", true) }; + const bt = try std.Thread.spawn(.{}, Background.run, .{&bg}); + waitActive(io, &slow, 1); + try std.posix.raise(.TERM); + // the drain has begun once a new request is refused + var refused: Reply = undefined; + while (true) { + refused = try request(a, io, port, "POST", "/v1/chat/completions", try chatBody(a, "@short", false)); + if (refused.status == 503) break; + try testing.expectEqual(@as(u16, 200), refused.status); // only before the signal lands + } + try testing.expect(std.mem.indexOf(u8, refused.head, "Retry-After: 30") != null); + try testing.expect(std.mem.indexOf(u8, refused.body, "\"code\": \"server_restarting\"") != null); + const anth = try request(a, io, port, "POST", "/v1/messages", "{\"model\":\"x\",\"max_tokens\":5,\"messages\":[{\"role\":\"user\",\"content\":\"hi\"}]}"); + try testing.expectEqual(@as(u16, 503), anth.status); + try testing.expect(std.mem.indexOf(u8, anth.body, "overloaded_error") != null); + const health = try request(a, io, port, "GET", "/health", ""); + try testing.expectEqual(@as(u16, 503), health.status); + try testing.expect(std.mem.indexOf(u8, health.head, "Retry-After: 30") != null); + try testing.expect(std.mem.indexOf(u8, health.body, "\"status\": \"draining\"") != null); + bt.join(); + t.join(); + try testing.expect(!bg.failed); + try testing.expectEqual(@as(u16, 200), bg.reply.?.status); + try testing.expect(std.mem.indexOf(u8, bg.reply.?.body, "\"finish_reason\": \"stop\"") != null); + try testing.expect(std.mem.indexOf(u8, bg.reply.?.body, "restarting") == null); + try testing.expect(std.mem.endsWith(u8, std.mem.trimEnd(u8, bg.reply.?.body, "\r\n"), "data: [DONE]")); + try testing.expectEqual(@as(u8, 0), x.code); + try testing.expectEqual(@as(u32, 0), slow.halts); + try testing.expectEqual(@as(u32, 1), slow.closes); + } + // 2. replies that outlast the drain end "server restarting"; the engine is halted, then closed once + { + try env.put("TF_DRAIN_S", "0.3"); + var slow: Slow = .{ .gpa = gpa, .io = io }; + var x: Served = .{}; + const t = try start(gpa, io, &env, &slow, &tt, &x); + const port = listen_port.load(.acquire); + var streamed: Background = .{ .a = a, .io = io, .port = port, .body = try chatBody(a, "@long", true) }; + var whole: Background = .{ .a = a, .io = io, .port = port, .body = try chatBody(a, "@long", false) }; + const st = try std.Thread.spawn(.{}, Background.run, .{&streamed}); + const wt = try std.Thread.spawn(.{}, Background.run, .{&whole}); + waitActive(io, &slow, 2); + try std.posix.raise(.TERM); + st.join(); + wt.join(); + t.join(); + try testing.expectEqual(@as(u8, 0), x.code); + try testing.expectEqual(@as(u32, 1), slow.halts); + try testing.expectEqual(@as(u32, 1), slow.closes); + try testing.expectEqual(@as(u32, 1), slow.halts_at_close); + try testing.expectEqualStrings(routes.restarting, slow.reason); + const sb = streamed.reply.?.body; + try testing.expectEqual(@as(u16, 200), streamed.reply.?.status); + try testing.expect(std.mem.indexOf(u8, sb, "\"type\": \"server_error\"") != null); + try testing.expect(std.mem.indexOf(u8, sb, routes.restarting) != null); + try testing.expect(std.mem.indexOf(u8, sb, "data: [DONE]") != null); + try testing.expect(whole.reply.?.status >= 500); + try testing.expect(std.mem.indexOf(u8, whole.reply.?.body, routes.restarting) != null); + } +} diff --git a/zig/src/server/main.zig b/zig/src/server/main.zig index 4b85fca8cd..2b0de6f365 100644 --- a/zig/src/server/main.zig +++ b/zig/src/server/main.zig @@ -77,8 +77,10 @@ pub fn main(init: std.process.Init) !u8 { return e; } orelse return fail(problem); defer up.text.deinit(); - defer up.engine.close(up.engine.ctx); + var closer: Closer = .{ .opened = up.engine }; + defer closer.closeOnce(); return serve.run(gpa, io, args, .{ + .stop = .{ .ctx = &closer, .halt = if (up.engine.halt != null) Closer.halt else null, .close = Closer.close }, .engine = up.engine.engine, .text = up.text.text(), .served = hub.servedName(args.name, args.model, dir), @@ -89,6 +91,27 @@ pub fn main(init: std.process.Init) !u8 { }); } +/// Closes once: serve.run closes the engine before freeing the server, main's defer covers paths that end before it. +const Closer = struct { + opened: engines.Opened, + closed: bool = false, + + fn closeOnce(c: *Closer) void { + if (c.closed) return; + c.closed = true; + c.opened.close(c.opened.ctx); + } + + fn close(ctx: *anyopaque) void { + closeOnce(@ptrCast(@alignCast(ctx))); + } + + fn halt(ctx: *anyopaque, reason: []const u8) void { + const c: *Closer = @ptrCast(@alignCast(ctx)); + if (c.opened.halt) |h| h(c.opened.ctx, reason); + } +}; + /// HfText.load with every failure named in `problem`, so main tells the text's failure from the engine's. fn loadText(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, template: ?[]const u8, pa: std.mem.Allocator, problem: *[]const u8) !*hf_text.HfText { return hf_text.HfText.load(gpa, io, dir, template, pa, problem) catch |e| { diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 0c0f60224d..af6e3057a5 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -43,6 +43,7 @@ test { _ = @import("stream_preflight_test.zig"); _ = @import("stream_dsml_test.zig"); _ = @import("logprobs_test.zig"); + _ = @import("drain_test.zig"); _ = @import("late_system_test.zig"); _ = @import("compact_test.zig"); _ = @import("completion_prompt_test.zig"); diff --git a/zig/src/server/routes.zig b/zig/src/server/routes.zig index 44981cdda9..8fe21318e1 100644 --- a/zig/src/server/routes.zig +++ b/zig/src/server/routes.zig @@ -63,15 +63,51 @@ fn discardBody(conn: *Conn, a: Allocator) void { fn post(srv: *Server, conn: *Conn, a: Allocator) void { const route = auth.routePath(conn.path); - if (std.mem.eql(u8, route, "/v1/slide/learn")) return slide.learn(srv, conn, a); - if (std.mem.endsWith(u8, route, "/decisions")) return decisions.post(srv, conn, a); - if (anthropic.route(conn.path)) return anthropic.post(srv, conn, a); - if (responses.route(route)) |rid| if (rid.len == 0) return responses.post(srv, conn, a); for (tokens.paths) |r| if (std.mem.eql(u8, route, r)) return tokens.post(srv, conn, a, std.mem.endsWith(u8, route, "/detokenize")); + const is_anthropic = anthropic.route(conn.path); + const is_responses = if (responses.route(route)) |rid| rid.len == 0 else false; + const is_decisions = std.mem.endsWith(u8, route, "/decisions"); + const is_slide = std.mem.eql(u8, route, "/v1/slide/learn"); const is_chat = std.mem.endsWith(u8, route, "/chat/completions"); - if (!is_chat and !std.mem.endsWith(u8, route, "/completions")) { + if (!is_slide and !is_anthropic and !is_responses and !is_decisions and !is_chat and !std.mem.endsWith(u8, route, "/completions")) { discardBody(conn, a); return unknown(conn, a); } + // counted before the drain flag is read: a stop that sees no request in progress sees every later one refused + _ = srv.generating.fetchAdd(1, .seq_cst); + defer _ = srv.generating.fetchSub(1, .seq_cst); + if (srv.draining.load(.seq_cst)) { + discardBody(conn, a); + return draining(conn, is_anthropic); + } + if (is_slide) return slide.learn(srv, conn, a); + if (is_decisions) return decisions.post(srv, conn, a); + if (is_anthropic) return anthropic.post(srv, conn, a); + if (is_responses) return responses.post(srv, conn, a); openai.post(srv, conn, a, is_chat); } + +/// Seconds a refused client is told to wait while the server drains for a restart (Retry-After). +pub const retry_after_s = 30; +pub const restarting = "server restarting: retry shortly"; + +/// A request that arrives while a stop drains: 503, Retry-After, the API's own error shape; the connection closes. +pub fn draining(conn: *Conn, anthropic_shape: bool) void { + const body = if (anthropic_shape) + "{\"type\": \"error\", \"error\": {\"type\": \"overloaded_error\", \"message\": \"" ++ restarting ++ "\"}}" + else + "{\"error\": {\"message\": \"" ++ restarting ++ "\", \"type\": \"server_error\", \"param\": null, \"code\": \"server_restarting\"}}"; + sendDraining(conn, body); +} + +/// 503 with Retry-After and Connection: close, `body` as JSON. +pub fn sendDraining(conn: *Conn, body: []const u8) void { + conn.close = true; + conn.startResponse(503, null) catch return; + conn.addHeader("Content-Type", "application/json") catch return; + conn.addHeader("Retry-After", std.fmt.comptimePrint("{d}", .{retry_after_s})) catch return; + var len: [24]u8 = undefined; + conn.addHeader("Content-Length", std.fmt.bufPrint(&len, "{d}", .{body.len}) catch unreachable) catch return; + conn.addHeader("Connection", "close") catch return; + conn.finish(body) catch {}; +} diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index 6974118c28..e83a9bbf97 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -9,15 +9,19 @@ const live = @import("live.zig"); const listener_mod = @import("listener.zig"); const model_text = @import("model_text.zig"); const server_mod = @import("server.zig"); +const routes = @import("routes.zig"); const Allocator = std.mem.Allocator; const posix = std.posix; var stop_requested: std.atomic.Value(bool) = .init(false); /// The signal that asked for the stop (0: none yet), for the shutdown line var stop_signal: std.atomic.Value(i32) = .init(0); +/// Stop signals so far: a second one cuts the drain short +var stop_count: std.atomic.Value(u32) = .init(0); fn onStop(sig: posix.SIG) callconv(.c) void { stop_signal.store(@intCast(@backingInt(sig)), .release); + _ = stop_count.fetchAdd(1, .acq_rel); stop_requested.store(true, .release); } @@ -84,6 +88,15 @@ pub const Setup = struct { on_listen: ?*const fn (port: u16) void = null, /// The checkpoint family's own chat rules (family.zig), or null. family: ?@import("family.zig").Family = null, + /// Run by `run` so the engine stops before the server is freed; null: the caller closes it after `run`. + stop: ?EngineStop = null, +}; + +/// `halt` (optional) ends the rounds at a boundary at the drain's deadline, so `close` never meets a running round. +pub const EngineStop = struct { + ctx: *anyopaque, + halt: ?*const fn (ctx: *anyopaque, reason: []const u8) void = null, + close: *const fn (ctx: *anyopaque) void, }; fn env(s: Setup, name: []const u8) ?[]const u8 { @@ -93,6 +106,9 @@ fn env(s: Setup, name: []const u8) ?[]const u8 { /// Serves until a stop signal; the process exit status. pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { + stop_requested.store(false, .release); + stop_count.store(0, .release); + stop_signal.store(0, .release); var arena: std.heap.ArenaAllocator = .init(gpa); defer arena.deinit(); const a = arena.allocator(); @@ -115,6 +131,10 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { std.debug.print("tensorfold: TENSORFOLD_SEED_SALT={s}: an integer\n", .{t}); return 1; } else 0; + const drain_s = seconds(env(s, "TF_DRAIN_S"), 20) orelse { + std.debug.print("tensorfold: TF_DRAIN_S: seconds, a number >= 0\n", .{}); + return 1; + }; const config: server_mod.Config = .{ .served_name = s.served, .model_dir = s.model_dir, @@ -177,21 +197,56 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { var ticker: live.Ticker = .{ .engine = s.engine }; const drawing = live.wanted(io, if (env(s, "TENSORFOLD_NO_LIVE")) |v| std.mem.eql(u8, v, "1") else false); if (drawing) ticker.start(); - const accept = std.Thread.spawn(.{}, server_mod.Server.serve, .{ srv, lis, &stop_requested }) catch return 1; + // the listener stays open through the drain: new requests get 503 + Retry-After, /health "draining" + var accept_stop: std.atomic.Value(bool) = .init(false); + const accept = std.Thread.spawn(.{}, server_mod.Server.serve, .{ srv, lis, &accept_stop }) catch return 1; + while (!stop_requested.load(.acquire)) std.Io.sleep(io, .fromMilliseconds(50), .awake) catch {}; + drain(srv, io, s.stop, drain_s); + accept_stop.store(true, .release); accept.join(); lis.close(); - // before the replies still running fail (their rounds' CUDA calls fail once the exit tears the driver down) - log.line("shutting down ({s}): no new requests; {d} connection(s) open, waiting up to 2 s before exiting", .{ signalName(stop_signal.load(.acquire)), srv.open_connections.load(.acquire) }); if (drawing) ticker.finish(); + // the engine stops while nothing runs on it: no round meets a CUDA driver torn down by the exit + if (s.stop) |st| st.close(st.ctx); var waited: u32 = 0; while (srv.open_connections.load(.acquire) > 0 and waited < 40) : (waited += 1) std.Io.sleep(io, .fromMilliseconds(50), .awake) catch {}; if (srv.open_connections.load(.acquire) > 0) { - log.line("shutting down: exiting with {d} connection(s) still open; their requests fail", .{srv.open_connections.load(.acquire)}); - std.process.exit(0); // replies still open: end without freeing what they read + log.line("shutting down: exiting with {d} idle connection(s) still open", .{srv.open_connections.load(.acquire)}); + std.process.exit(0); // connections still open: end without freeing what they read } + log.line("shutting down: stopped", .{}); return 0; } +/// Work in progress gets `drain_s` to finish (new work 503s) before the rest is halted; a second signal cuts it short. +fn drain(srv: *server_mod.Server, io: std.Io, stop: ?EngineStop, drain_s: f64) void { + srv.draining.store(true, .seq_cst); + const signals = stop_count.load(.acquire); + log.line("shutting down ({s}): draining: new requests get 503; {d} request(s) in progress may finish within {d:.1} s", .{ signalName(stop_signal.load(.acquire)), srv.generating.load(.seq_cst), drain_s }); + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + const deadline = t0 + @as(i96, @intFromFloat(drain_s * 1e9)); + while (srv.generating.load(.seq_cst) > 0 and std.Io.Clock.awake.now(io).toNanoseconds() < deadline and stop_count.load(.acquire) == signals) + std.Io.sleep(io, .fromMilliseconds(20), .awake) catch {}; + const left = srv.generating.load(.seq_cst); + if (left == 0) { + log.line("shutting down: drained", .{}); + return; + } + const cut = if (stop_count.load(.acquire) != signals) "a second stop signal" else "the drain deadline (TF_DRAIN_S)"; + log.line("shutting down: {s}: {d} request(s) still in progress end \"{s}\" at the next round boundary", .{ cut, left, routes.restarting }); + if (stop) |st| if (st.halt) |h| h(st.ctx, routes.restarting); + var waited: u32 = 0; + while (srv.generating.load(.seq_cst) > 0 and waited < 100) : (waited += 1) std.Io.sleep(io, .fromMilliseconds(20), .awake) catch {}; +} + +/// A seconds setting (TF_DRAIN_S): unset or blank keeps the default; null: not a number >= 0. +fn seconds(raw: ?[]const u8, default: f64) ?f64 { + const t = std.mem.trim(u8, raw orelse "", " \t"); + if (t.len == 0) return default; + const x = std.fmt.parseFloat(f64, t) catch return null; + return if (x >= 0) x else null; +} + /// ``--host`` as an address: an IP literal, else a name the resolver knows (localhost first). fn resolve(io: std.Io, host: []const u8, port: u16) ?std.Io.net.IpAddress { if (std.Io.net.IpAddress.parse(host, port)) |ip| return ip else |_| {} diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index fc90467612..79efe48e94 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -89,6 +89,10 @@ pub const Server = struct { next_id: std.atomic.Value(u64) = .init(1), /// Connections open now; a stop waits for them before freeing what they read. open_connections: std.atomic.Value(u32) = .init(0), + /// A stop is draining (serve.zig): new generation requests get 503 with Retry-After, /health answers "draining" + draining: std.atomic.Value(bool) = .init(false), + /// Generation requests in progress (routes.post), what a drain waits for; counted before `draining` is read + generating: std.atomic.Value(u32) = .init(0), /// The idle keepalive on the engine's queue, armed by --keep-warm; null when off or not Metal. keepalive: ?*api.keepalive.Keepalive = null, preparing: std.atomic.Value(i64) = .init(0), diff --git a/zig/src/server/status_routes.zig b/zig/src/server/status_routes.zig index 89ee5c8488..9c1f93a07e 100644 --- a/zig/src/server/status_routes.zig +++ b/zig/src/server/status_routes.zig @@ -11,6 +11,8 @@ const Allocator = std.mem.Allocator; /// With keys on, health names nothing about the model. pub fn health(srv: *Server, conn: *Conn, a: Allocator) !void { + // draining: the engine is not asked, it may be stopping + if (srv.draining.load(.seq_cst)) return routes.sendDraining(conn, try std.fmt.allocPrint(a, "{{\"status\": \"draining\", \"requests_in_progress\": {d}}}", .{srv.generating.load(.seq_cst)})); if (srv.keys) |k| if (k.enabled()) return conn.sendJson(200, "{\"status\": \"ok\"}"); const o = try json.newObject(a); try o.put(a, "status", .{ .string = "ok" }); From 602aad59afb3e3ee037f246e6345a5969d2f24b6 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 21:41:15 +0700 Subject: [PATCH 20/43] Fix RoCE test smaps lookup for merged mappings --- zig/src/fabric/tp/roce_inflight_test.zig | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/zig/src/fabric/tp/roce_inflight_test.zig b/zig/src/fabric/tp/roce_inflight_test.zig index 6441431794..fd7be53262 100644 --- a/zig/src/fabric/tp/roce_inflight_test.zig +++ b/zig/src/fabric/tp/roce_inflight_test.zig @@ -67,7 +67,7 @@ test "a QP whose inline capacity cannot take the 4-byte flag post is refused" { try std.testing.expect(rv.inlineFits(.{ .max_send_wr = 1, .max_recv_wr = 1, .max_send_sge = 1, .max_recv_sge = 1, .max_inline_data = 4 })); } -/// /proc/self/smaps' VmFlags of the mapping that starts at `at`. +/// /proc/self/smaps' VmFlags of the mapping that contains `at`. fn vmFlags(at: usize, buf: []u8) ![]const u8 { const fd = std.c.open("/proc/self/smaps", .{}); if (fd < 0) return error.NoSmaps; @@ -78,12 +78,15 @@ fn vmFlags(at: usize, buf: []u8) ![]const u8 { if (got <= 0) break; n += @intCast(got); } - var head: [32]u8 = undefined; - const want = try std.fmt.bufPrint(&head, "{x}-", .{at}); var lines = std.mem.splitScalar(u8, buf[0..n], '\n'); var inside = false; while (lines.next()) |l| { - if (std.mem.startsWith(u8, l, want)) inside = true; + const sp = std.mem.indexOfScalar(u8, l, ' ') orelse l.len; + if (std.mem.indexOfScalar(u8, l[0..sp], '-')) |dash| { + const lo = std.fmt.parseInt(usize, l[0..dash], 16) catch continue; + const hi = std.fmt.parseInt(usize, l[dash + 1 .. sp], 16) catch continue; + inside = lo <= at and at < hi; + } if (inside and std.mem.startsWith(u8, l, "VmFlags:")) return l; } return error.NotFound; @@ -100,7 +103,8 @@ test "the region is DONTFORK and NOHUGEPAGE before registration" { const before = try vmFlags(@intFromPtr(m.ptr), buf); try std.testing.expect(std.mem.indexOf(u8, before, " dc") == null and std.mem.indexOf(u8, before, " nh") == null); try rv.prepareRegion(m); - const flags = try vmFlags(@intFromPtr(m.ptr), buf); + // An interior page stays inside the VMA even when a neighbouring mapping merges with it. + const flags = try vmFlags(@intFromPtr(m.ptr) + 4096, buf); try std.testing.expect(std.mem.indexOf(u8, flags, " dc") != null); try std.testing.expect(std.mem.indexOf(u8, flags, " nh") != null); } From 23bf7d88c200395b5fa0bc05355dee30e82b73c6 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:29:01 +0700 Subject: [PATCH 21/43] Server tests: a stop that lands between a request's count and its drain check refuses the request (503) and the drain waits for it. A tests-only seam (routes.counted_hook) delivers SIGTERM in that gap; with the drain check moved before the count the test fails. --- zig/src/server/drain_test.zig | 62 +++++++++++++++++++++++++++++++++++ zig/src/server/routes.zig | 4 +++ 2 files changed, 66 insertions(+) diff --git a/zig/src/server/drain_test.zig b/zig/src/server/drain_test.zig index 75106abcb9..2bdc9f4a59 100644 --- a/zig/src/server/drain_test.zig +++ b/zig/src/server/drain_test.zig @@ -300,3 +300,65 @@ test "SIGTERM drains: 503 + Retry-After and /health draining while a reply finis try testing.expect(std.mem.indexOf(u8, whole.reply.?.body, routes.restarting) != null); } } + +/// The race seam's state: the stop delivered between a request's count and its drain check. +const Gap = struct { + var io: std.Io = undefined; + var port: u16 = 0; + var fired: std.atomic.Value(bool) = .init(false); + var saw_draining: bool = false; + var drain_waited: bool = false; + var served: *Served = undefined; + + fn hook() void { + if (fired.swap(true, .acq_rel)) return; + std.posix.raise(.TERM) catch return; + var arena = std.heap.ArenaAllocator.init(testing.allocator); + defer arena.deinit(); + // the drain has begun (it reads the count after it sets the flag): /health says so + var tries: u32 = 0; + while (tries < 500) : (tries += 1) { + const h = request(arena.allocator(), io, port, "GET", "/health", "") catch break; + if (h.status == 503 and std.mem.indexOf(u8, h.body, "draining") != null) { + saw_draining = true; + break; + } + std.Io.sleep(io, .fromMilliseconds(5), .awake) catch {}; + } + // this request is counted: the drain may not finish while it is between its count and its check + std.Io.sleep(io, .fromMilliseconds(300), .awake) catch {}; + drain_waited = served.code == 255; + } +}; + +test "a stop that lands between a request's count and its drain check: the request is refused (503) and the drain waited for it, never lost" { + const gpa = testing.allocator; + const io = testing.io; + var tt: TestText = .{}; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var env = std.process.Environ.Map.init(gpa); + defer env.deinit(); + try env.put("TENSORFOLD_NO_LIVE", "1"); + try env.put("TF_DRAIN_S", "10"); + var slow: Slow = .{ .gpa = gpa, .io = io }; + var x: Served = .{}; + const t = try start(gpa, io, &env, &slow, &tt, &x); + Gap.io = io; + Gap.port = listen_port.load(.acquire); + Gap.served = &x; + Gap.fired.store(false, .release); + routes.counted_hook = Gap.hook; + defer routes.counted_hook = null; + const r = try request(a, io, Gap.port, "POST", "/v1/chat/completions", try chatBody(a, "@short", false)); + t.join(); + try testing.expect(Gap.fired.load(.acquire)); + try testing.expect(Gap.saw_draining); + try testing.expect(Gap.drain_waited); + try testing.expectEqual(@as(u16, 503), r.status); + try testing.expect(std.mem.indexOf(u8, r.body, "server_restarting") != null); + try testing.expectEqual(@as(u32, 0), slow.closes -| 1); // closed once, after + try testing.expectEqual(@as(u32, 0), slow.active.load(.acquire)); // nothing reached the engine + try testing.expectEqual(@as(u8, 0), x.code); +} diff --git a/zig/src/server/routes.zig b/zig/src/server/routes.zig index 8fe21318e1..75b3b18a18 100644 --- a/zig/src/server/routes.zig +++ b/zig/src/server/routes.zig @@ -76,6 +76,7 @@ fn post(srv: *Server, conn: *Conn, a: Allocator) void { // counted before the drain flag is read: a stop that sees no request in progress sees every later one refused _ = srv.generating.fetchAdd(1, .seq_cst); defer _ = srv.generating.fetchSub(1, .seq_cst); + if (@import("builtin").is_test) if (counted_hook) |h| h(); if (srv.draining.load(.seq_cst)) { discardBody(conn, a); return draining(conn, is_anthropic); @@ -87,6 +88,9 @@ fn post(srv: *Server, conn: *Conn, a: Allocator) void { openai.post(srv, conn, a, is_chat); } +/// Tests only: the gap between a request's count and its drain check, where a stop must still see it counted. +pub var counted_hook: ?*const fn () void = null; + /// Seconds a refused client is told to wait while the server drains for a restart (Retry-After). pub const retry_after_s = 30; pub const restarting = "server restarting: retry shortly"; From 890dfc6442e1aa04d6ada6b9978a3aadef66495e Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Mon, 12 Oct 2026 06:30:24 +0700 Subject: [PATCH 22/43] TP RoCE: run_hosts.sh gives NCCL the RoCE picker's ports when neither TF_TP_ROCE_HCA nor NCCL_IB_HCA is set - `tf-tp-test hca` (host only, no GPU): both ranks run the RoCE pick over the bootstrap, the same rule the wire uses, and each prints its own devices as `NCCL_IB_HCA=name,...` (no GIDs). It exits 1 with the reason when the host has no usable port or no port on a /24 the peer shares, so no rank falls back to its first ports for NCCL. - run_hosts.sh runs it when neither variable is set, exports the names for every step and logs each rank's pick; a failed or malformed `hca` stops the runner (exit 3) with the reason; `hca` is capped at TF_TP_PLAN_TIMEOUT + 31 s, so its own reason (no peer joined in time, port in use, the peer closed, a bootstrap error) comes first. An explicit NCCL_IB_HCA, an explicit TF_TP_ROCE_HCA (NCCL_IB_HCA derived from it) or both keep their behaviour; the many-devices warning is gone. Both hosts use the same mode: with one naming its ports and the other not, the latter's `hca` finds no peer and stops after TF_TP_PLAN_TIMEOUT with that reason (README). - roce_pick: `select` / `selectFrom` return how the devices were chosen (`resolve` is unchanged for the wire), `joinNames` and `hcaLine` give the NCCL list. Host tests: the port scan over a stub verbs table (state, Ethernet, the IPv4 RoCE v2 GID index), then both ranks through a real bootstrap on the two hosts #636's receipt reports (modelled, documentation addresses), an explicit rank, and a host without ports. --- zig/src/fabric/tp/README.md | 8 ++ zig/src/fabric/tp/gpu_main.zig | 40 +++++- zig/src/fabric/tp/roce_pick.zig | 216 ++++++++++++++++++++++++++++++-- zig/src/fabric/tp/run_hosts.sh | 26 +++- 4 files changed, 276 insertions(+), 14 deletions(-) diff --git a/zig/src/fabric/tp/README.md b/zig/src/fabric/tp/README.md index a0b4af5342..7ad00ef7d5 100644 --- a/zig/src/fabric/tp/README.md +++ b/zig/src/fabric/tp/README.md @@ -36,6 +36,14 @@ Each device pairs with the peer device on the same IPv4 /24. `NCCL_IB_HCA` should name the same devices without the `:gid` suffix. In NCCL, `:N` after a device name selects port `N`, not a GID index. +`run_hosts.sh` sets `NCCL_IB_HCA` from `TF_TP_ROCE_HCA` when only that is set. +When neither is set, it first runs `tf-tp-test hca` on both ranks: the RoCE pick, made with the peer over the bootstrap. +The same rule picks the wire's devices, and each rank prints its own as `NCCL_IB_HCA=name,...`. +`hca` needs no GPU, and it exits 1 when no port shares a /24 with the peer, so no host falls back to its first ports. +The runner logs each rank's pick and stops with that reason when `hca` fails. +Both hosts must use the same mode: if one names its ports and the other names none, only the latter runs `hca`. +It finds no peer and stops after `TF_TP_PLAN_TIMEOUT` (default 300 s) with that reason. + Two DGX Spark nodes, with the second ConnectX-7 port of each cabled to the other (both PCIe functions of that port), and the first port cabled to a third node: diff --git a/zig/src/fabric/tp/gpu_main.zig b/zig/src/fabric/tp/gpu_main.zig index 07d7c7afb9..82fc1edd5b 100644 --- a/zig/src/fabric/tp/gpu_main.zig +++ b/zig/src/fabric/tp/gpu_main.zig @@ -1,4 +1,4 @@ -//! `tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | kernel`: one rank of a TP=2 test. +//! `tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | hca | kernel`: one rank of TP=2. const std = @import("std"); const cuda = @import("cuda"); const tp = @import("tp"); @@ -11,7 +11,7 @@ const max_bytes = (4 << 20) + 4096; pub fn main(init: std.process.Init) !u8 { const args = try init.minimal.args.toSlice(init.arena.allocator()); if (args.len < 2) { - std.debug.print("usage: tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | kernel\n", .{}); + std.debug.print("usage: tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | hca | kernel\n", .{}); return 2; } // `kernel`: whether the build embeds the mailbox kernel (-Dtp-nvcc or -Dtp-fatbin), for scripts to check first @@ -21,6 +21,7 @@ pub fn main(init: std.process.Init) !u8 { } const cfg = try tp.Config.fromEnv(); if (std.mem.eql(u8, args[1], "roce-host")) return @import("roce_host.zig").run(cfg, init.gpa); + if (std.mem.eql(u8, args[1], "hca")) return hca(cfg, init.io); if (std.mem.eql(u8, args[1], "failfast") and (!cfg.failfast or args.len > 5) and !hostOnlyAllowed(cfg)) { std.debug.print("REFUSED: failfast with TF_TP_FAILFAST=0 or a hold is host-only: TF_TP_ALLOW_HOST_ONLY=1 and no RDMA wire\n", .{}); return 2; @@ -50,6 +51,41 @@ pub fn main(init: std.process.Init) !u8 { return 2; } +/// `hca` (host only, no GPU): both ranks' RoCE pick over the bootstrap; this rank's `NCCL_IB_HCA=...` line on stdout. +fn hca(cfg: tp.Config, io: std.Io) !u8 { + var b = tp.Bootstrap.open(cfg) catch |e| { + const s = @as(f64, @floatFromInt(cfg.connect_timeout_ns)) / std.time.ns_per_s; + switch (e) { + // a peer that never runs `hca` (it names its ports, or is not up) leaves this rank waiting until here + error.Timeout => std.debug.print("FAIL hca rank {d}: no peer joined within TF_TP_PLAN_TIMEOUT ({d:.1} s): is its runner up, and leaving TF_TP_ROCE_HCA and NCCL_IB_HCA unset as this one does?\n", .{ cfg.rank, s }), + error.AddressInUse => std.debug.print("FAIL hca rank {d}: port {d} (TF_TP_PORT + 7) is in use on this host\n", .{ cfg.rank, cfg.link_port }), + error.Closed => std.debug.print("FAIL hca rank {d}: the peer closed the bootstrap connection\n", .{cfg.rank}), + else => std.debug.print("FAIL hca rank {d}: bootstrap: {t}\n", .{ cfg.rank, e }), + } + return 1; + }; + defer b.close(); + var specs: [tp.roce_verbs.max_hcas]tp.roce_verbs.HcaSpec = undefined; + const pick = tp.roce_pick.select(cfg, &b, &specs) catch |e| { + std.debug.print("FAIL hca rank {d}: {t} while picking the ports with the peer\n", .{ cfg.rank, e }); + return 1; + }; + var buf: [tp.roce_pick.list_bytes + 16]u8 = undefined; + const line = tp.roce_pick.hcaLine(pick, &specs, &buf) catch |e| { + const why = switch (e) { + error.NoRoceDevice => "no active Ethernet RDMA port with an IPv4 RoCE v2 GID on this host", + error.NoSharedSubnet => "no RoCE port here shares an IPv4 /24 with the peer's", + }; + std.debug.print("FAIL hca rank {d}: {s}: name the peer's ports in TF_TP_ROCE_HCA or NCCL_IB_HCA\n", .{ cfg.rank, why }); + return 1; + }; + var names: [tp.roce_pick.list_bytes]u8 = undefined; + const how = if (pick.explicit) "TF_TP_ROCE_HCA" else "on a /24 the peer shares"; + std.debug.print("INFO hca rank {d}: {s} ({s})\n", .{ cfg.rank, try tp.roce_pick.joinNames(specs[0..pick.n], true, &names), how }); + try std.Io.File.stdout().writeStreamingAll(io, line); + return 0; +} + fn suite(s: *tp.Session, d: *const cuda.Driver, gpa: std.mem.Allocator, cfg: tp.Config) !u8 { var c = try checks.Ctx.init(d, gpa, cfg.rank, max_bytes); defer c.deinit(); diff --git a/zig/src/fabric/tp/roce_pick.zig b/zig/src/fabric/tp/roce_pick.zig index 627d0aefe1..b9355c401a 100644 --- a/zig/src/fabric/tp/roce_pick.zig +++ b/zig/src/fabric/tp/roce_pick.zig @@ -3,6 +3,7 @@ const std = @import("std"); const log = @import("log.zig"); const rv = @import("roce_verbs.zig"); const bootstrap = @import("bootstrap.zig"); +const sock = @import("sock.zig"); const Config = @import("config.zig").Config; const Verbs = @import("verbs").Verbs; @@ -10,7 +11,8 @@ const Verbs = @import("verbs").Verbs; pub const max_ports = 8; const blob_bytes = 1 + 4 * max_ports; -pub const Pick = struct { n: usize, shared: bool }; +/// How many devices, and how: on a /24 the peer shares (`shared`), named by TF_TP_ROCE_HCA (`explicit`), or neither. +pub const Pick = struct { n: usize, shared: bool, explicit: bool = false }; fn usable(p: rv.Port) bool { return p.active and p.ethernet and p.ipv4 != null; @@ -83,6 +85,11 @@ fn decode(blob: []const u8, out: [][4]u8) usize { /// The devices for the verbs wire. A collective: every rank calls it, including ranks that set TF_TP_ROCE_HCA. pub fn resolve(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSpec) !usize { + return (try select(cfg, b, out)).n; +} + +/// `resolve` with how the devices were chosen, over this host's ports (a collective, as `resolve`). +pub fn select(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSpec) !Pick { var ports: [max_ports]rv.Port = undefined; var n_ports: usize = 0; if (Verbs.open(null)) |opened| { @@ -90,19 +97,55 @@ pub fn resolve(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSp defer v.close(); n_ports = rv.scanPorts(&v, &ports); } else |_| {} + return selectFrom(cfg, b, ports[0..n_ports], out); +} + +/// `select` over given ports (the scan's, or a test's fixture ports). +pub fn selectFrom(cfg: Config, b: *bootstrap.Bootstrap, ports: []const rv.Port, out: *[rv.max_hcas]rv.HcaSpec) !Pick { var mine: [blob_bytes]u8 = undefined; var buf: [bootstrap.max_ranks * blob_bytes]u8 = undefined; var blobs: [bootstrap.max_ranks][]const u8 = undefined; - const all = try b.allGather(encode(ports[0..n_ports], &mine), &buf, &blobs); - if (cfg.roce_hca) |text| return rv.parseSpecs(text, out); + const all = try b.allGather(encode(ports, &mine), &buf, &blobs); + if (cfg.roce_hca) |text| return .{ .n = try rv.parseSpecs(text, out), .shared = false, .explicit = true }; var peer: [bootstrap.max_ranks * max_ports][4]u8 = undefined; var n_peer: usize = 0; for (all, 0..) |blob, r| if (r != b.rank) { n_peer += decode(blob, peer[n_peer..]); }; - const pick = choose(ports[0..n_ports], peer[0..n_peer], out, cfg.roce_hcas); + const pick = choose(ports, peer[0..n_peer], out, cfg.roce_hcas); if (!pick.shared and pick.n > 0) log.warn("tp roce: no RoCE port shares an IPv4 /24 with the peer's; using the first {d} active (set TF_TP_ROCE_HCA)", .{pick.n}); - return pick.n; + return pick; +} + +/// Device names joined by commas, with each `:gid` when `gid` (TF_TP_ROCE_HCA's form) or without (NCCL_IB_HCA's). +pub fn joinNames(specs: []const rv.HcaSpec, gid: bool, buf: []u8) error{WriteFailed}![]const u8 { + var w: std.Io.Writer = .fixed(buf); + for (specs, 0..) |s, i| { + try w.print("{s}{s}", .{ if (i == 0) "" else ",", s.name[0..s.name_len] }); + if (gid) if (s.gid) |g| try w.print(":{d}", .{g}); + } + return w.buffered(); +} + +/// Room for `joinNames` of every device a rank stripes over (a name is under 64 bytes, a GID index at most 3 digits). +pub const list_bytes = rv.max_hcas * (64 + 4); + +pub const HcaError = error{ + /// no active Ethernet port with an IPv4 RoCE v2 GID on this host + NoRoceDevice, + /// this host's ports share no IPv4 /24 with the peer's: NCCL would not reach the peer on them either + NoSharedSubnet, +}; + +/// `tf-tp-test hca`'s line for run_hosts.sh: `NCCL_IB_HCA=` the picked devices; never the first local ports alone. +pub fn hcaLine(pick: Pick, specs: []const rv.HcaSpec, buf: *[list_bytes + 16]u8) HcaError![]const u8 { + if (pick.n == 0) return error.NoRoceDevice; + if (!pick.shared and !pick.explicit) return error.NoSharedSubnet; + const prefix = "NCCL_IB_HCA="; + @memcpy(buf[0..prefix.len], prefix); + const joined = joinNames(specs[0..pick.n], false, buf[prefix.len .. buf.len - 1]) catch unreachable; // list_bytes holds them + buf[prefix.len + joined.len] = '\n'; + return buf[0 .. prefix.len + joined.len + 1]; } fn port(name: []const u8, active: bool, ipv4: ?[4]u8) rv.Port { @@ -114,9 +157,7 @@ fn port(name: []const u8, active: bool, ipv4: ?[4]u8) rv.Port { } fn names(specs: []const rv.HcaSpec, buf: []u8) []const u8 { - var w: std.Io.Writer = .fixed(buf); - for (specs, 0..) |s, i| w.print("{s}{s}", .{ if (i == 0) "" else ",", s.name[0..s.name_len] }) catch unreachable; - return w.buffered(); + return joinNames(specs, false, buf) catch unreachable; } // A host with two neighbours: port f0 cabled to one Spark (both functions on one /24), port f1 to the other. @@ -184,3 +225,162 @@ test "two rails on one /24 count once per peer port; no shared /24 keeps the lis try std.testing.expectEqual(@as(usize, 0), decode(&.{ 2, 1, 2, 3, 4 }, &out)); try std.testing.expectEqual(@as(?u8, 3), specs[0].gid); } + +test "hca's line: the picked names for NCCL; no device or no shared /24 is refused, an explicit list is taken" { + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var line: [list_bytes + 16]u8 = undefined; + var peer_buf: [max_ports][4]u8 = undefined; + const shared = choose(&a_ports, listed(&b_ports, &peer_buf), &specs, 2); + try std.testing.expectEqualStrings("NCCL_IB_HCA=rocep1s0f1,roceP2p1s0f1\n", try hcaLine(shared, &specs, &line)); + var buf: [list_bytes]u8 = undefined; + try std.testing.expectEqualStrings("rocep1s0f1:3,roceP2p1s0f1:3", try joinNames(specs[0..shared.n], true, &buf)); + // the old fallback (first active ports, no /24 in common) is what NCCL failed on: refused, never printed + try std.testing.expectError(error.NoSharedSubnet, hcaLine(choose(&a_ports, &.{.{ 10, 9, 9, 9 }}, &specs, 2), &specs, &line)); + try std.testing.expectError(error.NoRoceDevice, hcaLine(choose(&.{}, &.{}, &specs, 2), &specs, &line)); + const n = try rv.parseSpecs("rocep1s0f1:3, roceP2p1s0f1", &specs); + try std.testing.expectEqualStrings("NCCL_IB_HCA=rocep1s0f1,roceP2p1s0f1\n", try hcaLine(.{ .n = n, .shared = false, .explicit = true }, &specs, &line)); +} + +/// One RDMA device for `rv.scanPorts` through a stub verbs table: port 1 up or down, Ethernet, mlx5's GIDs 0-3. +const FakeDev = struct { name: [:0]const u8, up: bool, ipv4: [4]u8 = @splat(0) }; + +const ScanFixture = struct { + const vabi = @import("verbs").abi; + var host: []const FakeDev = &.{}; + var list: [max_ports:null]?*vabi.Device = undefined; + + fn handle(i: usize) usize { + return 0x10000 * (i + 1); + } + fn index(p: *const anyopaque) usize { + return @intFromPtr(p) / 0x10000 - 1; + } + fn getList(num: *c_int) callconv(.c) ?[*:null]?*vabi.Device { + for (host, 0..) |_, i| list[i] = @ptrFromInt(handle(i)); + list[host.len] = null; + num.* = @intCast(host.len); + return &list; + } + fn freeList(_: [*:null]?*vabi.Device) callconv(.c) void {} + fn devName(d: *vabi.Device) callconv(.c) ?[*:0]const u8 { + return host[index(d)].name.ptr; + } + fn openDev(d: *vabi.Device) callconv(.c) ?*vabi.Context { + return @ptrFromInt(@intFromPtr(d)); + } + fn closeDev(_: *vabi.Context) callconv(.c) c_int { + return 0; + } + fn queryPort(c: *vabi.Context, _: u8, a: *vabi.PortAttr) callconv(.c) c_int { + a.* = std.mem.zeroes(vabi.PortAttr); + a.state = if (host[index(c)].up) vabi.port_active else 1; + a.link_layer = 2; // Ethernet + a.gid_tbl_len = 4; + return 0; + } + fn queryGid(c: *vabi.Context, _: u8, i: c_int, g: *vabi.Gid) callconv(.c) c_int { + g.* = .{ .raw = @splat(0) }; + if (i < 2) { + g.raw[0] = 0xfe; + g.raw[1] = 0x80; + g.raw[15] = @intCast(index(c) + 1); + } else { + g.raw[10] = 0xff; + g.raw[11] = 0xff; + g.raw[12..16].* = host[index(c)].ipv4; + } + return 0; + } + fn gidEx(_: *vabi.Context, _: u32, i: u32, e: *vabi.GidEntry, _: u32, _: usize) callconv(.c) c_int { + e.* = std.mem.zeroes(vabi.GidEntry); + e.gid_type = if (i % 2 == 1) 2 else 1; // odd indices RoCE v2 + return 0; + } + + /// `rv.scanPorts` over `devs` (names that no real /sys/class/infiniband has, so sysfs never answers for them). + fn scan(devs: []const FakeDev, out: *[max_ports]rv.Port) []const rv.Port { + host = devs; + var v: Verbs = undefined; + v.api.ibv_get_device_list = getList; + v.api.ibv_free_device_list = freeList; + v.api.ibv_get_device_name = devName; + v.api.ibv_open_device = openDev; + v.api.ibv_close_device = closeDev; + v.api.ibv_query_port = queryPort; + v.api.ibv_query_gid = queryGid; + v.gid_ex = gidEx; + return out[0..rv.scanPorts(&v, out)]; + } +}; + +/// #636's reported two GB10 hosts, modelled (no capture): device order and picks reported; addresses, states assumed. +const hosts = [2][4]FakeDev{ .{ + .{ .name = "fx_rocep1s0f0", .up = true, .ipv4 = .{ 192, 0, 2, 1 } }, + .{ .name = "fx_rocep1s0f1", .up = false }, + .{ .name = "fx_roceP2p1s0f0", .up = true, .ipv4 = .{ 198, 51, 100, 1 } }, + .{ .name = "fx_roceP2p1s0f1", .up = false }, +}, .{ + .{ .name = "fx_rocep1s0f0", .up = false }, + .{ .name = "fx_rocep1s0f1", .up = true, .ipv4 = .{ 192, 0, 2, 2 } }, + .{ .name = "fx_roceP2p1s0f0", .up = false }, + .{ .name = "fx_roceP2p1s0f1", .up = true, .ipv4 = .{ 198, 51, 100, 2 } }, +} }; + +test "the port scan reads a host's devices through verbs: state, Ethernet, the IPv4 RoCE v2 GID index" { + var out: [max_ports]rv.Port = undefined; + const a = ScanFixture.scan(&hosts[0], &out); + try std.testing.expectEqual(@as(usize, 4), a.len); + try std.testing.expectEqualStrings("fx_rocep1s0f0", a[0].spec.name[0..a[0].spec.name_len]); + try std.testing.expect(a[0].active and a[0].ethernet and !a[1].active); + try std.testing.expectEqual(@as(?[4]u8, .{ 192, 0, 2, 1 }), a[0].ipv4); + try std.testing.expectEqual(@as(?u8, 3), a[0].spec.gid); // the v2 IPv4 GID, not v1's index 2 + try std.testing.expectEqual(@as(?[4]u8, null), a[1].ipv4); +} + +/// One rank's `selectFrom` over a real bootstrap (rank 0 on `l`), its hca line or error into `out`. +fn hcaRank(cfg: Config, l: ?sock.Fd, ports: []const rv.Port, out: *[list_bytes + 16]u8, line: *HcaError![]const u8, err: *?anyerror) void { + var b = (if (l) |fd| bootstrap.Bootstrap.acceptOn(fd, 2, 5 * std.time.ns_per_s) else bootstrap.Bootstrap.open(cfg)) catch |e| { + err.* = e; + return; + }; + defer b.close(); + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + const pick = selectFrom(cfg, &b, ports, &specs) catch |e| { + err.* = e; + return; + }; + line.* = hcaLine(pick, &specs, out); +} + +fn hcaPair(cfgs: [2]Config, ports: [2][]const rv.Port, lines: *[2]HcaError![]const u8, bufs: *[2][list_bytes + 16]u8) !void { + const l = try bootstrap.testListener(); + defer sock.close(l.fd); + var errs: [2]?anyerror = @splat(null); + var c1 = cfgs[1]; + c1.link_port = l.port; + c1.connect_timeout_ns = 5 * std.time.ns_per_s; + const t = try std.Thread.spawn(.{}, hcaRank, .{ c1, null, ports[1], &bufs[1], &lines[1], &errs[1] }); + hcaRank(cfgs[0], l.fd, ports[0], &bufs[0], &lines[0], &errs[0]); + t.join(); + for (errs) |e| if (e) |x| return x; +} + +test "hca over a real bootstrap: both ranks print their peer-linked ports; an explicit rank keeps its list; no ports fail both" { + var lines: [2]HcaError![]const u8 = undefined; + var bufs: [2][list_bytes + 16]u8 = undefined; + var scanned: [2][max_ports]rv.Port = undefined; + const ports = [2][]const rv.Port{ ScanFixture.scan(&hosts[0], &scanned[0]), ScanFixture.scan(&hosts[1], &scanned[1]) }; + const auto: [2]Config = .{ .{ .rank = 0 }, .{ .rank = 1 } }; + try hcaPair(auto, ports, &lines, &bufs); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f0,fx_roceP2p1s0f0\n", try lines[0]); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f1,fx_roceP2p1s0f1\n", try lines[1]); + // TF_TP_ROCE_HCA on one rank: its list as given (no GID for NCCL); it still sends its ports to the other + const one: [2]Config = .{ .{ .rank = 0, .roce_hca = "fx_rocep1s0f0:3" }, .{ .rank = 1 } }; + try hcaPair(one, ports, &lines, &bufs); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f0\n", try lines[0]); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f1,fx_roceP2p1s0f1\n", try lines[1]); + // a host without usable ports: it has none, its peer shares no /24 with it, and both refuse + try hcaPair(auto, .{ &.{}, ports[1] }, &lines, &bufs); + try std.testing.expectError(error.NoRoceDevice, lines[0]); + try std.testing.expectError(error.NoSharedSubnet, lines[1]); +} diff --git a/zig/src/fabric/tp/run_hosts.sh b/zig/src/fabric/tp/run_hosts.sh index 5f00168f83..75f0c7900b 100644 --- a/zig/src/fabric/tp/run_hosts.sh +++ b/zig/src/fabric/tp/run_hosts.sh @@ -7,13 +7,31 @@ mkdir -p "$OUT" # NCCL uses the RoCE wire's devices unless NCCL_IB_HCA is set: the names without `:gid` (NCCL reads `:N` as a port) [ -n "${TF_TP_ROCE_HCA:-}" ] && [ -z "${NCCL_IB_HCA:-}" ] && NCCL_IB_HCA=$(echo "$TF_TP_ROCE_HCA" | sed 's/:[0-9]*//g; s/ //g') [ -n "${TF_TP_ROCE_HCA:-}" ] && export TF_TP_ROCE_HCA -[ -n "${NCCL_IB_HCA:-}" ] && export NCCL_IB_HCA "$BIN" kernel > /dev/null 2>&1 || { echo "error: $BIN has no mailbox kernel: rebuild with -Dtp-nvcc=PATH or -Dtp-fatbin=PATH" >&2; exit 2; } ulimit -l unlimited 2> /dev/null || echo "warning: memlock is $(ulimit -l) KiB: ibv_reg_mr of the RoCE region may fail" >&2 -DEVICES=$(ibv_devices 2> /dev/null | tail -n +3 | grep -c .) -if [ "$DEVICES" -gt 1 ] && [ -z "${TF_TP_ROCE_HCA:-}${NCCL_IB_HCA:-}" ]; then - echo "warning: $DEVICES RDMA devices and no TF_TP_ROCE_HCA / NCCL_IB_HCA: name the peer's ports (README.md)" >&2 +# neither named: NCCL takes the ports the RoCE picker chose with the peer (`tf-tp-test hca`, both ranks), not its own +if [ -z "${TF_TP_ROCE_HCA:-}${NCCL_IB_HCA:-}" ]; then + # hca gives up after TF_TP_PLAN_TIMEOUT (default 300 s) with its reason; the cap here only catches a hang past that + CAP=$(awk 'BEGIN { t = ARGV[1] + 0; if (t <= 0) t = 300; printf "%d", int(t) + 31 }' "${TF_TP_PLAN_TIMEOUT:-300}") + echo "hca: rank $RANK picking the ports with the peer (up to $CAP s)" + PICK=$(TF_TP_RANK=$RANK TF_TP_MASTER=$MASTER TF_TP_PORT=29700 timeout "$CAP" "$BIN" hca 2> "$OUT/hca.log") + rc=$? + if [ "$rc" -eq 124 ]; then + echo "error: $BIN hca still running after $CAP s: killed ($OUT/hca.log)" >&2 + exit 3 + fi + if [ "$rc" -ne 0 ]; then + echo "error: $BIN hca exited $rc: $(grep -h '^FAIL\|^error' "$OUT/hca.log" | tail -n 1 | grep . || tail -n 1 "$OUT/hca.log") ($OUT/hca.log)" >&2 + exit 3 + fi + if ! echo "$PICK" | grep -Eqx 'NCCL_IB_HCA=[A-Za-z0-9_.-]+(,[A-Za-z0-9_.-]+)*'; then + echo "error: $BIN hca printed '$PICK', not NCCL_IB_HCA=name,... ($OUT/hca.log)" >&2 + exit 3 + fi + NCCL_IB_HCA=${PICK#NCCL_IB_HCA=} + echo "hca: rank $RANK NCCL_IB_HCA=$NCCL_IB_HCA ($(grep -h '^INFO hca' "$OUT/hca.log"))" fi +[ -n "${NCCL_IB_HCA:-}" ] && export NCCL_IB_HCA { hostname; nvidia-smi -L; ibv_devices 2>&1; ulimit -l; env | grep -E '^(TF_TP_ROCE_HCA|NCCL_IB_HCA|NCCL_SOCKET_IFNAME)='; } > "$OUT/host.txt" 2>&1 PORT=29700 step() { # NAME BACKEND ARGS...: this host's rank; log OUT/NAME.log ending with an EXIT line From 5410eac2a92cba50e63fe0ffef370d3379ad7fd4 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:42:56 +0700 Subject: [PATCH 23/43] Session memory for paged-KV families: page pool, prefix trie, RAM residency budget, streamed NVMe tier. zig/src/sessions keeps a chat's KV pages after its request ends: entries are found by a trie over chained page digests, held in RAM within a residency budget (chats and pool bytes), and the least recently used idle chat is parked to an NVMe tier and streamed back on resume. Ranks that split KV run every store call in the same order, so their books stay identical. It complements core/prompt_cache, which snapshots a whole family state; a family that uses it reports the resident bytes as Memory.cache and its budget as Info.prompt_cache_plan. No family uses it yet; zig build test-sessions runs the host tests, the NVMe tier on real files. --- build.zig | 2 + zig/build/sessions.zig | 13 + zig/src/sessions/disk.zig | 587 ++++++++++++++++++++++++++++++++ zig/src/sessions/disk_test.zig | 125 +++++++ zig/src/sessions/format.zig | 197 +++++++++++ zig/src/sessions/freeset.zig | 114 +++++++ zig/src/sessions/jobs.zig | 150 ++++++++ zig/src/sessions/lanes.zig | 336 ++++++++++++++++++ zig/src/sessions/pagestore.zig | 129 +++++++ zig/src/sessions/pool.zig | 475 ++++++++++++++++++++++++++ zig/src/sessions/prefix.zig | 457 +++++++++++++++++++++++++ zig/src/sessions/sessions.zig | 33 ++ zig/src/sessions/store.zig | 570 +++++++++++++++++++++++++++++++ zig/src/sessions/store_test.zig | 216 ++++++++++++ 14 files changed, 3404 insertions(+) create mode 100644 zig/build/sessions.zig create mode 100644 zig/src/sessions/disk.zig create mode 100644 zig/src/sessions/disk_test.zig create mode 100644 zig/src/sessions/format.zig create mode 100644 zig/src/sessions/freeset.zig create mode 100644 zig/src/sessions/jobs.zig create mode 100644 zig/src/sessions/lanes.zig create mode 100644 zig/src/sessions/pagestore.zig create mode 100644 zig/src/sessions/pool.zig create mode 100644 zig/src/sessions/prefix.zig create mode 100644 zig/src/sessions/sessions.zig create mode 100644 zig/src/sessions/store.zig create mode 100644 zig/src/sessions/store_test.zig diff --git a/build.zig b/build.zig index d9ec51a4e0..5d2b15b96d 100644 --- a/build.zig +++ b/build.zig @@ -13,6 +13,7 @@ const exl3_build = @import("zig/build/exl3.zig"); const exlr8_build = @import("zig/build/exlr8.zig"); const golden_build = @import("zig/build/golden.zig"); const tp_build = @import("zig/build/tp.zig"); +const sessions_build = @import("zig/build/sessions.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -191,6 +192,7 @@ pub fn build(b: *std.Build) void { b.step("test-learn-warning", "CPU production-log refusal, backoff and recovery test").dependOn(&warning_run.step); grammar_build.steps(b, target, xgrammar); tp_build.add(b, target, optimize, test_step); + sessions_build.add(b, test_step); } /// `zig build native -Dcpu=apple_m1`: tensorfold-native with the Metal engines for the Python package's bundle (a native M5 build traps on M1-M4). diff --git a/zig/build/sessions.zig b/zig/build/sessions.zig new file mode 100644 index 0000000000..a5d8ea6f3c --- /dev/null +++ b/zig/build/sessions.zig @@ -0,0 +1,13 @@ +//! Session memory's build: the family-neutral `sessions` module (zig/src/sessions) and its host tests on any OS. + +const std = @import("std"); + +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path("zig/src/sessions/sessions.zig"), .target = target, .optimize = optimize, .link_libc = true }); +} + +pub fn add(b: *std.Build, test_step: *std.Build.Step) void { + const run = b.addRunArtifact(b.addTest(.{ .root_module = module(b, b.graph.host, .debug) })); + b.step("test-sessions", "Session memory's host tests: pool, prefix trie, RAM tier, NVMe tier on real files (no GPU)").dependOn(&run.step); + test_step.dependOn(&run.step); +} diff --git a/zig/src/sessions/disk.zig b/zig/src/sessions/disk.zig new file mode 100644 index 0000000000..f5ab58f06f --- /dev/null +++ b/zig/src/sessions/disk.zig @@ -0,0 +1,587 @@ +//! The NVMe tier: a file an entry under //rank/, written via .tmp + rename, streamed via the lanes. + +const std = @import("std"); +const linux = std.os.linux; +const Io = std.Io; +const lanes_mod = @import("lanes.zig"); +const Lanes = lanes_mod.Lanes; +const Staging = lanes_mod.Staging; +const Op = lanes_mod.Op; +const up = lanes_mod.up; +const format = @import("format.zig"); +const Layout = format.Layout; +const Head = format.Head; +const Seg = format.Seg; +const PageStore = @import("pagestore.zig").PageStore; +const pool_mod = @import("pool.zig"); +const Family = pool_mod.Family; +const Digest = @import("prefix.zig").Digest; + +pub const none = std.math.maxInt(u32); + +pub const Options = struct { + root: []const u8, + compat: Digest, + rank: u32 = 0, + world: u32 = 1, + /// bytes of files the tier keeps (least recently used files without dependents go first) + budget: u64 = 64 << 30, + /// entries shorter than this are dropped, not parked + min_tokens: u64 = 1024, + /// O_DIRECT where the filesystem takes it (buffered with the page cache dropped otherwise) + direct: bool = true, + /// staging buffers and their size: the host transient of any park or restore + staging: u32 = 16, + chunk: u32 = 8 << 20, + lanes: u32 = 4, + /// delta files a chain may stack before a park writes the entry whole + max_chain: u32 = 8, +}; + +/// A piece of bounded state (rings, carries, taps): small, host bytes. +pub const Blob = struct { + name: []const u8, + bytes: []const u8, +}; + +/// What a park writes: the entry, its tokens from first_page on, and this rank's pages per family, in logical order. +pub const WriteSrc = struct { + key: Digest, + /// the file holding pages [0, first_page) (a delta), or null (a whole entry: first_page 0) + base: ?Digest = null, + tag: u32, + kind: u32 = 0, + page: u32, + pos: u64, + first_page: u32 = 0, + npages: u32, + /// tokens [first_page x page, pos) + ids: []const i32, + store: PageStore, + fence: u64 = 0, + families: []const Family, + /// per family: family-tensor pages of this rank's logical pages [first_page, npages), in logical order + pages: []const []const u32, + blobs: []const Blob = &.{}, +}; + +/// One file of a restore: read logical pages [lo, hi) of it into the slot. +pub const Extent = struct { + key: Digest, + first_page: u32, + lo: u32, + hi: u32, +}; + +/// Where a restore puts rows: per family, the family-tensor page of each of the slot's logical pages. +pub const ReadSink = struct { + store: PageStore, + families: []const Family, + page: u32, + world: u32, + rank: u32, + /// per family: family-tensor pages for logical pages [0, npages) + pages: []const []const u32, +}; + +/// The bounded state a restore returns (owned). +pub const Restored = struct { + blobs: []Blob, + + pub fn deinit(r: *Restored, gpa: std.mem.Allocator) void { + for (r.blobs) |b| { + gpa.free(b.name); + gpa.free(b.bytes); + } + gpa.free(r.blobs); + } +}; + +/// The index record of one file. +pub const Rec = struct { + key: Digest, + base: u32 = none, + tag: u32, + kind: u32, + pos: u64, + first_page: u32, + npages: u32, + size: u64, + /// files stacked on this one + refs: u32 = 0, + /// files under it, this one included + depth: u32 = 1, + older: u32 = none, + newer: u32 = none, + live: bool = true, +}; + +pub const Stats = struct { + writes: u64 = 0, + reads: u64 = 0, + bytes_written: u64 = 0, + bytes_read: u64 = 0, + bad: u64 = 0, + evicted: u64 = 0, +}; + +pub const Disk = struct { + gpa: std.mem.Allocator, + io: Io, + opts: Options, + dir: Io.Dir, + staging: Staging, + lanes: *Lanes, + recs: std.ArrayList(Rec) = .empty, + free_recs: std.ArrayList(u32) = .empty, + by_key: std.AutoHashMapUnmanaged(Digest, u32) = .empty, + /// least recently used first + oldest: u32 = none, + newest: u32 = none, + used: u64 = 0, + stats: Stats = .{}, + + pub fn open(gpa: std.mem.Allocator, io: Io, opts: Options) !*Disk { + const d = try gpa.create(Disk); + errdefer gpa.destroy(d); + var name_buf: [64]u8 = undefined; + const sub = try std.fmt.bufPrint(&name_buf, "{x}/rank{d}", .{ opts.compat[0..8], opts.rank }); + var root = try Io.Dir.cwd().createDirPathOpen(io, opts.root, .{}); + defer root.close(io); + const dir = try root.createDirPathOpen(io, sub, .{ .open_options = .{ .iterate = true } }); + errdefer dir.close(io); + d.* = .{ .gpa = gpa, .io = io, .opts = opts, .dir = dir, .staging = try Staging.init(gpa, io, @max(2, opts.staging), up(opts.chunk)), .lanes = undefined }; + errdefer d.staging.deinit(); + d.lanes = try Lanes.start(gpa, io, &d.staging, @max(1, opts.lanes)); + return d; + } + + pub fn close(d: *Disk) void { + d.lanes.shutdown(); + d.staging.deinit(); + d.dir.close(d.io); + d.recs.deinit(d.gpa); + d.free_recs.deinit(d.gpa); + d.by_key.deinit(d.gpa); + const gpa = d.gpa; + gpa.destroy(d); + } + + pub fn accepts(d: *const Disk, pos: u64) bool { + return pos >= d.opts.min_tokens; + } + + pub fn get(d: *Disk, key: Digest) ?*Rec { + const i = d.by_key.get(key) orelse return null; + return &d.recs.items[i]; + } + + fn fileName(key: Digest, ext: []const u8, buf: *[48]u8) [:0]const u8 { + const n = (std.fmt.bufPrint(buf[0..47], "{x}.{s}", .{ &key, ext }) catch unreachable).len; + buf[n] = 0; + return buf[0..n :0]; + } + + fn openFile(d: *Disk, name: [:0]const u8, write: bool) !struct { fd: linux.fd_t, direct: bool } { + var flags: linux.O = .{ .ACCMODE = if (write) .WRONLY else .RDONLY, .CLOEXEC = true }; + if (write) { + flags.CREAT = true; + flags.TRUNC = true; + } + if (d.opts.direct) { + var f2 = flags; + f2.DIRECT = true; + const rc = linux.openat(d.dir.handle, name, f2, 0o644); + if (linux.errno(rc) == .SUCCESS) return .{ .fd = @intCast(rc), .direct = true }; + } + const rc = linux.openat(d.dir.handle, name, flags, 0o644); + if (linux.errno(rc) != .SUCCESS) return if (linux.errno(rc) == .NOENT) error.FileNotFound else error.OpenFailed; + return .{ .fd = @intCast(rc), .direct = false }; + } + + /// The layout a WriteSrc gets: segments, chunk sizes (whole pages a chunk for pool segments). + pub fn plan(d: *Disk, src: *const WriteSrc) !Layout { + var segs: std.ArrayList(Seg) = .empty; + defer segs.deinit(d.gpa); + const cap = d.opts.chunk; + try segs.append(d.gpa, .{ .name = format.segName("ids"), .kind = .ids, .family = 0, .nbytes = src.ids.len * 4, .chunk = cap }); + for (src.families, 0..) |f, i| { + const pb: u32 = @intCast(f.pageBytes(src.page)); + const k = @max(1, cap / pb); + if (@as(u64, k) * pb > d.staging.size) return error.PageTooLarge; + try segs.append(d.gpa, .{ .name = format.segName(f.name), .kind = .pool, .family = @intCast(i), .nbytes = @as(u64, src.pages[i].len) * pb, .chunk = k * pb }); + } + for (src.blobs) |b| try segs.append(d.gpa, .{ .name = format.segName(b.name), .kind = .blob, .family = 0, .nbytes = b.bytes.len, .chunk = cap }); + const head: Head = .{ .key = src.key, .base = src.base orelse @splat(0), .compat = d.opts.compat, .tag = src.tag, .kind = src.kind, .page = src.page, .world = d.opts.world, .rank = d.opts.rank, .first_page = src.first_page, .npages = src.npages, .nseg = 0, .ndigest = 0, .pos = src.pos }; + return Layout.build(d.gpa, head, segs.items); + } + + /// Writes an entry's file streamed (the pump copies pages to staging, lanes hash and write); returns its bytes. + pub fn writeFile(d: *Disk, src: *const WriteSrc) !u64 { + var lay = try d.plan(src); + defer lay.deinit(d.gpa); + var tb: [48]u8 = undefined; + var fb: [48]u8 = undefined; + const tmp = fileName(src.key, "tmp", &tb); + const f = try d.openFile(tmp, true); + var ok = false; + defer { + _ = linux.close(f.fd); + if (!ok) _ = linux.unlinkat(d.dir.handle, tmp, 0); + } + try src.store.waitFence(src.fence); + var op = try Op.init(d.gpa, d.io, &d.staging); + defer op.deinit(); + errdefer op.wait() catch {}; + for (lay.segs) |*s| { + for (0..s.nchunks) |ci| { + if (op.failed()) |e| return e; + const c = s.chunkAt(@intCast(ci)); + const sb = try d.staging.get(); + errdefer d.staging.put(sb); + const buf = d.staging.buf(sb)[0..c.data]; + const at = @as(u64, ci) * s.chunk; + switch (s.kind) { + .ids => @memcpy(buf, std.mem.sliceAsBytes(src.ids)[at..][0..c.data]), + .blob => @memcpy(buf, blobOf(src, s).bytes[at..][0..c.data]), + .pool => { + const pb: u32 = @intCast(src.families[s.family].pageBytes(src.page)); + const p0: usize = @intCast(at / pb); + try src.store.read(s.family, src.pages[s.family][p0..][0 .. c.data / pb], buf); + }, + _ => unreachable, + } + d.lanes.submit(.{ .kind = .write, .fd = f.fd, .staging = sb, .data = c.data, .span = c.span, .off = c.off, .digest = &lay.digests[s.first_digest + ci], .op = &op }); + } + } + try op.wait(); + const hb = try d.staging.get(); + defer d.staging.put(hb); + if (lay.head_span > d.staging.size) return error.HeaderTooLarge; + lay.encode(d.staging.buf(hb)); + try lanes_mod.pwriteAll(f.fd, d.staging.buf(hb)[0..lay.head_span], 0); + if (linux.errno(linux.fdatasync(f.fd)) != .SUCCESS) return error.SyncFailed; + if (!f.direct) _ = linux.fadvise(f.fd, 0, 0, linux.POSIX_FADV.DONTNEED); + if (linux.errno(linux.renameat(d.dir.handle, tmp, d.dir.handle, fileName(src.key, "tfs", &fb))) != .SUCCESS) return error.RenameFailed; + ok = true; + return lay.head.total; + } + + fn blobOf(src: *const WriteSrc, s: *const Seg) Blob { + for (src.blobs) |b| if (std.mem.eql(u8, b.name, s.nameOf())) return b; + unreachable; + } + + /// A file's header, checked. + pub fn readHead(d: *Disk, fd: linux.fd_t) !Layout { + const hb = try d.staging.get(); + defer d.staging.put(hb); + const buf = d.staging.buf(hb); + try lanes_mod.preadAll(fd, buf[0..4096], 0); + const hn = std.mem.readInt(u32, buf[12..16], .little); + if (up(hn) > buf.len) return error.BadHeader; + if (hn > 4096) try lanes_mod.preadAll(fd, buf[4096..up(hn)], 4096); + return Layout.decode(d.gpa, buf[0..up(hn)]); + } + + /// Reads pages [lo, hi) of one restore file into the sink, each chunk checked first; blobs too when want_blobs. + pub fn readFile(d: *Disk, ext: Extent, sink: *const ReadSink, want_blobs: bool) !?Restored { + var nb: [48]u8 = undefined; + const f = try d.openFile(fileName(ext.key, "tfs", &nb), false); + defer _ = linux.close(f.fd); + var lay = try d.readHead(f.fd); + defer lay.deinit(d.gpa); + if (!std.mem.eql(u8, &lay.head.key, &ext.key) or lay.head.first_page != ext.first_page or lay.head.page != sink.page) return error.BadHeader; + var st: linux.Statx = undefined; + if (linux.errno(linux.statx(f.fd, "", linux.AT.EMPTY_PATH, .{ .SIZE = true }, &st)) != .SUCCESS or st.size < lay.head.total) return error.Truncated; + var out: std.ArrayList(Blob) = .empty; + errdefer { + for (out.items) |b| { + d.gpa.free(b.name); + d.gpa.free(b.bytes); + } + out.deinit(d.gpa); + } + var op = try Op.init(d.gpa, d.io, &d.staging); + defer op.deinit(); + errdefer op.wait() catch {}; + // (segment, chunk) jobs; the pump keeps at most staging - 1 in flight (one buffer for the header path) + var jobs: std.ArrayList(struct { seg: u32, chunk: u32 }) = .empty; + defer jobs.deinit(d.gpa); + for (lay.segs, 0..) |*s, si| switch (s.kind) { + .pool => { + if (s.family >= sink.families.len) return error.BadHeader; + const fam = sink.families[s.family]; + const pb: u32 = @intCast(fam.pageBytes(sink.page)); + if (s.chunk % pb != 0 or !std.mem.eql(u8, s.nameOf(), fam.name)) return error.BadHeader; + const need = ownedIn(fam, sink, ext.first_page, ext.hi) - ownedIn(fam, sink, ext.first_page, ext.lo); + const skip = ownedIn(fam, sink, ext.first_page, ext.lo); + if (s.nbytes < (@as(u64, skip) + need) * pb) return error.BadHeader; + const kpp = s.chunk / pb; + const c0 = skip / kpp; + const c1 = if (need == 0) c0 else (skip + need + kpp - 1) / kpp; + for (c0..c1) |c| try jobs.append(d.gpa, .{ .seg = @intCast(si), .chunk = @intCast(c) }); + }, + .blob => if (want_blobs) { + const buf = try d.gpa.alloc(u8, s.nbytes); + errdefer d.gpa.free(buf); + const name = try d.gpa.dupe(u8, s.nameOf()); + errdefer d.gpa.free(name); + try out.append(d.gpa, .{ .name = name, .bytes = buf }); + for (0..s.nchunks) |c| try jobs.append(d.gpa, .{ .seg = @intCast(si), .chunk = @intCast(c) }); + }, + else => {}, + }; + const window = d.staging.count - 1; + var sent: usize = 0; + var seen: usize = 0; + var bytes: u64 = 0; + while (seen < jobs.items.len) { + while (sent < jobs.items.len and sent - seen < window and op.failed() == null) : (sent += 1) { + const jb = jobs.items[sent]; + const s = &lay.segs[jb.seg]; + const c = s.chunkAt(jb.chunk); + d.lanes.submit(.{ .kind = .read, .fd = f.fd, .staging = try d.staging.get(), .data = c.data, .span = c.span, .off = c.off, .digest = &lay.digests[s.first_digest + jb.chunk], .op = &op, .tag = @intCast(sent) }); + bytes += c.span; + } + const j = op.next() orelse break; + seen += 1; + defer d.staging.put(j.staging); + if (op.failed() != null) continue; + const jb = jobs.items[j.tag]; + const s = &lay.segs[jb.seg]; + const data = d.staging.buf(j.staging)[0..j.data]; + if (s.kind == .blob) { + for (out.items) |b| if (std.mem.eql(u8, b.name, s.nameOf())) @memcpy(@constCast(b.bytes)[@as(u64, jb.chunk) * s.chunk ..][0..data.len], data); + continue; + } + try d.drain(s, jb.chunk, data, ext, sink); + } + try op.wait(); + if (!f.direct) _ = linux.fadvise(f.fd, 0, 0, linux.POSIX_FADV.DONTNEED); + d.stats.bytes_read += bytes; + if (!want_blobs) { + out.deinit(d.gpa); + return null; + } + return .{ .blobs = try out.toOwnedSlice(d.gpa) }; + } + + /// A checked pool chunk's pages that fall in [lo, hi) into the sink's pages. + fn drain(d: *Disk, s: *const Seg, chunk: u32, data: []const u8, ext: Extent, sink: *const ReadSink) !void { + _ = d; + const fam = sink.families[s.family]; + const pb: u32 = @intCast(fam.pageBytes(sink.page)); + const kpp = s.chunk / pb; + const skip = ownedIn(fam, sink, ext.first_page, ext.lo); + const end = ownedIn(fam, sink, ext.first_page, ext.hi); + const a = @max(chunk * kpp, skip); + const b = @min(chunk * kpp + @as(u32, @intCast(data.len / pb)), end); + if (a >= b) return; + var list: [512]u32 = undefined; + var n: u32 = 0; + var j = a; + while (j < b) : (j += 1) { + list[n] = sink.pages[s.family][logicalOf(fam, sink, ext.first_page, j)]; + n += 1; + if (n == list.len or j + 1 == b) { + const first = j + 1 - n; + try sink.store.write(s.family, list[0..n], data[(first - chunk * kpp) * pb ..][0 .. n * pb]); + n = 0; + } + } + } + + /// Records a written file in the index (round thread, at the plan's point); its base gains a dependent. + pub fn register(d: *Disk, src: *const WriteSrc, size: u64) !u32 { + const id = try d.add(src.key, src.base, src.tag, src.kind, src.pos, src.first_page, src.npages, size); + d.stats.writes += 1; + d.stats.bytes_written += size; + return id; + } + + fn add(d: *Disk, key: Digest, base: ?Digest, tag: u32, kind: u32, pos: u64, first_page: u32, npages: u32, size: u64) !u32 { + var r: Rec = .{ .key = key, .tag = tag, .kind = kind, .pos = pos, .first_page = first_page, .npages = npages, .size = size }; + if (base) |b| { + const bi = d.by_key.get(b) orelse return error.NoBase; + r.base = bi; + r.depth = d.recs.items[bi].depth + 1; + } + const id: u32 = if (d.free_recs.pop()) |i| i else blk: { + try d.recs.append(d.gpa, undefined); + break :blk @intCast(d.recs.items.len - 1); + }; + try d.by_key.put(d.gpa, key, id); + if (r.base != none) d.recs.items[r.base].refs += 1; + d.recs.items[id] = r; + d.link(id); + d.used += size; + return id; + } + + /// A valid file found on restart: its header, size and verified ids (tokens [first_page x page, pos), owned). + pub const Found = struct { + head: Head, + size: u64, + mtime: i64, + ids: []i32, + + fn older(_: void, a: Found, b: Found) bool { + return a.mtime < b.mtime or (a.mtime == b.mtime and a.head.first_page < b.head.first_page); + } + }; + + /// Indexes the directory on restart: drops .tmp and bad files, checks ids, registers bases before their deltas. + pub fn scan(d: *Disk, out: *std.ArrayList(Found)) !void { + var it = d.dir.iterate(); + var nb: [48]u8 = undefined; + while (try it.next(d.io)) |ent| { + if (ent.name.len >= nb.len) continue; + @memcpy(nb[0..ent.name.len], ent.name); + nb[ent.name.len] = 0; + const name = nb[0..ent.name.len :0]; + if (std.mem.endsWith(u8, name, ".tmp")) { + _ = linux.unlinkat(d.dir.handle, name, 0); + continue; + } + if (!std.mem.endsWith(u8, name, ".tfs")) continue; + const f = d.found(name) catch { + d.stats.bad += 1; + _ = linux.unlinkat(d.dir.handle, name, 0); + continue; + }; + try out.append(d.gpa, f); + } + std.mem.sort(Found, out.items, {}, Found.older); + var i: usize = 0; + while (i < out.items.len) { + const f = out.items[i]; + const base: ?Digest = if (f.head.hasBase()) f.head.base else null; + if (base != null and d.by_key.get(base.?) == null) { + d.stats.bad += 1; + var kb: [48]u8 = undefined; + _ = linux.unlinkat(d.dir.handle, fileName(f.head.key, "tfs", &kb), 0); + d.gpa.free(f.ids); + _ = out.orderedRemove(i); + continue; + } + _ = try d.add(f.head.key, base, f.head.tag, f.head.kind, f.head.pos, f.head.first_page, f.head.npages, f.size); + i += 1; + } + } + + fn found(d: *Disk, name: [:0]const u8) !Found { + const f = try d.openFile(name, false); + defer _ = linux.close(f.fd); + var lay = try d.readHead(f.fd); + defer lay.deinit(d.gpa); + const h = lay.head; + var kb: [48]u8 = undefined; + if (!std.mem.eql(u8, fileName(h.key, "tfs", &kb), name)) return error.BadHeader; + if (!std.mem.eql(u8, &h.compat, &d.opts.compat) or h.rank != d.opts.rank or h.world != d.opts.world) return error.Foreign; + var st: linux.Statx = undefined; + if (linux.errno(linux.statx(f.fd, "", linux.AT.EMPTY_PATH, .{ .SIZE = true, .MTIME = true }, &st)) != .SUCCESS or st.size < h.total) return error.Truncated; + const seg = lay.find(.ids, "ids") orelse return error.BadHeader; + if (seg.nbytes % 4 != 0 or seg.nbytes / 4 != h.pos - @as(u64, h.first_page) * h.page) return error.BadHeader; + const ids = try d.gpa.alloc(i32, seg.nbytes / 4); + errdefer d.gpa.free(ids); + const sb = try d.staging.get(); + defer d.staging.put(sb); + for (0..seg.nchunks) |ci| { + const c = seg.chunkAt(@intCast(ci)); + const buf = d.staging.buf(sb); + try lanes_mod.preadAll(f.fd, buf[0..c.span], c.off); + var dg: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(buf[0..c.data], &dg, .{}); + if (!std.mem.eql(u8, &dg, &lay.digests[seg.first_digest + ci])) return error.ChecksumMismatch; + @memcpy(std.mem.sliceAsBytes(ids)[@as(usize, ci) * seg.chunk ..][0..c.data], buf[0..c.data]); + } + return .{ .head = h, .size = st.size, .mtime = st.mtime.sec * std.time.ns_per_s + st.mtime.nsec, .ids = ids }; + } + + fn link(d: *Disk, id: u32) void { + const r = &d.recs.items[id]; + r.older = d.newest; + r.newer = none; + if (d.newest != none) d.recs.items[d.newest].newer = id else d.oldest = id; + d.newest = id; + } + + fn unlink(d: *Disk, id: u32) void { + const r = &d.recs.items[id]; + if (r.older != none) d.recs.items[r.older].newer = r.newer else d.oldest = r.newer; + if (r.newer != none) d.recs.items[r.newer].older = r.older else d.newest = r.older; + } + + /// Marks a file (and its bases) used now: O(chain). + pub fn touch(d: *Disk, key: Digest) void { + var id = d.by_key.get(key) orelse return; + while (id != none) : (id = d.recs.items[id].base) { + d.unlink(id); + d.link(id); + } + // bases end up newest: walk order keeps the entry itself older than its bases, so trims reach dependents first + } + + /// The files a restore of key reads, base first, with the logical pages each gives. + pub fn extents(d: *Disk, key: Digest, out: *std.ArrayList(Extent)) !void { + out.clearRetainingCapacity(); + var id = d.by_key.get(key) orelse return error.FileNotFound; + var hi = d.recs.items[id].npages; + while (id != none) { + const r = &d.recs.items[id]; + try out.append(d.gpa, .{ .key = r.key, .first_page = r.first_page, .lo = r.first_page, .hi = hi }); + hi = r.first_page; + id = r.base; + } + std.mem.reverse(Extent, out.items); + } + + /// Deletes a file without dependents (error when one stacks on it). + pub fn remove(d: *Disk, key: Digest) !void { + const id = d.by_key.get(key) orelse return; + const r = &d.recs.items[id]; + if (r.refs > 0) return error.HasDependents; + if (r.base != none) d.recs.items[r.base].refs -= 1; + d.unlink(id); + _ = d.by_key.remove(key); + d.used -= r.size; + r.live = false; + var nb: [48]u8 = undefined; + _ = linux.unlinkat(d.dir.handle, fileName(key, "tfs", &nb), 0); + try d.free_recs.append(d.gpa, id); + } + + /// Deletes least recently used files without dependents while over budget; their keys go to out. + pub fn trim(d: *Disk, out: *std.ArrayList(Digest)) !void { + var id = d.oldest; + while (d.used > d.opts.budget and id != none) { + const next = d.recs.items[id].newer; + if (d.recs.items[id].refs == 0) { + const key = d.recs.items[id].key; + try d.remove(key); + try out.append(d.gpa, key); + d.stats.evicted += 1; + id = d.oldest; // a removed file may free its base + continue; + } + id = next; + } + } +}; + +/// Owned (this rank's) logical pages of family f in [first, k): the file's index of logical page k. +pub fn ownedIn(f: Family, sink: *const ReadSink, first: u32, k: u32) u32 { + if (!(f.split and sink.world > 1)) return k - first; + var by: [pool_mod.max_world]u32 = undefined; + pool_mod.span(first, k, sink.world, &by); + return by[sink.rank]; +} + +/// The logical page of a family's j-th stored page of a file that starts at first. +pub fn logicalOf(f: Family, sink: *const ReadSink, first: u32, j: u32) u32 { + if (!(f.split and sink.world > 1)) return first + j; + const w = sink.world; + const r0 = first + (sink.rank + w - first % w) % w; // the first owned logical page at or after first + return r0 + j * w; +} diff --git a/zig/src/sessions/disk_test.zig b/zig/src/sessions/disk_test.zig new file mode 100644 index 0000000000..6ac0525e45 --- /dev/null +++ b/zig/src/sessions/disk_test.zig @@ -0,0 +1,125 @@ +//! The NVMe tier on real files: whole and delta files round-trip bit for bit, split ranks, damage. + +const std = @import("std"); +const testing = std.testing; +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Family = pool_mod.Family; +const HostPool = @import("pagestore.zig").HostPool; +const disk = @import("disk.zig"); + +const fams = [_]Family{ + .{ .name = "comp.2", .ratio = 2, .row_bytes = 584, .split = true }, + .{ .name = "index_k.2", .ratio = 2, .row_bytes = 132 }, + .{ .name = "comp.20", .ratio = 1, .row_bytes = 584, .split = true }, +}; + +/// Fills a page of every family with bytes that name (family, logical page, byte). +fn fill(h: *HostPool, p: *const Pool, logical: u32, pg: u32, salt: u8) void { + for (fams, 0..) |f, fi| { + if (f.split and !p.owns(pg)) continue; + const v = h.pageOf(@intCast(fi), p.familyPage(f, pg)); + for (v, 0..) |*b, i| b.* = @truncate(i *% 131 +% logical *% 7 +% fi *% 29 +% salt); + } +} + +fn familyPages(gpa: std.mem.Allocator, p: *const Pool, pages: []const u32, first: u32, owned_only: bool) ![][]u32 { + const out = try gpa.alloc([]u32, fams.len); + for (fams, 0..) |f, fi| { + var l: std.ArrayList(u32) = .empty; + for (pages[first..], first..) |pg, k| { + if (f.split and p.world > 1 and owned_only and k % p.world != p.rank) continue; + try l.append(gpa, p.familyPage(f, pg)); + } + out[fi] = try l.toOwnedSlice(gpa); + } + return out; +} + +fn freePages(gpa: std.mem.Allocator, x: [][]u32) void { + for (x) |l| gpa.free(l); + gpa.free(x); +} + +fn roundTrip(world: u32, rank: u32) !void { + const gpa = testing.allocator; + const io = testing.io; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + const root = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(root); + var p = try Pool.init(gpa, .{ .families = &fams, .page = 256 }, 64 * 256, world, rank); + defer p.deinit(); + var h = try HostPool.init(gpa, &p); + defer h.deinit(); + // small chunks so pool segments span many chunks; few staging buffers so the pump waits on the lanes + const d = try disk.Disk.open(gpa, io, .{ .root = root, .compat = @splat(5), .rank = rank, .world = world, .chunk = 3 * 74_752, .staging = 3, .lanes = 3, .min_tokens = 1 }); + defer d.close(); + const a = try p.newSlot(32 * 256); + const pos: u64 = 9 * 256 + 100; + try a.ensure(pos); + for (a.mapped(), 0..) |pg, k| fill(&h, &p, @intCast(k), pg, 1); + const ids = try gpa.alloc(i32, pos); + defer gpa.free(ids); + for (ids, 0..) |*t, i| t.* = @intCast(i * 3 + 1); + var blob: [5000]u8 = undefined; + for (&blob, 0..) |*x, i| x.* = @truncate(9 - i % 5); + const fp = try familyPages(gpa, &p, a.mapped(), 0, true); + defer freePages(gpa, fp); + const src: disk.WriteSrc = .{ .key = @splat(1), .tag = 4, .page = 256, .pos = pos, .npages = a.len, .ids = ids, .store = h.store(), .families = &fams, .pages = fp, .blobs = &.{.{ .name = "rings", .bytes = &blob }} }; + const size = try d.writeFile(&src); + _ = try d.register(&src, size); + // a newer turn stacks a delta on it: pages [9, 14) only (page 9 was partial in the base) + const pos2: u64 = 13 * 256 + 7; + try a.ensure(pos2); + for (a.mapped()[9..], 9..) |pg, k| fill(&h, &p, @intCast(k), pg, 2); + const ids2 = try gpa.alloc(i32, pos2); + defer gpa.free(ids2); + for (ids2, 0..) |*t, i| t.* = @intCast(i * 3 + 1); + const fp2 = try familyPages(gpa, &p, a.mapped(), 9, true); + defer freePages(gpa, fp2); + const src2: disk.WriteSrc = .{ .key = @splat(2), .base = @splat(1), .tag = 4, .page = 256, .pos = pos2, .first_page = 9, .npages = a.len, .ids = ids2[9 * 256 ..], .store = h.store(), .families = &fams, .pages = fp2, .blobs = &.{.{ .name = "rings", .bytes = blob[0..77] }} }; + const size2 = try d.writeFile(&src2); + _ = try d.register(&src2, size2); + try testing.expect(size2 < size); + try testing.expectError(error.HasDependents, d.remove(@splat(1))); + // restore the delta onto other pages through its chain + const b = try p.newSlot(32 * 256); + try b.ensure(pos2); + var ext: std.ArrayList(disk.Extent) = .empty; + defer ext.deinit(gpa); + try d.extents(@splat(2), &ext); + try testing.expectEqual(@as(usize, 2), ext.items.len); + try testing.expectEqual(@as(u32, 9), ext.items[0].hi); + const tp = try familyPages(gpa, &p, b.mapped(), 0, false); + defer freePages(gpa, tp); + const sink: disk.ReadSink = .{ .store = h.store(), .families = &fams, .page = 256, .world = world, .rank = rank, .pages = tp }; + try testing.expect(try d.readFile(ext.items[0], &sink, false) == null); + var got = (try d.readFile(ext.items[1], &sink, true)).?; + defer got.deinit(gpa); + try testing.expectEqualSlices(u8, blob[0..77], got.blobs[0].bytes); + for (a.mapped(), b.mapped()) |pa, pb| for (fams, 0..) |f, fi| { + if (f.split and !p.owns(pa)) continue; + try testing.expectEqualSlices(u8, h.pageOf(@intCast(fi), p.familyPage(f, pa)), h.pageOf(@intCast(fi), p.familyPage(f, pb))); + }; + // damage one chunk: the read refuses it before or as its rows land + var nb: [48]u8 = undefined; + const key2: [16]u8 = @splat(2); + const name = try std.fmt.bufPrint(&nb, "{x}.tfs", .{&key2}); + const f = try d.dir.openFile(io, name, .{ .mode = .read_write }); + var one = [_]u8{0}; + _ = try f.readPositional(io, &.{&one}, 3 * 4096 + 100); + one[0] ^= 0xff; + try f.writePositionalAll(io, &one, 3 * 4096 + 100); + f.close(io); + try testing.expectError(error.ChecksumMismatch, d.readFile(ext.items[1], &sink, true)); +} + +test "whole and delta files round-trip bit for bit (replicated)" { + try roundTrip(1, 0); +} + +test "split ranks write and restore their own pages (rank 0 and 1 of 2)" { + try roundTrip(2, 0); + try roundTrip(2, 1); +} diff --git a/zig/src/sessions/format.zig b/zig/src/sessions/format.zig new file mode 100644 index 0000000000..2ac391d156 --- /dev/null +++ b/zig/src/sessions/format.zig @@ -0,0 +1,197 @@ +//! A session file (format 2): a header, then segments of 4 KiB-aligned chunks, each with its own SHA-256. + +const std = @import("std"); +const Sha256 = std.crypto.hash.sha2.Sha256; +const up = @import("lanes.zig").up; + +pub const magic = "TFSESS02".*; +pub const version: u32 = 2; + +pub const SegKind = enum(u32) { ids = 1, pool = 2, blob = 3, _ }; + +pub const Head = extern struct { + magic: [8]u8 = magic, + version: u32 = version, + /// bytes of Head + segments + digests (before the pad) + head_bytes: u32 = 0, + /// SHA-256 of those bytes with this field zero + head_sha: [32]u8 = @splat(0), + key: [16]u8, + /// the entry whose files hold logical pages [0, first_page); zero: none + base: [16]u8 = @splat(0), + /// what makes these bytes mean the same thing to a build (family, knobs, layout, split) + compat: [16]u8, + tag: u32, + kind: u32, + page: u32, + world: u32, + rank: u32, + first_page: u32, + npages: u32, + nseg: u32, + ndigest: u32, + _pad: u32 = 0, + pos: u64, + total: u64 = 0, + + pub fn hasBase(h: *const Head) bool { + return !std.mem.allEqual(u8, &h.base, 0); + } +}; + +pub const Seg = extern struct { + name: [40]u8, + kind: SegKind, + /// pool: the family's index in the layout + family: u32, + nbytes: u64, + offset: u64 = 0, + /// bytes a chunk (the last may be shorter) and file bytes a chunk (4 KiB multiple) + chunk: u32, + stride: u32 = 0, + first_digest: u32 = 0, + nchunks: u32 = 0, + + pub fn nameOf(s: *const Seg) []const u8 { + return std.mem.sliceTo(&s.name, 0); + } + + /// Chunk i's data bytes and file span. + pub fn chunkAt(s: *const Seg, i: u32) struct { off: u64, data: u32, span: u32 } { + const data: u32 = @intCast(@min(@as(u64, s.chunk), s.nbytes - @as(u64, i) * s.chunk)); + return .{ .off = s.offset + @as(u64, i) * s.stride, .data = data, .span = @intCast(up(data)) }; + } +}; + +pub const Digest32 = [32]u8; + +pub fn segName(name: []const u8) [40]u8 { + var n: [40]u8 = @splat(0); + @memcpy(n[0..@min(name.len, 39)], name[0..@min(name.len, 39)]); + return n; +} + +/// A file's layout: offsets, chunks and the digest slots, before or after its chunks are written. +pub const Layout = struct { + head: Head, + segs: []Seg, + digests: []Digest32, + /// file bytes of the padded header + head_span: u64, + + /// Places segments (each with name, kind, family, nbytes, chunk) after the header. + pub fn build(gpa: std.mem.Allocator, head: Head, segs_in: []const Seg) !Layout { + const segs = try gpa.dupe(Seg, segs_in); + errdefer gpa.free(segs); + var nd: u32 = 0; + for (segs) |*s| { + if (s.chunk == 0) return error.BadSegment; + s.nchunks = @intCast((s.nbytes + s.chunk - 1) / s.chunk); + s.stride = @intCast(up(s.chunk)); + s.first_digest = nd; + nd += s.nchunks; + } + var h = head; + h.nseg = @intCast(segs.len); + h.ndigest = nd; + h.head_bytes = @intCast(@sizeOf(Head) + segs.len * @sizeOf(Seg) + nd * @sizeOf(Digest32)); + const span = up(h.head_bytes); + var off = span; + for (segs) |*s| { + s.offset = off; + if (s.nchunks > 0) off += @as(u64, s.nchunks - 1) * s.stride + up(s.nbytes - @as(u64, s.nchunks - 1) * s.chunk); + } + h.total = off; + const d = try gpa.alloc(Digest32, nd); + @memset(d, @splat(0)); + return .{ .head = h, .segs = segs, .digests = d, .head_span = span }; + } + + pub fn deinit(l: *Layout, gpa: std.mem.Allocator) void { + gpa.free(l.segs); + gpa.free(l.digests); + l.* = undefined; + } + + /// The header's bytes with its checksum into out (len >= head_span; the pad is zeroed). + pub fn encode(l: *Layout, out: []u8) void { + @memset(out[0..l.head_span], 0); + var h = l.head; + h.head_sha = @splat(0); + const hb = std.mem.asBytes(&h); + const sb = std.mem.sliceAsBytes(l.segs); + const db = std.mem.sliceAsBytes(l.digests); + @memcpy(out[0..hb.len], hb); + @memcpy(out[hb.len..][0..sb.len], sb); + @memcpy(out[hb.len + sb.len ..][0..db.len], db); + var sha: [32]u8 = undefined; + Sha256.hash(out[0..h.head_bytes], &sha, .{}); + @memcpy(out[@offsetOf(Head, "head_sha")..][0..32], &sha); + l.head.head_sha = sha; + } + + /// The header read back from bytes (at least head_bytes of them), every field and offset checked. + pub fn decode(gpa: std.mem.Allocator, bytes: []const u8) !Layout { + if (bytes.len < @sizeOf(Head)) return error.BadHeader; + var h: Head = undefined; + @memcpy(std.mem.asBytes(&h), bytes[0..@sizeOf(Head)]); + if (!std.mem.eql(u8, &h.magic, &magic) or h.version != version) return error.BadHeader; + const want = @sizeOf(Head) + @as(u64, h.nseg) * @sizeOf(Seg) + @as(u64, h.ndigest) * @sizeOf(Digest32); + if (h.head_bytes != want or bytes.len < want) return error.BadHeader; + var check: [32]u8 = undefined; + var sha = Sha256.init(.{}); + sha.update(bytes[0..@offsetOf(Head, "head_sha")]); + sha.update(&@as([32]u8, @splat(0))); + sha.update(bytes[@offsetOf(Head, "head_sha") + 32 .. want]); + sha.final(&check); + if (!std.mem.eql(u8, &check, &h.head_sha)) return error.ChecksumMismatch; + const segs = try gpa.alloc(Seg, h.nseg); + errdefer gpa.free(segs); + @memcpy(std.mem.sliceAsBytes(segs), bytes[@sizeOf(Head)..][0 .. segs.len * @sizeOf(Seg)]); + const d = try gpa.alloc(Digest32, h.ndigest); + errdefer gpa.free(d); + @memcpy(std.mem.sliceAsBytes(d), bytes[@sizeOf(Head) + segs.len * @sizeOf(Seg) ..][0 .. d.len * 32]); + var nd: u64 = 0; + for (segs) |s| { + if (s.chunk == 0 or s.stride != up(s.chunk) or s.first_digest != nd or s.offset % 4096 != 0) return error.BadHeader; + if (s.nchunks != (s.nbytes + s.chunk - 1) / s.chunk) return error.BadHeader; + if (s.nchunks > 0 and s.offset + @as(u64, s.nchunks - 1) * s.stride + up(s.nbytes - @as(u64, s.nchunks - 1) * s.chunk) > h.total) return error.BadHeader; + nd += s.nchunks; + } + if (nd != h.ndigest) return error.BadHeader; + return .{ .head = h, .segs = segs, .digests = d, .head_span = up(h.head_bytes) }; + } + + pub fn find(l: *const Layout, kind: SegKind, name: []const u8) ?*const Seg { + for (l.segs) |*s| if (s.kind == kind and std.mem.eql(u8, s.nameOf(), name)) return s; + return null; + } +}; + +test "a layout encodes, decodes, and refuses a damaged header" { + const gpa = std.testing.allocator; + const head: Head = .{ .key = @splat(1), .compat = @splat(2), .tag = 3, .kind = 0, .page = 256, .world = 2, .rank = 1, .first_page = 4, .npages = 9, .nseg = 0, .ndigest = 0, .pos = 2200 }; + var l = try Layout.build(gpa, head, &.{ + .{ .name = segName("ids"), .kind = .ids, .family = 0, .nbytes = 4 * (2200 - 1024), .chunk = 1 << 20 }, + .{ .name = segName("comp.20"), .kind = .pool, .family = 0, .nbytes = 3 * 149_504, .chunk = 2 * 149_504 }, + .{ .name = segName("taps"), .kind = .blob, .family = 0, .nbytes = 0, .chunk = 4096 }, + }); + defer l.deinit(gpa); + try std.testing.expectEqual(@as(u32, 1 + 2 + 0), l.head.ndigest); + try std.testing.expectEqual(@as(u64, 4096), l.segs[0].offset); + try std.testing.expectEqual(@as(u64, 4096 + up(4 * 1176)), l.segs[1].offset); + const c1 = l.segs[1].chunkAt(1); + try std.testing.expectEqual(@as(u32, 149_504), c1.data); + try std.testing.expectEqual(l.segs[1].offset + up(2 * 149_504), c1.off); + try std.testing.expectEqual(l.segs[2].offset, l.head.total); + l.digests[1][5] = 9; + var buf: [8192]u8 = undefined; + l.encode(&buf); + var back = try Layout.decode(gpa, &buf); + defer back.deinit(gpa); + try std.testing.expectEqual(l.head.total, back.head.total); + try std.testing.expectEqual(@as(u8, 9), back.digests[1][5]); + try std.testing.expectEqualStrings("comp.20", back.find(.pool, "comp.20").?.nameOf()); + buf[200] ^= 0x40; + try std.testing.expectError(error.ChecksumMismatch, Layout.decode(gpa, &buf)); +} diff --git a/zig/src/sessions/freeset.zig b/zig/src/sessions/freeset.zig new file mode 100644 index 0000000000..b8cce323af --- /dev/null +++ b/zig/src/sessions/freeset.zig @@ -0,0 +1,114 @@ +//! A set of free page indices, lowest first in O(log64 n): a bit a page, then a bit a non-empty word a level up. + +const std = @import("std"); + +const max_levels = 6; // 64^6 pages: more than any pool + +pub const FreeSet = struct { + /// levels[0] holds a bit a page, levels[k] a bit a non-empty word of levels[k - 1] + levels: [max_levels][]u64 = @splat(&.{}), + depth: u8 = 0, + n: u32 = 0, + free: u32 = 0, + + /// All of 0 .. n - 1 free. + pub fn init(gpa: std.mem.Allocator, n: u32) !FreeSet { + var s: FreeSet = .{ .n = n }; + errdefer s.deinit(gpa); + var bits: usize = @max(n, 1); + while (true) { + const words = (bits + 63) / 64; + s.levels[s.depth] = try gpa.alloc(u64, words); + @memset(s.levels[s.depth], 0); + s.depth += 1; + if (words == 1) break; + bits = words; + } + for (0..n) |i| s.setBit(@intCast(i)); + s.free = n; + return s; + } + + pub fn deinit(s: *FreeSet, gpa: std.mem.Allocator) void { + for (s.levels[0..s.depth]) |l| gpa.free(l); + s.* = .{}; + } + + pub fn count(s: *const FreeSet) u32 { + return s.free; + } + + pub fn contains(s: *const FreeSet, i: u32) bool { + return s.levels[0][i / 64] & (@as(u64, 1) << @intCast(i % 64)) != 0; + } + + /// The lowest free index, taken; null when none is free. + pub fn takeLowest(s: *FreeSet) ?u32 { + if (s.free == 0) return null; + var idx: usize = 0; + var lv: usize = s.depth; + while (lv > 0) { + lv -= 1; + const w = s.levels[lv][idx]; + idx = idx * 64 + @ctz(w); + } + const i: u32 = @intCast(idx); + s.clearBit(i); + s.free -= 1; + return i; + } + + /// Takes a given free index (error when it is not free). + pub fn take(s: *FreeSet, i: u32) error{NotFree}!void { + if (i >= s.n or !s.contains(i)) return error.NotFree; + s.clearBit(i); + s.free -= 1; + } + + /// Gives an index back (error when it is already free: a double free). + pub fn give(s: *FreeSet, i: u32) error{DoubleFree}!void { + if (i >= s.n or s.contains(i)) return error.DoubleFree; + s.setBit(i); + s.free += 1; + } + + fn setBit(s: *FreeSet, i: u32) void { + var idx: usize = i; + for (s.levels[0..s.depth]) |l| { + const was = l[idx / 64]; + l[idx / 64] = was | (@as(u64, 1) << @intCast(idx % 64)); + if (was != 0) return; // the word was non-empty: the levels above already say so + idx /= 64; + } + } + + fn clearBit(s: *FreeSet, i: u32) void { + var idx: usize = i; + for (s.levels[0..s.depth]) |l| { + l[idx / 64] &= ~(@as(u64, 1) << @intCast(idx % 64)); + if (l[idx / 64] != 0) return; // the word still has a free bit: the levels above stay set + idx /= 64; + } + } +}; + +test "lowest first, give back, double free refused, across levels" { + const gpa = std.testing.allocator; + var s = try FreeSet.init(gpa, 70_000); // three levels + defer s.deinit(gpa); + try std.testing.expectEqual(@as(u8, 3), s.depth); + for (0..5000) |i| try std.testing.expectEqual(@as(?u32, @intCast(i)), s.takeLowest()); + try s.give(4097); + try s.give(17); + try std.testing.expectError(error.DoubleFree, s.give(17)); + try std.testing.expectEqual(@as(?u32, 17), s.takeLowest()); + try std.testing.expectEqual(@as(?u32, 4097), s.takeLowest()); + try std.testing.expectEqual(@as(?u32, 5000), s.takeLowest()); + try s.take(69_999); + try std.testing.expectError(error.NotFree, s.take(69_999)); + try std.testing.expectEqual(@as(u32, 70_000 - 5001 - 1), s.count()); + while (s.takeLowest()) |_| {} + try std.testing.expectEqual(@as(u32, 0), s.count()); + try s.give(69_999); + try std.testing.expectEqual(@as(?u32, 69_999), s.takeLowest()); +} diff --git a/zig/src/sessions/jobs.zig b/zig/src/sessions/jobs.zig new file mode 100644 index 0000000000..5b445f20e7 --- /dev/null +++ b/zig/src/sessions/jobs.zig @@ -0,0 +1,150 @@ +//! Parks and restores off the round: a tier thread runs each job's file I/O until the plan settles it (Store.settle). + +const std = @import("std"); +const Io = std.Io; +const disk_mod = @import("disk.zig"); +const Disk = disk_mod.Disk; + +pub const Kind = enum { park, restore }; +pub const State = enum(u8) { queued, running, done, failed }; + +pub const Job = struct { + kind: Kind, + /// the store's entry + entry: u32, + state: std.atomic.Value(State) = .init(.queued), + err: ?anyerror = null, + started_ns: i96 = 0, + elapsed_ns: u64 = 0, + /// park: the file to write (owned buffers below) and its size once written + src: disk_mod.WriteSrc = undefined, + size: u64 = 0, + /// restore: files base first, the sink, the bounded state read + extents: []disk_mod.Extent = &.{}, + sink: disk_mod.ReadSink = undefined, + restored: ?disk_mod.Restored = null, + /// restore: the slot it fills + slot: u32 = 0, + /// pool pages the job holds until settled + held: []u32 = &.{}, + /// owned buffers + ids: []i32 = &.{}, + lists: [][]u32 = &.{}, + blobs: []disk_mod.Blob = &.{}, + + pub fn done(j: *const Job) bool { + const s = j.state.load(.acquire); + return s == .done or s == .failed; + } + + pub fn deinit(j: *Job, gpa: std.mem.Allocator) void { + gpa.free(j.held); + gpa.free(j.ids); + for (j.lists) |l| gpa.free(l); + gpa.free(j.lists); + for (j.blobs) |b| { + gpa.free(b.name); + gpa.free(b.bytes); + } + gpa.free(j.blobs); + gpa.free(j.extents); + if (j.restored) |*r| r.deinit(gpa); + gpa.destroy(j); + } + + fn run(j: *Job, d: *Disk, io: Io) void { + j.state.store(.running, .release); + const t0 = Io.Clock.awake.now(io).toNanoseconds(); + j.err = switch (j.kind) { + .park => if (d.writeFile(&j.src)) |n| blk: { + j.size = n; + break :blk null; + } else |e| e, + .restore => j.restore(d), + }; + j.elapsed_ns = @intCast(Io.Clock.awake.now(io).toNanoseconds() - t0); + j.state.store(if (j.err == null) .done else .failed, .release); + } + + fn restore(j: *Job, d: *Disk) ?anyerror { + for (j.extents, 0..) |e, i| { + const last = i + 1 == j.extents.len; + const r = d.readFile(e, &j.sink, last) catch |err| return err; + if (last) j.restored = r; + } + j.sink.store.publish() catch |err| return err; + return null; + } +}; + +/// The tier thread and its queue (null thread: jobs run inline). +pub const Runner = struct { + gpa: std.mem.Allocator, + io: Io, + disk: *Disk, + thread: ?std.Thread = null, + queue: std.ArrayList(*Job) = .empty, + head: usize = 0, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + stop: bool = false, + + pub fn init(gpa: std.mem.Allocator, io: Io, d: *Disk) Runner { + return .{ .gpa = gpa, .io = io, .disk = d }; + } + + /// Starts the tier thread: jobs then run off the caller's thread. + pub fn startThread(r: *Runner) !void { + if (r.thread != null) return; + r.thread = try std.Thread.spawn(.{ .stack_size = 1 << 20 }, loop, .{r}); + } + + pub fn deinit(r: *Runner) void { + if (r.thread) |t| { + r.mutex.lockUncancelable(r.io); + r.stop = true; + r.cond.broadcast(r.io); + r.mutex.unlock(r.io); + t.join(); + } + r.queue.deinit(r.gpa); + } + + pub fn submit(r: *Runner, j: *Job) !void { + if (r.thread == null) return j.run(r.disk, r.io); + r.mutex.lockUncancelable(r.io); + defer r.mutex.unlock(r.io); + if (r.head == r.queue.items.len) { + r.queue.clearRetainingCapacity(); + r.head = 0; + } + try r.queue.append(r.gpa, j); + r.cond.signal(r.io); + } + + /// Waits for a job (the plan settles it next). + pub fn wait(r: *Runner, j: *Job) void { + while (!j.done()) { + r.mutex.lockUncancelable(r.io); + defer r.mutex.unlock(r.io); + if (!j.done()) r.cond.waitTimeout(r.io, &r.mutex, .{ .duration = .{ .raw = .fromMilliseconds(5), .clock = .awake } }) catch {}; + } + } + + fn loop(r: *Runner) void { + while (true) { + const j = blk: { + r.mutex.lockUncancelable(r.io); + defer r.mutex.unlock(r.io); + while (r.head == r.queue.items.len and !r.stop) r.cond.waitUncancelable(r.io, &r.mutex); + if (r.head == r.queue.items.len) return; + r.head += 1; + break :blk r.queue.items[r.head - 1]; + }; + j.run(r.disk, r.io); + r.mutex.lockUncancelable(r.io); + r.cond.broadcast(r.io); + r.mutex.unlock(r.io); + } + } +}; diff --git a/zig/src/sessions/lanes.zig b/zig/src/sessions/lanes.zig new file mode 100644 index 0000000000..38fad8cd89 --- /dev/null +++ b/zig/src/sessions/lanes.zig @@ -0,0 +1,336 @@ +//! The NVMe tier's lane threads, which hash and write (or read and check) 4 KiB chunks, and their staging buffers. + +const std = @import("std"); +const linux = std.os.linux; +const Io = std.Io; +const Sha256 = std.crypto.hash.sha2.Sha256; + +pub const block = 4096; + +pub fn up(n: u64) u64 { + return (n + block - 1) / block * block; +} + +/// Bounded, aligned staging buffers (allocated once; `release` gives the memory back while idle). +pub const Staging = struct { + io: Io, + size: usize, + count: u32, + bufs: [][]align(block) u8, + free_list: []u32, + free_n: u32, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + gpa: std.mem.Allocator, + + pub fn init(gpa: std.mem.Allocator, io: Io, count: u32, size: usize) !Staging { + if (count == 0 or size % block != 0) return error.BadStaging; + const bufs = try gpa.alloc([]align(block) u8, count); + errdefer gpa.free(bufs); + const fl = try gpa.alloc(u32, count); + errdefer gpa.free(fl); + for (bufs, 0..) |*b, i| { + b.* = &.{}; + fl[i] = @intCast(i); + } + return .{ .io = io, .size = size, .count = count, .bufs = bufs, .free_list = fl, .free_n = count, .gpa = gpa }; + } + + pub fn deinit(s: *Staging) void { + s.release(); + s.gpa.free(s.bufs); + s.gpa.free(s.free_list); + s.* = undefined; + } + + /// A free buffer's index, waiting for one; its memory is mapped on first use. + pub fn get(s: *Staging) !u32 { + s.mutex.lockUncancelable(s.io); + defer s.mutex.unlock(s.io); + while (s.free_n == 0) s.cond.waitUncancelable(s.io, &s.mutex); + s.free_n -= 1; + const i = s.free_list[s.free_n]; + if (s.bufs[i].len == 0) s.bufs[i] = try s.gpa.alignedAlloc(u8, .fromByteUnits(block), s.size); + return i; + } + + pub fn put(s: *Staging, i: u32) void { + s.mutex.lockUncancelable(s.io); + defer s.mutex.unlock(s.io); + s.free_list[s.free_n] = i; + s.free_n += 1; + s.cond.signal(s.io); + } + + pub fn buf(s: *Staging, i: u32) []align(block) u8 { + return s.bufs[i]; + } + + /// Unmaps every buffer (all must be free): the memory goes back between operations. + pub fn release(s: *Staging) void { + std.debug.assert(s.free_n == s.count); + for (s.bufs) |*b| if (b.len != 0) { + s.gpa.free(b.*); + b.* = &.{}; + }; + } + + pub fn bytes(s: *const Staging) u64 { + return @as(u64, s.count) * s.size; + } +}; + +pub const Kind = enum { write, read }; + +/// One chunk's I/O: buf[0 .. span] at file offset off (span: the chunk padded to 4 KiB; data: the chunk's own bytes). +pub const Job = struct { + kind: Kind, + fd: linux.fd_t, + staging: u32, + data: u32, + span: u32, + off: u64, + /// write: the chunk's digest goes here; read: the digest it must have + digest: *[32]u8, + op: *Op, + /// the pump's tag for a finished read + tag: u32 = 0, +}; + +/// A park's or restore's chunks in flight: how many are pending, the first error, and finished reads for the pump. +pub const Op = struct { + io: Io, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + pending: u32 = 0, + err: ?anyerror = null, + /// finished reads (their jobs), drained by the pump + done: std.ArrayList(Job) = .empty, + gpa: std.mem.Allocator, + /// write jobs hand their staging back themselves + staging: *Staging, + + pub fn init(gpa: std.mem.Allocator, io: Io, staging: *Staging) !Op { + var o: Op = .{ .io = io, .gpa = gpa, .staging = staging }; + try o.done.ensureTotalCapacity(gpa, staging.count); + return o; + } + + pub fn deinit(o: *Op) void { + o.done.deinit(o.gpa); + } + + fn finish(o: *Op, j: Job, err: ?anyerror) void { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + o.pending -= 1; + if (err) |e| { + if (o.err == null) o.err = e; + } + if (j.kind == .read) o.done.appendAssumeCapacity(j) else o.staging.put(j.staging); + o.cond.broadcast(o.io); + } + + /// The next finished read, waiting for one; null when none is pending. + pub fn next(o: *Op) ?Job { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + while (o.done.items.len == 0 and o.pending > 0) o.cond.waitUncancelable(o.io, &o.mutex); + return o.done.pop(); + } + + /// Waits for every pending job; returns the first error. + pub fn wait(o: *Op) !void { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + while (o.pending > 0) o.cond.waitUncancelable(o.io, &o.mutex); + if (o.err) |e| return e; + } + + pub fn failed(o: *Op) ?anyerror { + o.mutex.lockUncancelable(o.io); + defer o.mutex.unlock(o.io); + return o.err; + } +}; + +/// The lane threads and their queue. +pub const Lanes = struct { + io: Io, + gpa: std.mem.Allocator, + staging: *Staging, + threads: []std.Thread, + /// queued jobs: a ring as long as the staging set (every queued job holds a buffer) + ring: []Job, + head: usize = 0, + len: usize = 0, + mutex: Io.Mutex = .init, + cond: Io.Condition = .init, + stop: bool = false, + /// seconds the lanes spent hashing and in I/O (summed over lanes) + hash_ns: std.atomic.Value(u64) = .init(0), + io_ns: std.atomic.Value(u64) = .init(0), + + pub fn start(gpa: std.mem.Allocator, io: Io, staging: *Staging, n: u32) !*Lanes { + const l = try gpa.create(Lanes); + errdefer gpa.destroy(l); + const ring = try gpa.alloc(Job, staging.count); + errdefer gpa.free(ring); + l.* = .{ .io = io, .gpa = gpa, .staging = staging, .threads = try gpa.alloc(std.Thread, n), .ring = ring }; + errdefer gpa.free(l.threads); + var made: usize = 0; + errdefer { + l.halt(); + for (l.threads[0..made]) |t| t.join(); + } + for (l.threads) |*t| { + t.* = try std.Thread.spawn(.{ .stack_size = 1 << 20 }, run, .{l}); + made += 1; + } + return l; + } + + fn halt(l: *Lanes) void { + l.mutex.lockUncancelable(l.io); + l.stop = true; + l.cond.broadcast(l.io); + l.mutex.unlock(l.io); + } + + pub fn shutdown(l: *Lanes) void { + l.halt(); + for (l.threads) |t| t.join(); + l.gpa.free(l.ring); + l.gpa.free(l.threads); + l.gpa.destroy(l); + } + + /// Queues a job (each holds a staging buffer, so the ring never overflows). + pub fn submit(l: *Lanes, j: Job) void { + { + j.op.mutex.lockUncancelable(l.io); + defer j.op.mutex.unlock(l.io); + j.op.pending += 1; + } + l.mutex.lockUncancelable(l.io); + defer l.mutex.unlock(l.io); + std.debug.assert(l.len < l.ring.len); + l.ring[(l.head + l.len) % l.ring.len] = j; + l.len += 1; + l.cond.signal(l.io); + } + + fn run(l: *Lanes) void { + while (true) { + const j = blk: { + l.mutex.lockUncancelable(l.io); + defer l.mutex.unlock(l.io); + while (l.len == 0 and !l.stop) l.cond.waitUncancelable(l.io, &l.mutex); + if (l.len == 0) return; + const j = l.ring[l.head]; + l.head = (l.head + 1) % l.ring.len; + l.len -= 1; + break :blk j; + }; + j.op.finish(j, l.work(j)); + } + } + + fn work(l: *Lanes, j: Job) ?anyerror { + const b = l.staging.buf(j.staging); + var t0 = std.Io.Clock.awake.now(l.io).toNanoseconds(); + switch (j.kind) { + .write => { + Sha256.hash(b[0..j.data], j.digest, .{}); + const t1 = std.Io.Clock.awake.now(l.io).toNanoseconds(); + _ = l.hash_ns.fetchAdd(@intCast(t1 - t0), .monotonic); + t0 = t1; + @memset(b[j.data..j.span], 0); + pwriteAll(j.fd, b[0..j.span], j.off) catch |e| return e; + _ = l.io_ns.fetchAdd(@intCast(std.Io.Clock.awake.now(l.io).toNanoseconds() - t0), .monotonic); + }, + .read => { + preadAll(j.fd, b[0..j.span], j.off) catch |e| return e; + const t1 = std.Io.Clock.awake.now(l.io).toNanoseconds(); + _ = l.io_ns.fetchAdd(@intCast(t1 - t0), .monotonic); + var d: [32]u8 = undefined; + Sha256.hash(b[0..j.data], &d, .{}); + _ = l.hash_ns.fetchAdd(@intCast(std.Io.Clock.awake.now(l.io).toNanoseconds() - t1), .monotonic); + if (!std.mem.eql(u8, &d, j.digest)) return error.ChecksumMismatch; + }, + } + return null; + } +}; + +pub fn pwriteAll(fd: linux.fd_t, data: []const u8, off: u64) !void { + var done: usize = 0; + while (done < data.len) { + const rc = linux.pwrite(fd, data[done..].ptr, data.len - done, @intCast(off + done)); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .INTR => continue, + .NOSPC => return error.NoSpaceLeft, + else => return error.WriteFailed, + } + if (rc == 0) return error.WriteFailed; + done += rc; + } +} + +pub fn preadAll(fd: linux.fd_t, out: []u8, off: u64) !void { + var done: usize = 0; + while (done < out.len) { + const rc = linux.pread(fd, out[done..].ptr, out.len - done, @intCast(off + done)); + switch (linux.errno(rc)) { + .SUCCESS => {}, + .INTR => continue, + else => return error.ReadFailed, + } + if (rc == 0) return error.ShortRead; + done += rc; + } +} + +test "lanes write, read back and check chunks through bounded staging" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var st = try Staging.init(gpa, io, 3, 2 * block); + defer st.deinit(); + const l = try Lanes.start(gpa, io, &st, 2); + defer l.shutdown(); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const f = try tmp.dir.createFile(io, "lanes.bin", .{ .read = true }); + defer f.close(io); + const fd = f.handle; + var digests: [7][32]u8 = undefined; + var op = try Op.init(gpa, io, &st); + defer op.deinit(); + for (0..7) |i| { + const s = try st.get(); + const n: u32 = if (i == 6) 1000 else 2 * block - 7; + for (st.buf(s)[0..n], 0..) |*b, k| b.* = @truncate(k * 31 + i); + l.submit(.{ .kind = .write, .fd = fd, .staging = s, .data = n, .span = @intCast(up(n)), .off = i * 2 * block, .digest = &digests[i], .op = &op }); + } + try op.wait(); + var rop = try Op.init(gpa, io, &st); + defer rop.deinit(); + var seen: u32 = 0; + var sent: u32 = 0; + while (seen < 7) { + while (sent < 7 and sent - seen < 3) : (sent += 1) { + const n: u32 = if (sent == 6) 1000 else 2 * block - 7; + l.submit(.{ .kind = .read, .fd = fd, .staging = try st.get(), .data = n, .span = @intCast(up(n)), .off = sent * 2 * block, .digest = &digests[sent], .op = &rop, .tag = sent }); + } + const j = rop.next().?; + for (st.buf(j.staging)[0..j.data], 0..) |b, k| try std.testing.expectEqual(@as(u8, @truncate(k * 31 + j.tag)), b); + st.put(j.staging); + seen += 1; + } + try rop.wait(); + digests[2][0] ^= 1; // a damaged chunk is refused + l.submit(.{ .kind = .read, .fd = fd, .staging = try st.get(), .data = 2 * block - 7, .span = 2 * block, .off = 2 * 2 * block, .digest = &digests[2], .op = &rop }); + st.put(rop.next().?.staging); + try std.testing.expectError(error.ChecksumMismatch, rop.wait()); +} diff --git a/zig/src/sessions/pagestore.zig b/zig/src/sessions/pagestore.zig new file mode 100644 index 0000000000..0e4792699b --- /dev/null +++ b/zig/src/sessions/pagestore.zig @@ -0,0 +1,129 @@ +//! The bytes behind the pool's pages: a family's tensor of pages, read and written a page run at a time. + +const std = @import("std"); +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Family = pool_mod.Family; + +pub const PageStore = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + /// family fam's pages, back to back into dst (pages.len x a page's bytes) + read: *const fn (ptr: *anyopaque, fam: u32, pages: []const u32, dst: []u8) anyerror!void, + /// src (pages.len x a page's bytes) into family fam's pages + write: *const fn (ptr: *anyopaque, fam: u32, pages: []const u32, src: []const u8) anyerror!void, + /// one page's rows of family fam into another page (copy on write of a RAM resume's partial page) + copyPage: *const fn (ptr: *anyopaque, fam: u32, src: u32, dst: u32) anyerror!void, + /// round thread: a point after every row the compute stream has written so far + fence: *const fn (ptr: *anyopaque) anyerror!u64, + /// before reads: wait until the fence's rows are visible + waitFence: *const fn (ptr: *anyopaque, token: u64) anyerror!void, + /// after writes: make them visible to the compute stream + publish: *const fn (ptr: *anyopaque) anyerror!void, + }; + + pub fn read(s: PageStore, fam: u32, pages: []const u32, dst: []u8) !void { + return s.vtable.read(s.ptr, fam, pages, dst); + } + pub fn write(s: PageStore, fam: u32, pages: []const u32, src: []const u8) !void { + return s.vtable.write(s.ptr, fam, pages, src); + } + pub fn copyPage(s: PageStore, fam: u32, src: u32, dst: u32) !void { + return s.vtable.copyPage(s.ptr, fam, src, dst); + } + pub fn fence(s: PageStore) !u64 { + return s.vtable.fence(s.ptr); + } + pub fn waitFence(s: PageStore, token: u64) !void { + return s.vtable.waitFence(s.ptr, token); + } + pub fn publish(s: PageStore) !void { + return s.vtable.publish(s.ptr); + } +}; + +/// Every family's pages in host memory: [familyTensorPages x page bytes] a family, zeroed (the null page stays zero). +pub const HostPool = struct { + gpa: std.mem.Allocator, + page: u32, + families: []const Family, + tensors: [][]u8, + + pub fn init(gpa: std.mem.Allocator, p: *const Pool) !HostPool { + const t = try gpa.alloc([]u8, p.layout.families.len); + errdefer gpa.free(t); + var made: usize = 0; + errdefer for (t[0..made]) |x| gpa.free(x); + for (p.layout.families, 0..) |f, i| { + t[i] = try gpa.alloc(u8, @as(usize, p.familyTensorPages(f)) * f.pageBytes(p.page)); + @memset(t[i], 0); + made += 1; + } + return .{ .gpa = gpa, .page = p.page, .families = p.layout.families, .tensors = t }; + } + + pub fn deinit(h: *HostPool) void { + for (h.tensors) |x| h.gpa.free(x); + h.gpa.free(h.tensors); + h.* = undefined; + } + + pub fn store(h: *HostPool) PageStore { + return .{ .ptr = h, .vtable = &vtable }; + } + + /// One page of family fam (a view). + pub fn pageOf(h: *HostPool, fam: u32, pg: u32) []u8 { + const n: usize = @intCast(h.families[fam].pageBytes(h.page)); + return h.tensors[fam][pg * n ..][0..n]; + } + + const vtable: PageStore.VTable = .{ .read = read, .write = write, .copyPage = copyPage, .fence = fence, .waitFence = waitFence, .publish = publish }; + + fn read(ptr: *anyopaque, fam: u32, pages: []const u32, dst: []u8) anyerror!void { + const h: *HostPool = @ptrCast(@alignCast(ptr)); + const n: usize = @intCast(h.families[fam].pageBytes(h.page)); + if (dst.len != pages.len * n) return error.BadLength; + for (pages, 0..) |pg, i| @memcpy(dst[i * n ..][0..n], h.pageOf(fam, pg)); + } + + fn write(ptr: *anyopaque, fam: u32, pages: []const u32, src: []const u8) anyerror!void { + const h: *HostPool = @ptrCast(@alignCast(ptr)); + const n: usize = @intCast(h.families[fam].pageBytes(h.page)); + if (src.len != pages.len * n) return error.BadLength; + for (pages, 0..) |pg, i| @memcpy(h.pageOf(fam, pg), src[i * n ..][0..n]); + } + + fn copyPage(ptr: *anyopaque, fam: u32, src: u32, dst: u32) anyerror!void { + const h: *HostPool = @ptrCast(@alignCast(ptr)); + @memcpy(h.pageOf(fam, dst), h.pageOf(fam, src)); + } + + fn fence(_: *anyopaque) anyerror!u64 { + return 0; + } + fn waitFence(_: *anyopaque, _: u64) anyerror!void {} + fn publish(_: *anyopaque) anyerror!void {} +}; + +test "host pool pages round-trip through the interface" { + const gpa = std.testing.allocator; + const fams = [_]Family{ .{ .name = "comp", .ratio = 1, .row_bytes = 8 }, .{ .name = "ik", .ratio = 2, .row_bytes = 4 } }; + var p = try Pool.init(gpa, .{ .families = &fams, .page = 16 }, 64, 1, 0); + defer p.deinit(); + var h = try HostPool.init(gpa, &p); + defer h.deinit(); + const s = h.store(); + var src: [2 * 128]u8 = undefined; + for (&src, 0..) |*b, i| b.* = @truncate(i * 7 + 3); + try s.write(0, &.{ 3, 1 }, &src); + var dst: [2 * 128]u8 = undefined; + try s.read(0, &.{ 3, 1 }, &dst); + try std.testing.expectEqualSlices(u8, &src, &dst); + try s.copyPage(0, 3, 0); + try std.testing.expectEqualSlices(u8, src[0..128], h.pageOf(0, 0)); + try std.testing.expect(std.mem.allEqual(u8, h.pageOf(0, p.null_page), 0)); + try std.testing.expectError(error.BadLength, s.read(1, &.{0}, &dst)); +} diff --git a/zig/src/sessions/pool.zig b/zig/src/sessions/pool.zig new file mode 100644 index 0000000000..d3e67516af --- /dev/null +++ b/zig/src/sessions/pool.zig @@ -0,0 +1,475 @@ +//! The paged KV pool's host book: lowest-free pages per residue mod W, holder counts, slot page tables, quotas. + +const std = @import("std"); +const FreeSet = @import("freeset.zig").FreeSet; + +pub const Error = error{ PoolExhausted, PastCapacity, SlotNotEmpty, DoubleFree, BadPool, ResidueMismatch } || std.mem.Allocator.Error; + +/// A cache family: rows of `row_bytes` every `ratio` tokens; `split` families shard by logical page over the ranks. +pub const Family = struct { + name: []const u8, + ratio: u32, + row_bytes: u32, + split: bool = false, + + pub fn pageRows(f: Family, page: u32) u32 { + return page / f.ratio; + } + + /// Bytes of one page of this family (a whole page's rows). + pub fn pageBytes(f: Family, page: u32) u64 { + return @as(u64, f.pageRows(page)) * f.row_bytes; + } +}; + +/// The pool's families and page size: what a token costs and how pages are laid out. +pub const Layout = struct { + families: []const Family, + /// tokens a page (a power of two; 256 for every family we serve) + page: u32 = 256, + + /// A rank's bytes a token: the split families count 1 / W. + pub fn bytesPerToken(l: Layout, world: u32) f64 { + var b: f64 = 0; + for (l.families) |f| b += @as(f64, @floatFromInt(f.row_bytes)) / @as(f64, @floatFromInt(f.ratio)) / @as(f64, @floatFromInt(if (f.split and world > 1) world else 1)); + return b; + } + + /// A rank's average bytes a page (split families 1 / W): what the residency budget counts. + pub fn pageBytes(l: Layout, world: u32) u64 { + var b: u64 = 0; + for (l.families) |f| b += f.pageBytes(l.page) / (if (f.split and world > 1) world else 1); + return b; + } +}; + +pub fn pagesFor(tokens: u64, page: u32) u32 { + return @intCast((tokens + page - 1) / page); +} + +/// Logical pages [k0, k1) counted by residue mod `w` (O(w)). +pub fn span(k0: u32, k1: u32, w: u32, out: []u32) void { + for (out[0..w], 0..) |*o, r| { + const rr: u32 = @intCast(r); + const first = k0 + (rr + w - k0 % w) % w; + o.* = if (first >= k1) 0 else (k1 - 1 - first) / w + 1; + } +} + +pub const max_world = 8; + +/// A file whose logical pages [0, pages) match a slot's or entry's rows: a later park writes only past it. +pub const Origin = struct { + key: [16]u8, + pages: u32, +}; + +/// One slot's page table: logical page i -> physical page; the device tables follow through the dirty range. +pub const Slot = struct { + pool: *Pool, + id: u32, + capacity: u64, + max_pages: u32, + /// physical page of each mapped logical page; len is the mapped count + pages: []u32, + len: u32 = 0, + /// pages the admission reserved for this slot's request (null: no admission control) + quota: ?u32 = null, + /// pages taken past the quota, from pages nobody had reserved + over: u32 = 0, + /// table entries changed since the device side last synced them: [lo, hi) + dirty_lo: u32 = std.math.maxInt(u32), + dirty_hi: u32 = 0, + /// the file the leading pages came from (truncation shortens it, release clears it) + origin: ?Origin = null, + + pub fn mapped(s: *const Slot) []const u32 { + return s.pages[0..s.len]; + } + + pub fn tokens(s: *const Slot) u64 { + return @as(u64, s.len) * s.pool.page; + } + + fn mark(s: *Slot, lo: u32, hi: u32) void { + s.dirty_lo = @min(s.dirty_lo, lo); + s.dirty_hi = @max(s.dirty_hi, hi); + } + + /// The table range to upload since the last call (empty when nothing changed): entries past len are the null page. + pub fn takeDirty(s: *Slot) ?struct { lo: u32, hi: u32 } { + if (s.dirty_lo >= s.dirty_hi) return null; + defer { + s.dirty_lo = std.math.maxInt(u32); + s.dirty_hi = 0; + } + return .{ .lo = s.dirty_lo, .hi = s.dirty_hi }; + } + + /// The device table's value at logical page i (the page, or the null page past len). + pub fn tableAt(s: *const Slot, i: u32) u32 { + return if (i < s.len) s.pages[i] else s.pool.null_page; + } + + /// The split families' device table at logical page i: the local page this rank stores, else the discard page. + pub fn localTableAt(s: *const Slot, i: u32) u32 { + return s.pool.localPage(s.tableAt(i)); + } + + fn reservedSpan(s: *const Slot, out: []u32) void { + const q = s.quota orelse 0; + span(s.len, @max(s.len, q), s.pool.world, out); + } + + /// Maps every page holding positions < end (a no-op when they are: the common case). + pub fn ensure(s: *Slot, end: u64) Error!void { + const p = s.pool; + const need = pagesFor(end, p.page); + if (need <= s.len) return; + if (need > s.max_pages) return error.PastCapacity; + const w = p.world; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + if (s.quota) |q| if (need > q) { + // past the reservation: only pages nobody else reserved, residue by residue + var extra: [max_world]u32 = undefined; + span(@max(q, s.len), need, w, &extra); + for (0..w) |r| if (extra[r] > p.free[r].count() -| (p.reserved[r] - before[r])) return error.PoolExhausted; + s.over += need - @max(q, s.len); + }; + var want: [max_world]u32 = undefined; + span(s.len, need, w, &want); + for (0..w) |r| if (want[r] > p.free[r].count()) return error.PoolExhausted; + const lo = s.len; + for (lo..need) |k| { + const pg = p.takeAt(@intCast(k)) catch unreachable; // counted above + p.hold(pg, 0, 1); // the taker's hold becomes the slot's + s.pages[k] = pg; + } + s.len = need; + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(lo, need); + } + + /// Maps shared pages as the first pages of an empty slot (a RAM resume: read only, the slot writes past them). + pub fn adopt(s: *Slot, shared: []const u32) Error!void { + if (s.len != 0) return error.SlotNotEmpty; + if (shared.len > s.max_pages) return error.PastCapacity; + const p = s.pool; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + for (shared, 0..) |pg, k| if (pg >= p.npages or pg % p.world != k % p.world) return error.ResidueMismatch; + for (shared, 0..) |pg, k| { + p.hold(pg, 1, 1); + s.pages[k] = pg; + } + s.len = @intCast(shared.len); + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(0, s.len); + } + + /// Appends one page taken for logical page len (the copy-on-write page of a RAM resume). + pub fn push(s: *Slot, pg: u32) Error!void { + if (s.len >= s.max_pages) return error.PastCapacity; + const p = s.pool; + if (p.world > 1 and pg % p.world != s.len % p.world) return error.ResidueMismatch; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + p.hold(pg, 0, 1); // the taker's hold becomes the slot's + s.pages[s.len] = pg; + s.len += 1; + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(s.len - 1, s.len); + } + + /// Lets go of the pages past position keep (rows nothing reads again); returns how many. + pub fn truncate(s: *Slot, keep: u64) Error!u32 { + const p = s.pool; + // positions >= keep get written again: a page holding one no longer equals the origin file's + if (s.origin) |*o| o.pages = @min(o.pages, @as(u32, @intCast(keep / p.page))); + const k = pagesFor(keep, p.page); + if (k >= s.len) return 0; + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + const n = s.len - k; + for (s.pages[k..s.len]) |pg| try p.release(pg, 1); + const old = s.len; + s.len = k; + p.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + p.addReserved(&after); + s.mark(k, old); + return n; + } + + /// Sets the admission's reservation (null: none). + pub fn reserve(s: *Slot, quota: ?u32) void { + var before: [max_world]u32 = undefined; + s.reservedSpan(&before); + s.quota = quota; + s.pool.unreserve(&before); + var after: [max_world]u32 = undefined; + s.reservedSpan(&after); + s.pool.addReserved(&after); + } + + /// Frees the slot's pages and its reservation. + pub fn releaseAll(s: *Slot) Error!u32 { + s.reserve(null); + s.over = 0; + s.origin = null; + return s.truncate(0); + } +}; + +/// The pool's book: free sets by residue, holders a page, the slots and their reservations. +pub const Pool = struct { + gpa: std.mem.Allocator, + layout: Layout, + npages: u32, + page: u32, + /// log2 of page tokens + shift: u5, + /// ranks the split families are sharded over (1: every rank stores every row) + world: u32, + rank: u32, + /// the page unmapped table entries point at (all zeros: a stray read reads zeros) + null_page: u32, + /// a split family's local page count (npages / W), its local null page and its discard page (other ranks' rows) + local_pages: u32, + local_null: u32, + discard: u32, + free: [max_world]FreeSet = @splat(.{}), + /// pages each residue's slots reserved and have not mapped yet + reserved: [max_world]u32 = @splat(0), + /// holders a page (slots + entries + in-flight jobs) and how many are slots + refs: []u16, + slot_refs: []u16, + /// pages held by something and by no slot (session entries alone): what parking can free + entry_only: u32 = 0, + slots: std.ArrayList(*Slot) = .empty, + + pub fn init(gpa: std.mem.Allocator, layout: Layout, tokens: u64, world: u32, rank: u32) Error!Pool { + if (layout.page == 0 or layout.page & (layout.page - 1) != 0) return error.BadPool; + if (tokens % layout.page != 0 or tokens == 0) return error.BadPool; + if (world == 0 or world > max_world or rank >= world) return error.BadPool; + const npages: u32 = @intCast(tokens / layout.page); + if (npages % world != 0) return error.BadPool; + for (layout.families) |f| if (f.ratio == 0 or layout.page % f.ratio != 0) return error.BadPool; + var p: Pool = .{ + .gpa = gpa, + .layout = layout, + .npages = npages, + .page = layout.page, + .shift = @intCast(std.math.log2_int(u32, layout.page)), + .world = world, + .rank = rank, + .null_page = npages, + .local_pages = npages / world, + .local_null = npages / world, + .discard = npages / world + 1, + .refs = &.{}, + .slot_refs = &.{}, + }; + errdefer p.deinit(); + for (0..world) |r| p.free[r] = try FreeSet.init(gpa, npages / world); + p.refs = try gpa.alloc(u16, npages); + @memset(p.refs, 0); + p.slot_refs = try gpa.alloc(u16, npages); + @memset(p.slot_refs, 0); + return p; + } + + pub fn deinit(p: *Pool) void { + for (p.slots.items) |s| { + p.gpa.free(s.pages); + p.gpa.destroy(s); + } + p.slots.deinit(p.gpa); + for (0..p.world) |r| p.free[r].deinit(p.gpa); + p.gpa.free(p.refs); + p.gpa.free(p.slot_refs); + p.* = undefined; + } + + pub fn newSlot(p: *Pool, capacity: u64) Error!*Slot { + const s = try p.gpa.create(Slot); + errdefer p.gpa.destroy(s); + const mp = pagesFor(capacity, p.page); + s.* = .{ .pool = p, .id = @intCast(p.slots.items.len), .capacity = capacity, .max_pages = mp, .pages = try p.gpa.alloc(u32, mp) }; + errdefer p.gpa.free(s.pages); + try p.slots.append(p.gpa, s); + return s; + } + + /// A free physical page for logical page k (its residue when split), held once by the caller. + pub fn takeAt(p: *Pool, k: u32) Error!u32 { + const r = k % p.world; + const local = p.free[r].takeLowest() orelse return error.PoolExhausted; + const pg = local * p.world + r; + p.hold(pg, 1, 0); // the caller's hold (an entry-like holder until a slot maps it) + return pg; + } + + /// Holders change on page pg: dref total, dslot of them slots; keeps entry_only and gives the page back at zero. + fn hold(p: *Pool, pg: u32, dref: i32, dslot: i32) void { + const r0 = p.refs[pg]; + const s0 = p.slot_refs[pg]; + const r1: u16 = @intCast(@as(i32, r0) + dref); + const s1: u16 = @intCast(@as(i32, s0) + dslot); + p.refs[pg] = r1; + p.slot_refs[pg] = s1; + const was = r0 > 0 and s0 == 0; + const now = r1 > 0 and s1 == 0; + if (was and !now) p.entry_only -= 1; + if (!was and now) p.entry_only += 1; + } + + /// One holder lets go (slot: from a slot's table); the page returns to its free set when nobody holds it. + fn release(p: *Pool, pg: u32, slot: u1) Error!void { + if (p.refs[pg] == 0) return error.DoubleFree; + p.hold(pg, -1, -@as(i32, slot)); + if (p.refs[pg] == 0) p.free[pg % p.world].give(pg / p.world) catch return error.DoubleFree; + } + + /// One more holder (a session entry or an in-flight job) of each page. + pub fn share(p: *Pool, pages: []const u32) void { + for (pages) |pg| p.hold(pg, 1, 0); + } + + /// Those holders let go. + pub fn drop(p: *Pool, pages: []const u32) Error!void { + for (pages) |pg| try p.release(pg, 0); + } + + pub fn holders(p: *const Pool, pg: u32) u16 { + return p.refs[pg]; + } + + fn unreserve(p: *Pool, by: []const u32) void { + for (0..p.world) |r| p.reserved[r] -= by[r]; + } + + fn addReserved(p: *Pool, by: []const u32) void { + for (0..p.world) |r| p.reserved[r] += by[r]; + } + + /// Free pages no admission reserved (split: W x the scarcer residue, so any logical run fits). + pub fn available(p: *const Pool) u32 { + var low: u32 = std.math.maxInt(u32); + for (0..p.world) |r| low = @min(low, p.free[r].count() -| p.reserved[r]); + return low * p.world; + } + + pub fn freePages(p: *const Pool) u32 { + var n: u32 = 0; + for (0..p.world) |r| n += p.free[r].count(); + return n; + } + + /// Pages an admission needs for a request (prompt + reply + slack, at most the slot's capacity). + pub fn needPages(p: *const Pool, prompt: u64, max_tokens: u64, capacity: u64, slack: u64) u32 { + return pagesFor(@min(prompt + max_tokens + slack, capacity), p.page); + } + + /// This rank stores the split families' rows of page pg. + pub fn owns(p: *const Pool, pg: u32) bool { + return pg != p.null_page and pg % p.world == p.rank; + } + + /// A split family's local page of pg: its own when owned, the local null for the null page, else the discard page. + pub fn localPage(p: *const Pool, pg: u32) u32 { + if (pg == p.null_page) return p.local_null; + return if (pg % p.world == p.rank) pg / p.world else p.discard; + } + + /// The page a family's tensor indexes for physical page pg (split families: localPage). + pub fn familyPage(p: *const Pool, f: Family, pg: u32) u32 { + return if (f.split and p.world > 1) p.localPage(pg) else pg; + } + + /// Pages a family's device tensor holds: npages + null, or (split) local pages + local null + discard page. + pub fn familyTensorPages(p: *const Pool, f: Family) u32 { + return if (f.split and p.world > 1) p.local_pages + 2 else p.npages + 1; + } + + /// Bytes this rank stores of family f for logical pages [first, first + n) of a slot or entry. + pub fn familyBytes(p: *const Pool, f: Family, first: u32, n: u32) u64 { + if (!(f.split and p.world > 1)) return @as(u64, n) * f.pageBytes(p.page); + var by: [max_world]u32 = undefined; + span(first, first + n, p.world, &by); + return @as(u64, by[p.rank]) * f.pageBytes(p.page); + } + + /// Bytes of every family's device tensor on this rank. + pub fn tensorBytes(p: *const Pool) u64 { + var b: u64 = 0; + for (p.layout.families) |f| b += @as(u64, p.familyTensorPages(f)) * f.pageBytes(p.page); + return b; + } + + /// Pages held by session entries alone and their bytes on this rank (the residency budget's measure): O(1). + pub fn entryOnlyBytes(p: *const Pool) u64 { + return @as(u64, p.entry_only) * p.layout.pageBytes(p.world); + } +}; + +test "residues, holders, quotas and the device tables" { + const gpa = std.testing.allocator; + const fams = [_]Family{ .{ .name = "comp", .ratio = 2, .row_bytes = 584, .split = true }, .{ .name = "index_k", .ratio = 2, .row_bytes = 132 } }; + var p = try Pool.init(gpa, .{ .families = &fams }, 16 * 256, 2, 1); + defer p.deinit(); + try std.testing.expectEqual(@as(u32, 8 + 2), p.familyTensorPages(fams[0])); + try std.testing.expectEqual(@as(u32, 16 + 1), p.familyTensorPages(fams[1])); + try std.testing.expectApproxEqAbs(@as(f64, 146 + 66), p.layout.bytesPerToken(2), 1e-9); + const a = try p.newSlot(8 * 256); + const b = try p.newSlot(8 * 256); + try a.ensure(3 * 256 + 1); // 4 pages: residues 0,1,0,1 + try std.testing.expectEqualSlices(u32, &.{ 0, 1, 2, 3 }, a.mapped()); + try b.ensure(256); + try std.testing.expectEqualSlices(u32, &.{4}, b.mapped()); + try std.testing.expectEqual(p.discard, a.localTableAt(0)); // page 0 is rank 0's + try std.testing.expectEqual(@as(u32, 0), a.localTableAt(1)); // page 1 -> rank 1's local page 0 + try std.testing.expectEqual(p.local_null, a.localTableAt(7)); + const d = a.takeDirty().?; + try std.testing.expectEqual(@as(u32, 0), d.lo); + try std.testing.expectEqual(@as(u32, 4), d.hi); + try std.testing.expect(a.takeDirty() == null); + // an entry shares a's pages; a lets go; the pages stay held by the entry alone + p.share(a.mapped()[0..3]); + try std.testing.expectEqual(@as(u32, 0), p.entry_only); + _ = try a.releaseAll(); + try std.testing.expectEqual(@as(u32, 3), p.entry_only); + try std.testing.expectEqual(@as(u32, 16 - 1 - 3), p.freePages()); + try p.drop(&.{ 0, 1, 2 }); + try std.testing.expectEqual(@as(u32, 0), p.entry_only); + try std.testing.expectError(error.DoubleFree, p.drop(&.{0})); + // quotas: b reserves 6 pages; available counts W x the scarcer residue + b.reserve(6); + var by: [2]u32 = undefined; + span(1, 6, 2, &by); + try std.testing.expectEqual(@as(u32, 2), by[0]); + try std.testing.expectEqual(@as(u32, 3), by[1]); + try std.testing.expectEqual(@as(u32, 2 * @min(7 - 2, 8 - 3)), p.available()); + try b.ensure(6 * 256); + try std.testing.expectEqual(@as(u32, 0), p.reserved[0] + p.reserved[1]); + for (b.mapped(), 0..) |pg, k| try std.testing.expectEqual(k % 2, pg % 2); + // RAM resume: adopt full pages of the right residues, refuse the wrong ones + const c = try p.newSlot(8 * 256); + try std.testing.expectError(error.ResidueMismatch, c.adopt(&.{b.pages[1]})); + p.share(b.mapped()[0..2]); + try c.adopt(b.mapped()[0..2]); + try c.push(try p.takeAt(2)); + try std.testing.expectEqual(@as(u16, 3), p.holders(b.pages[0])); + try std.testing.expectEqual(@as(u32, 0), c.pages[2] % 2); +} diff --git a/zig/src/sessions/prefix.zig b/zig/src/sessions/prefix.zig new file mode 100644 index 0000000000..28c639fbd7 --- /dev/null +++ b/zig/src/sessions/prefix.zig @@ -0,0 +1,457 @@ +//! Session entries by prefix: a trie over chained page digests, one node per distinct run of full pages a tag. + +const std = @import("std"); + +pub const Digest = [16]u8; +pub const none = std.math.maxInt(u32); +const Blake = std.crypto.hash.blake2.Blake2b128; + +/// Chained digests of the full pages of ids under a tag: out[j] determines tag and ids[0 .. page (j + 1)]. +pub fn chain(tag: u32, ids: []const i32, page: u32, out: []Digest) void { + var prev = root(tag); + for (out, 0..) |*o, j| { + o.* = step(prev, ids[j * page ..][0..page]); + prev = o.*; + } +} + +/// The digest of a tag's empty prefix. +pub fn root(tag: u32) Digest { + var h = Blake.init(.{}); + h.update("tf-sess-root"); + h.update(std.mem.asBytes(&tag)); + return h.finalResult(); +} + +/// The next page's digest. +pub fn step(prev: Digest, page_ids: []const i32) Digest { + var h = Blake.init(.{}); + h.update(&prev); + h.update(std.mem.sliceAsBytes(page_ids)); + return h.finalResult(); +} + +/// An entry's identity: its node's digest (tag and full pages), its length and tail. +pub fn entryKey(node: Digest, pos: u64, tail: []const i32) Digest { + var h = Blake.init(.{}); + h.update("tf-sess-entry"); + h.update(&node); + h.update(std.mem.asBytes(&pos)); + h.update(std.mem.sliceAsBytes(tail)); + return h.finalResult(); +} + +pub const Tier = enum { ram, any }; + +pub fn Index(comptime Extra: type) type { + return struct { + const Self = @This(); + + pub const Node = struct { + digest: Digest, + parent: u32, + depth: u32, + first_child: u32 = none, + next_sib: u32 = none, + prev_sib: u32 = none, + first_entry: u32 = none, + /// entries in the subtree (this node's included): all, and those in RAM + sub_all: u32 = 0, + sub_ram: u32 = 0, + /// RAM entries at this node itself + own_ram: u32 = 0, + /// the tokens of the page leading into this node, while sub_ram > 0 + tokens: ?[]i32 = null, + }; + + pub const Entry = struct { + key: Digest, + tag: u32, + pos: u64, + node: u32, + tail: []i32, + ram: bool, + next_at: u32 = none, + prev_at: u32 = none, + live: bool = true, + extra: Extra, + }; + + /// An entry inserted, or the one already there with the same key. + pub const Placed = struct { id: u32, new: bool }; + + gpa: std.mem.Allocator, + page: u32, + nodes: std.ArrayList(Node) = .empty, + free_nodes: std.ArrayList(u32) = .empty, + by_digest: std.AutoHashMapUnmanaged(Digest, u32) = .empty, + entries: std.ArrayList(Entry) = .empty, + free_entries: std.ArrayList(u32) = .empty, + by_key: std.AutoHashMapUnmanaged(Digest, u32) = .empty, + /// page-sized token blocks, reused + blocks: std.ArrayList([]i32) = .empty, + + pub fn init(gpa: std.mem.Allocator, page: u32) Self { + return .{ .gpa = gpa, .page = page }; + } + + pub fn deinit(x: *Self) void { + for (x.nodes.items) |n| if (n.tokens) |t| x.gpa.free(t); + for (x.blocks.items) |b| x.gpa.free(b); + for (x.entries.items) |e| if (e.live) x.gpa.free(e.tail); + x.nodes.deinit(x.gpa); + x.free_nodes.deinit(x.gpa); + x.by_digest.deinit(x.gpa); + x.entries.deinit(x.gpa); + x.free_entries.deinit(x.gpa); + x.by_key.deinit(x.gpa); + x.blocks.deinit(x.gpa); + x.* = undefined; + } + + pub fn get(x: *Self, e: u32) *Entry { + return &x.entries.items[e]; + } + + pub fn node(x: *Self, n: u32) *Node { + return &x.nodes.items[n]; + } + + pub fn lookup(x: *const Self, key: Digest) ?u32 { + return x.by_key.get(key); + } + + fn newNode(x: *Self, digest: Digest, parent: u32, depth: u32) !u32 { + const id: u32 = if (x.free_nodes.pop()) |i| i else blk: { + try x.nodes.append(x.gpa, undefined); + break :blk @intCast(x.nodes.items.len - 1); + }; + x.nodes.items[id] = .{ .digest = digest, .parent = parent, .depth = depth }; + try x.by_digest.put(x.gpa, digest, id); + if (parent != none) { + const p = &x.nodes.items[parent]; + x.nodes.items[id].next_sib = p.first_child; + if (p.first_child != none) x.nodes.items[p.first_child].prev_sib = id; + p.first_child = id; + } + return id; + } + + fn takeBlock(x: *Self) ![]i32 { + if (x.blocks.pop()) |b| return b; + return x.gpa.alloc(i32, x.page); + } + + /// Adds an entry for ids (chain: its full pages' digests), or returns the one with the same key. + pub fn insert(x: *Self, tag: u32, ids: []const i32, digests: []const Digest, ram: bool, extra: Extra) !Placed { + const full: u32 = @intCast(ids.len / x.page); + std.debug.assert(digests.len >= full); + const r = root(tag); + var n = x.by_digest.get(r) orelse try x.newNode(r, none, 0); + for (0..full) |j| { + const d = digests[j]; + n = x.by_digest.get(d) orelse try x.newNode(d, n, @intCast(j + 1)); + } + const r2 = try x.place(tag, n, ids.len, ids[full * x.page ..], extra); + if (r2.new and ram) try x.setRam(r2.id, true, ids); + return r2; + } + + /// Adds a disk entry continuing from node anchor (none: the tag's root); suffix: tokens [depth x page, pos). + pub fn insertFrom(x: *Self, tag: u32, anchor: u32, suffix: []const i32, extra: Extra) !Placed { + var n = if (anchor != none) anchor else blk: { + const r = root(tag); + break :blk x.by_digest.get(r) orelse try x.newNode(r, none, 0); + }; + const d0 = x.nodes.items[n].depth; + const full: u32 = @intCast(suffix.len / x.page); + for (0..full) |j| { + const d = step(x.nodes.items[n].digest, suffix[j * x.page ..][0..x.page]); + n = x.by_digest.get(d) orelse try x.newNode(d, n, d0 + @as(u32, @intCast(j)) + 1); + } + return x.place(tag, n, @as(u64, d0) * x.page + suffix.len, suffix[full * x.page ..], extra); + } + + /// The node at depth on an entry's path (its first `depth` full pages). + pub fn nodeAt(x: *Self, id: u32, depth: u32) ?u32 { + var n = x.entries.items[id].node; + while (n != none and x.nodes.items[n].depth > depth) n = x.nodes.items[n].parent; + return if (n != none and x.nodes.items[n].depth == depth) n else null; + } + + fn place(x: *Self, tag: u32, n: u32, pos: u64, tail_ids: []const i32, extra: Extra) !Placed { + const key = entryKey(x.nodes.items[n].digest, pos, tail_ids); + if (x.by_key.get(key)) |e| return .{ .id = e, .new = false }; + const tail = try x.gpa.dupe(i32, tail_ids); + errdefer x.gpa.free(tail); + const id: u32 = if (x.free_entries.pop()) |i| i else blk: { + try x.entries.append(x.gpa, undefined); + break :blk @intCast(x.entries.items.len - 1); + }; + try x.by_key.put(x.gpa, key, id); + const nd = &x.nodes.items[n]; + x.entries.items[id] = .{ .key = key, .tag = tag, .pos = pos, .node = n, .tail = tail, .ram = false, .next_at = nd.first_entry, .extra = extra }; + if (nd.first_entry != none) x.entries.items[nd.first_entry].prev_at = id; + nd.first_entry = id; + x.addPath(n, 1, 0); + return .{ .id = id, .new = true }; + } + + fn addPath(x: *Self, from: u32, d_all: i32, d_ram: i32) void { + var n = from; + while (n != none) { + const nd = &x.nodes.items[n]; + nd.sub_all = @intCast(@as(i64, nd.sub_all) + d_all); + nd.sub_ram = @intCast(@as(i64, nd.sub_ram) + d_ram); + n = nd.parent; + } + } + + /// Moves an entry in or out of RAM; going in needs its ids (to keep its path's tokens). + pub fn setRam(x: *Self, id: u32, ram: bool, ids: ?[]const i32) !void { + const e = &x.entries.items[id]; + if (e.ram == ram) return; + if (ram) { + const all = ids orelse return error.IdsNeeded; + var n = e.node; + while (n != none and x.nodes.items[n].depth > 0) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + if (nd.tokens != null) continue; + const b = try x.takeBlock(); + @memcpy(b, all[(nd.depth - 1) * x.page ..][0..x.page]); + nd.tokens = b; + } + } + e.ram = ram; + x.nodes.items[e.node].own_ram = if (ram) x.nodes.items[e.node].own_ram + 1 else x.nodes.items[e.node].own_ram - 1; + x.addPath(e.node, 0, if (ram) 1 else -1); + if (!ram) x.dropTokens(e.node); + } + + fn dropTokens(x: *Self, from: u32) void { + var n = from; + while (n != none) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + if (nd.sub_ram > 0) return; + if (nd.tokens) |t| { + x.blocks.append(x.gpa, t) catch x.gpa.free(t); + nd.tokens = null; + } + } + } + + /// Removes an entry and the nodes no entry needs any more. + pub fn remove(x: *Self, id: u32) void { + x.setRam(id, false, null) catch unreachable; + const e = &x.entries.items[id]; + const n = e.node; + if (e.prev_at != none) x.entries.items[e.prev_at].next_at = e.next_at else x.nodes.items[n].first_entry = e.next_at; + if (e.next_at != none) x.entries.items[e.next_at].prev_at = e.prev_at; + _ = x.by_key.remove(e.key); + x.gpa.free(e.tail); + e.live = false; + x.free_entries.append(x.gpa, id) catch {}; + x.addPath(n, -1, 0); + var m = n; + while (m != none and x.nodes.items[m].sub_all == 0) { + const nd = x.nodes.items[m]; + if (nd.parent != none) { + if (nd.prev_sib != none) x.nodes.items[nd.prev_sib].next_sib = nd.next_sib else x.nodes.items[nd.parent].first_child = nd.next_sib; + if (nd.next_sib != none) x.nodes.items[nd.next_sib].prev_sib = nd.prev_sib; + } + _ = x.by_digest.remove(nd.digest); + x.free_nodes.append(x.gpa, m) catch {}; + m = nd.parent; + } + } + + /// The entry's tail matches the prompt there and the entry is shorter (its full pages matched by digest). + fn tailOf(e: *const Entry, prompt: []const i32) bool { + const at: usize = @intCast(e.pos - e.tail.len); + return e.pos < prompt.len and std.mem.eql(i32, e.tail, prompt[at..][0..e.tail.len]); + } + + /// The longest entry of tag whose ids strictly prefix prompt (RAM first on a tie); digests: the prompt's chain. + pub fn find(x: *Self, tag: u32, prompt: []const i32, digests: []const Digest, tier: Tier) ?u32 { + const full: u32 = @intCast(prompt.len / x.page); + const r = x.by_digest.get(root(tag)) orelse return null; + // the deepest node on the prompt's path: each prefix of an entry's path has a node: binary search + var lo: u32 = 0; + var hi: u32 = full; + while (lo < hi) { + const mid = lo + (hi - lo + 1) / 2; + if (x.by_digest.contains(digests[mid - 1])) lo = mid else hi = mid - 1; + } + var n: u32 = if (lo == 0) r else x.by_digest.get(digests[lo - 1]).?; + while (n != none) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + if (nd.sub_all == 0 or (tier == .ram and nd.own_ram == 0)) continue; + var best: u32 = none; + var e = nd.first_entry; + while (e != none) : (e = x.entries.items[e].next_at) { + const en = &x.entries.items[e]; + if (tier == .ram and !en.ram) continue; + if (!tailOf(en, prompt)) continue; + if (best == none or en.pos > x.entries.items[best].pos or (en.pos == x.entries.items[best].pos and en.ram)) best = e; + } + if (best != none) return best; + } + return null; + } + + /// A newer entry of id's chain holds every position it does (tier .ram: in RAM), so id need not be written. + pub fn covered(x: *Self, id: u32, tier: Tier) bool { + const e = &x.entries.items[id]; + const nd = &x.nodes.items[e.node]; + var f = nd.first_entry; + while (f != none) : (f = x.entries.items[f].next_at) { + const o = &x.entries.items[f]; + if (f == id or (tier == .ram and !o.ram)) continue; + if (o.pos > e.pos and std.mem.startsWith(i32, o.tail, e.tail)) return true; + } + var c = nd.first_child; + while (c != none) : (c = x.nodes.items[c].next_sib) { + const ch = &x.nodes.items[c]; + const n = if (tier == .ram) ch.sub_ram else ch.sub_all; + if (n == 0) continue; + if (e.tail.len == 0) return true; + if (ch.tokens) |t| if (std.mem.startsWith(i32, t, e.tail)) return true; + } + return false; + } + + /// How many RAM entries extend a RAM entry (longer, same tag, its ids their prefix): branch points have more. + pub fn extensions(x: *Self, id: u32) u32 { + const e = &x.entries.items[id]; + const nd = &x.nodes.items[e.node]; + var n: u32 = 0; + var f = nd.first_entry; + while (f != none) : (f = x.entries.items[f].next_at) { + const o = &x.entries.items[f]; + if (f != id and o.ram and o.pos > e.pos and std.mem.startsWith(i32, o.tail, e.tail)) n += 1; + } + var c = nd.first_child; + while (c != none) : (c = x.nodes.items[c].next_sib) { + const ch = &x.nodes.items[c]; + if (ch.sub_ram == 0) continue; + if (std.mem.startsWith(i32, ch.tokens.?, e.tail)) n += ch.sub_ram; + } + return n; + } + + /// A RAM entry no other RAM entry extends: a chat's newest state. + pub fn isLeaf(x: *Self, id: u32) bool { + return x.entries.items[id].ram and !x.covered(id, .ram); + } + + /// A RAM leaf's older turns (the RAM entries it extends) into out; returns the chat's recency (newest `used`). + pub fn members(x: *Self, leaf: u32, out: *std.ArrayList(u32), used: *const fn (*const Extra) u64) !u64 { + const l = &x.entries.items[leaf]; + var recency = used(&l.extra); + var below: u32 = none; // the path's node under n: its tokens are the page an entry at n must prefix + var n = l.node; + while (n != none) : ({ + below = n; + n = x.nodes.items[n].parent; + }) { + const nd = &x.nodes.items[n]; + if (nd.own_ram == 0) continue; + const page_ids: []const i32 = if (below == none) l.tail else x.nodes.items[below].tokens.?; + var f = nd.first_entry; + while (f != none) : (f = x.entries.items[f].next_at) { + const o = &x.entries.items[f]; + if (f == leaf or !o.ram or o.pos >= l.pos) continue; + if (!std.mem.startsWith(i32, page_ids, o.tail)) continue; + try out.append(x.gpa, f); + recency = @max(recency, used(&o.extra)); + } + } + return recency; + } + + /// A RAM entry's ids, rebuilt from its path's pages and tail, into out (len = pos). + pub fn idsOf(x: *Self, id: u32, out: []i32) !void { + const e = &x.entries.items[id]; + if (out.len != e.pos) return error.BadLength; + const at = e.pos - e.tail.len; + @memcpy(out[at..], e.tail); + var n = e.node; + while (n != none and x.nodes.items[n].depth > 0) : (n = x.nodes.items[n].parent) { + const nd = &x.nodes.items[n]; + const t = nd.tokens orelse return error.NotInRam; + @memcpy(out[(nd.depth - 1) * x.page ..][0..x.page], t); + } + } + }; +} + +const T = Index(u64); + +fn usedOf(e: *const u64) u64 { + return e.*; +} + +fn seq(gpa: std.mem.Allocator, n: usize, salt: i32) ![]i32 { + const v = try gpa.alloc(i32, n); + for (v, 0..) |*t, i| t.* = @as(i32, @intCast(i % 997)) * 3 + salt; + return v; +} + +fn add(x: *T, tag: u32, ids: []const i32, ram: bool, used: u64) !u32 { + var d: [64]Digest = undefined; + chain(tag, ids, x.page, d[0 .. ids.len / x.page]); + return (try x.insert(tag, ids, d[0 .. ids.len / x.page], ram, used)).id; +} + +test "longest strict prefix, ties, tags and coverage" { + const gpa = std.testing.allocator; + var x = T.init(gpa, 16); + defer x.deinit(); + const chat = try seq(gpa, 200, 0); + defer gpa.free(chat); + const t1 = try add(&x, 7, chat[0..40], true, 1); // a prompt's replay point + const t2 = try add(&x, 7, chat[0..72], true, 2); // a turn's end + const t3 = try add(&x, 7, chat[0..128], false, 3); // on disk only, page aligned + _ = try add(&x, 8, chat[0..100], true, 4); // another tag + var d: [16]Digest = undefined; + chain(7, chat[0..40], 16, d[0..2]); + const again = try x.insert(7, chat[0..40], d[0..2], true, 9); + try std.testing.expect(again.id == t1 and !again.new); + chain(7, chat[0..150], 16, d[0..9]); + try std.testing.expectEqual(@as(?u32, t3), x.find(7, chat[0..150], d[0..9], .any)); + try std.testing.expectEqual(@as(?u32, t2), x.find(7, chat[0..150], d[0..9], .ram)); + try std.testing.expectEqual(@as(?u32, t2), x.find(7, chat[0..128], d[0..8], .any)); // strict: 128 is not a prefix of 128 + chain(7, chat[0..72], 16, d[0..4]); + try std.testing.expectEqual(@as(?u32, t1), x.find(7, chat[0..72], d[0..4], .any)); + var other = try gpa.dupe(i32, chat[0..150]); + defer gpa.free(other); + other[50] = -1; // diverges inside t2's tail page: t1 still prefixes it + chain(7, other, 16, d[0..9]); + try std.testing.expectEqual(@as(?u32, t1), x.find(7, other, d[0..9], .any)); + // t3 extends t2 on disk only: its page's tokens are gone, so a tail there is undecided (the store marks covered_by) + try std.testing.expect(x.covered(t1, .ram) and !x.covered(t2, .any) and !x.covered(t2, .ram) and !x.covered(t3, .any)); + const t4 = try add(&x, 7, chat[0..64], false, 5); // page aligned: any deeper entry covers it + try std.testing.expect(x.covered(t4, .any) and x.covered(t4, .ram)); + x.remove(t4); + try std.testing.expect(x.isLeaf(t2) and !x.isLeaf(t1)); + try std.testing.expectEqual(@as(u32, 1), x.extensions(t1)); + try std.testing.expectEqual(@as(u32, 0), x.extensions(t2)); + var m: std.ArrayList(u32) = .empty; + defer m.deinit(gpa); + try std.testing.expectEqual(@as(u64, 2), try x.members(t2, &m, usedOf)); + try std.testing.expectEqualSlices(u32, &.{t1}, m.items); + const ids = try gpa.alloc(i32, 72); + defer gpa.free(ids); + try x.idsOf(t2, ids); + try std.testing.expectEqualSlices(i32, chat[0..72], ids); + // going to disk drops the path's tokens once no RAM entry needs them; removal prunes the nodes + try x.setRam(t2, false, null); + try x.setRam(t1, false, null); + try std.testing.expectError(error.NotInRam, x.idsOf(t2, ids)); + x.remove(t3); + x.remove(t2); + x.remove(t1); + chain(7, chat[0..150], 16, d[0..9]); + try std.testing.expectEqual(@as(?u32, null), x.find(7, chat[0..150], d[0..9], .any)); + try std.testing.expectEqual(@as(u32, 1 + 6), x.by_digest.count()); // tag 8's root + its 6 full pages +} diff --git a/zig/src/sessions/sessions.zig b/zig/src/sessions/sessions.zig new file mode 100644 index 0000000000..7bdc46b386 --- /dev/null +++ b/zig/src/sessions/sessions.zig @@ -0,0 +1,33 @@ +//! Session memory, family-neutral: the KV pool book, page bytes, prefix entries, the RAM tier and the NVMe tier. + +pub const freeset = @import("freeset.zig"); +pub const pool = @import("pool.zig"); +pub const Pool = pool.Pool; +pub const Slot = pool.Slot; +pub const Family = pool.Family; +pub const Layout = pool.Layout; +pub const pagestore = @import("pagestore.zig"); +pub const PageStore = pagestore.PageStore; +pub const HostPool = pagestore.HostPool; +pub const prefix = @import("prefix.zig"); +pub const lanes = @import("lanes.zig"); +pub const format = @import("format.zig"); +pub const disk = @import("disk.zig"); +pub const Disk = disk.Disk; +pub const jobs = @import("jobs.zig"); +pub const store = @import("store.zig"); +pub const Store = store.Store; + +test { + _ = freeset; + _ = pool; + _ = pagestore; + _ = prefix; + _ = lanes; + _ = format; + _ = disk; + _ = jobs; + _ = store; + _ = @import("disk_test.zig"); + _ = @import("store_test.zig"); +} diff --git a/zig/src/sessions/store.zig b/zig/src/sessions/store.zig new file mode 100644 index 0000000000..851ce509aa --- /dev/null +++ b/zig/src/sessions/store.zig @@ -0,0 +1,570 @@ +//! The session store: RAM entries holding pool pages and bounded state, a residency budget, and the NVMe tier. + +const std = @import("std"); +const Io = std.Io; +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Slot = pool_mod.Slot; +const PageStore = @import("pagestore.zig").PageStore; +const prefix = @import("prefix.zig"); +const Digest = prefix.Digest; +const disk_mod = @import("disk.zig"); +const Disk = disk_mod.Disk; +const jobs = @import("jobs.zig"); +const Job = jobs.Job; + +pub const none = prefix.none; + +/// A slot's bounded state (rings, carries, lookback, taps) as the family snapshots it: opaque to the store. +pub const Bounded = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + nbytes: *const fn (ptr: *anyopaque) u64, + /// host copies of its pieces for a park (owned by the caller: name and bytes allocated with gpa) + blobs: *const fn (ptr: *anyopaque, gpa: std.mem.Allocator) anyerror![]disk_mod.Blob, + release: *const fn (ptr: *anyopaque) void, + }; +}; + +pub const Policy = struct { + /// bounded state the RAM tier keeps (least recently used entries go past it) + ram_bytes: u64 = 512 << 20, + /// idle chats resident in the pool (0: no limit) and pool bytes held by entries alone (0: no limit) + chats: u32 = 0, + resident_bytes: u64 = 0, + /// park an evicted chat's older turns too (default: drop them, the leaf's file covers them) + park_prefixes: bool = false, + /// drop, never write, an entry a newer entry of its chain covers + skip_covered: bool = true, + /// write a newer turn as a delta on the file its leading pages came from + delta: bool = true, + /// parks and restores on the tier thread (settled by the plan) + async_io: bool = false, +}; + +pub const Extra = struct { + /// RAM: the pool pages it holds, logical [0, npages) + pages: []u32 = &.{}, + bounded: ?Bounded = null, + /// a file with this key is registered on the NVMe tier + file: bool = false, + used: u64 = 0, + origin: ?pool_mod.Origin = null, + /// a newer entry of its chain was parked from RAM: writing this one again is redundant + covered_by: bool = false, + kind: u32 = 0, + /// an unsettled park of it + job: ?*Job = null, + older: u32 = none, + newer: u32 = none, +}; + +const Index = prefix.Index(Extra); + +pub const Hit = struct { id: u32, ram: bool, pos: u64 }; + +pub const Stats = struct { + saves: u64 = 0, + dups: u64 = 0, + hits_ram: u64 = 0, + hits_disk: u64 = 0, + parks: u64 = 0, + delta_parks: u64 = 0, + drops: u64 = 0, + covered_drops: u64 = 0, + chat_evictions: u64 = 0, + park_ns: u64 = 0, + restore_ns: u64 = 0, +}; + +pub const Store = struct { + gpa: std.mem.Allocator, + io: Io, + pool: *Pool, + pages: PageStore, + index: Index, + disk: ?*Disk, + /// heap-held: the tier thread keeps its address + runner: ?*jobs.Runner = null, + policy: Policy, + clock: u64 = 0, + bounded_total: u64 = 0, + /// RAM entries, least recently used first + oldest: u32 = none, + newest: u32 = none, + ram_count: u32 = 0, + pending: std.ArrayList(*Job) = .empty, + scratch: std.ArrayList(Digest) = .empty, + stats: Stats = .{}, + + pub fn init(gpa: std.mem.Allocator, io: Io, pool: *Pool, pages: PageStore, d: ?*Disk, policy: Policy) !Store { + var s: Store = .{ .gpa = gpa, .io = io, .pool = pool, .pages = pages, .index = Index.init(gpa, pool.page), .disk = d, .policy = policy }; + if (d) |dd| { + const r = try gpa.create(jobs.Runner); + r.* = jobs.Runner.init(gpa, io, dd); + s.runner = r; + if (policy.async_io) r.startThread() catch |e| { + gpa.destroy(r); + return e; + }; + } + return s; + } + + pub fn deinit(s: *Store) void { + if (s.runner) |r| { + for (s.pending.items) |j| r.wait(j); + r.deinit(); + s.gpa.destroy(r); + } + for (s.pending.items) |j| j.deinit(s.gpa); + s.pending.deinit(s.gpa); + for (s.index.entries.items) |*e| if (e.live) { + if (e.extra.bounded) |b| b.vtable.release(b.ptr); + s.pool.drop(e.extra.pages) catch {}; + s.gpa.free(e.extra.pages); + }; + s.index.deinit(); + s.scratch.deinit(s.gpa); + s.* = undefined; + } + + fn digests(s: *Store, tag: u32, ids: []const i32) ![]const Digest { + const n = ids.len / s.pool.page; + try s.scratch.resize(s.gpa, n); + prefix.chain(tag, ids, s.pool.page, s.scratch.items); + return s.scratch.items; + } + + pub fn entry(s: *Store, id: u32) *Index.Entry { + return s.index.get(id); + } + + /// The longest stored entry (RAM or NVMe) whose ids strictly prefix prompt, with the same tag. + pub fn find(s: *Store, tag: u32, prompt: []const i32) !?Hit { + const d = try s.digests(tag, prompt); + const id = s.index.find(tag, prompt, d, .any) orelse return null; + const e = s.index.get(id); + return .{ .id = id, .ram = e.ram, .pos = e.pos }; + } + + fn link(s: *Store, id: u32) void { + const x = &s.index.get(id).extra; + x.older = s.newest; + x.newer = none; + if (s.newest != none) s.index.get(s.newest).extra.newer = id else s.oldest = id; + s.newest = id; + } + + fn unlink(s: *Store, id: u32) void { + const x = &s.index.get(id).extra; + if (x.older != none) s.index.get(x.older).extra.newer = x.newer else s.oldest = x.newer; + if (x.newer != none) s.index.get(x.newer).extra.older = x.older else s.newest = x.older; + } + + fn touch(s: *Store, id: u32) void { + s.clock += 1; + const e = s.index.get(id); + e.extra.used = s.clock; + if (e.ram) { + s.unlink(id); + s.link(id); + } + } + + /// An entry for a slot's state at len(ids), sharing its pages; null for a duplicate (bounded is released then). + pub fn save(s: *Store, slot: *Slot, bounded: Bounded, ids: []const i32, tag: u32, kind: u32) !?u32 { + if (ids.len == 0) return error.EmptyEntry; + const need = pool_mod.pagesFor(ids.len, s.pool.page); + if (slot.len < need) return error.SlotTooShort; + const d = try s.digests(tag, ids); + const ins = try s.index.insert(tag, ids, d, false, .{ .kind = kind }); + errdefer if (ins.new) s.index.remove(ins.id); + const e = s.index.get(ins.id); + if (e.ram) { + bounded.vtable.release(bounded.ptr); + s.touch(ins.id); + s.stats.dups += 1; + return null; + } + const pages = try s.gpa.dupe(u32, slot.mapped()[0..need]); + errdefer s.gpa.free(pages); + try s.index.setRam(ins.id, true, ids); + s.pool.share(pages); + const x = &s.index.get(ins.id).extra; + x.pages = pages; + x.bounded = bounded; + x.kind = kind; + if (slot.origin) |o| x.origin = .{ .key = o.key, .pages = @min(o.pages, @as(u32, @intCast(ids.len / s.pool.page))) }; + s.bounded_total += bounded.vtable.nbytes(bounded.ptr); + s.ram_count += 1; + s.link(ins.id); + s.touch(ins.id); + s.stats.saves += 1; + try s.trimRam(); + _ = try s.enforce(); + return ins.id; + } + + /// A RAM entry into an empty slot: full pages shared, the partial last page copied; returns its bounded state. + pub fn restoreRam(s: *Store, id: u32, slot: *Slot) !Bounded { + const e = s.index.get(id); + if (!e.ram) return error.NotInRam; + const full: u32 = @intCast(e.pos / s.pool.page); + try slot.adopt(e.extra.pages[0..full]); + errdefer _ = slot.releaseAll() catch {}; + if (e.pos % s.pool.page != 0) { + const src = e.extra.pages[full]; + const dst = try s.pool.takeAt(full); + for (s.pool.layout.families, 0..) |f, fi| { + if (f.split and !s.pool.owns(src)) continue; + try s.pages.copyPage(@intCast(fi), s.pool.familyPage(f, src), s.pool.familyPage(f, dst)); + } + try slot.push(dst); + } + if (e.extra.origin) |o| slot.origin = .{ .key = o.key, .pages = @min(o.pages, full) }; + s.touch(id); + s.stats.hits_ram += 1; + return e.extra.bounded.?; + } + + /// Starts an NVMe resume into an empty slot (pages mapped now, filled by the job); settle it before the slot runs. + pub fn beginRestore(s: *Store, id: u32, slot: *Slot) !*Job { + const d = s.disk orelse return error.NoDisk; + const e = s.index.get(id); + if (!e.extra.file) return error.NotOnDisk; + if (slot.len != 0) return error.SlotNotEmpty; + try slot.ensure(e.pos); + const j = try s.gpa.create(Job); + j.* = .{ .kind = .restore, .entry = id, .slot = slot.id }; + errdefer j.deinit(s.gpa); + var ext: std.ArrayList(disk_mod.Extent) = .empty; + errdefer ext.deinit(s.gpa); + try d.extents(e.key, &ext); + j.extents = try ext.toOwnedSlice(s.gpa); + const np = pool_mod.pagesFor(e.pos, s.pool.page); + j.lists = try s.familyLists(slot.mapped()[0..np], 0, false); + j.sink = .{ .store = s.pages, .families = s.pool.layout.families, .page = s.pool.page, .world = s.pool.world, .rank = s.pool.rank, .pages = j.lists }; + try s.pending.append(s.gpa, j); + try s.runner.?.submit(j); + return j; + } + + /// Per family: family-tensor pages of logical [first, len(pages)) (owned_only: this rank's of split families). + fn familyLists(s: *Store, pages: []const u32, first: u32, owned_only: bool) ![][]u32 { + const fams = s.pool.layout.families; + const out = try s.gpa.alloc([]u32, fams.len); + var made: usize = 0; + errdefer { + for (out[0..made]) |l| s.gpa.free(l); + s.gpa.free(out); + } + for (fams, 0..) |f, fi| { + const w = s.pool.world; + const sharded = f.split and w > 1 and owned_only; + var n: usize = pages.len - first; + if (sharded) { + var by: [pool_mod.max_world]u32 = undefined; + pool_mod.span(first, @intCast(pages.len), w, &by); + n = by[s.pool.rank]; + } + const l = try s.gpa.alloc(u32, n); + var i: usize = 0; + var k: usize = first; + if (sharded) k = first + (s.pool.rank + w - first % w) % w; + while (k < pages.len) : (k += if (sharded) w else 1) { + l[i] = s.pool.familyPage(f, pages[k]); + i += 1; + } + out[fi] = l; + made += 1; + } + return out; + } + + /// Starts a park of a RAM entry (null: nothing to write, the entry is already settled out of RAM). + pub fn beginPark(s: *Store, id: u32) !?*Job { + const d = s.disk orelse return error.NoDisk; + const e = s.index.get(id); + if (!e.ram or e.extra.job != null) return error.NotParkable; + if (e.extra.file and d.get(e.key) != null) { + try s.leaveRam(id); + return null; + } + const page = s.pool.page; + const npages = pool_mod.pagesFor(e.pos, page); + var base: ?Digest = null; + var first: u32 = 0; + if (s.policy.delta) if (e.extra.origin) |o| if (o.pages > 0) { + var r = d.get(o.key); + while (r) |rec| : (r = if (rec.base == disk_mod.none) null else &d.recs.items[rec.base]) { + if (rec.first_page < o.pages) break; + } + if (r) |rec| if (rec.depth < d.opts.max_chain) { + base = rec.key; + first = o.pages; + }; + }; + const j = try s.gpa.create(Job); + j.* = .{ .kind = .park, .entry = id }; + errdefer j.deinit(s.gpa); + j.ids = try s.gpa.alloc(i32, e.pos); + try s.index.idsOf(id, j.ids); + j.lists = try s.familyLists(e.extra.pages, first, true); + j.blobs = try e.extra.bounded.?.vtable.blobs(e.extra.bounded.?.ptr, s.gpa); + j.src = .{ .key = e.key, .base = base, .tag = e.tag, .kind = e.extra.kind, .page = page, .pos = e.pos, .first_page = first, .npages = npages, .ids = j.ids[first * page ..], .store = s.pages, .fence = try s.pages.fence(), .families = s.pool.layout.families, .pages = j.lists, .blobs = j.blobs }; + if (base) |b| d.get(b).?.refs += 1; // held while the job runs: a trim cannot take the base + e.extra.job = j; + try s.pending.append(s.gpa, j); + try s.runner.?.submit(j); + return j; + } + + /// Settles a finished job at the plan's point: a park registers its file, a restore returns the bounded state. + pub fn settle(s: *Store, j: *Job) !?disk_mod.Restored { + s.runner.?.wait(j); + defer { + for (s.pending.items, 0..) |p, i| if (p == j) { + _ = s.pending.orderedRemove(i); + break; + }; + j.deinit(s.gpa); + } + const d = s.disk.?; + switch (j.kind) { + .park => { + if (j.src.base) |b| d.get(b).?.refs -= 1; + const e = s.index.get(j.entry); + e.extra.job = null; + if (j.err) |err| return err; // the entry stays in RAM; the .tmp is gone + _ = try d.register(&j.src, j.size); + e.extra.file = true; + s.stats.parks += 1; + if (j.src.base != null) s.stats.delta_parks += 1; + s.stats.park_ns += j.elapsed_ns; + if (s.policy.skip_covered) { + var m: std.ArrayList(u32) = .empty; + defer m.deinit(s.gpa); + _ = try s.index.members(j.entry, &m, usedOf); + for (m.items) |o| s.index.get(o).extra.covered_by = true; + } + try s.leaveRam(j.entry); + try s.trimDisk(); + return null; + }, + .restore => { + if (j.err) |err| { + // the slot as before the restore: empty, its reservation kept + _ = try s.pool.slots.items[j.slot].truncate(0); + return err; + } + const e = s.index.get(j.entry); + d.touch(e.key); + s.touch(j.entry); + s.pool.slots.items[j.slot].origin = .{ .key = e.key, .pages = @intCast(e.pos / s.pool.page) }; + s.stats.hits_disk += 1; + s.stats.restore_ns += j.elapsed_ns; + const r = j.restored; + j.restored = null; + return r; + }, + } + } + + /// Settles every pending job in submission order (one rank, or the plan's "settle all"). + pub fn settleAll(s: *Store) !void { + while (s.pending.items.len > 0) { + if (try s.settle(s.pending.items[0])) |r| { + var rr = r; + rr.deinit(s.gpa); + } + } + } + + fn usedOf(x: *const Extra) u64 { + return x.used; + } + + /// A RAM entry lets go of its pages and bounded state; it stays indexed if a file holds it. + fn leaveRam(s: *Store, id: u32) !void { + const e = s.index.get(id); + if (!e.ram) return; + try s.pool.drop(e.extra.pages); + s.gpa.free(e.extra.pages); + e.extra.pages = &.{}; + if (e.extra.bounded) |b| { + s.bounded_total -= b.vtable.nbytes(b.ptr); + b.vtable.release(b.ptr); + } + e.extra.bounded = null; + s.unlink(id); + s.ram_count -= 1; + try s.index.setRam(id, false, null); + if (!e.extra.file) s.index.remove(id); + } + + /// An entry whose file could not be read back: the file deleted, the entry gone unless RAM holds it (every rank). + pub fn forget(s: *Store, id: u32) !void { + const e = s.index.get(id); + if (e.extra.job != null) return error.Busy; + if (e.extra.file) if (s.disk) |d| { + d.stats.bad += 1; + // a delta's base stays while a newer file stacks on it (an opt-in delta mode) + d.remove(e.key) catch |err| return if (err == error.HasDependents) {} else err; + e.extra.file = false; + }; + if (!e.ram) s.index.remove(id); + } + + /// Drops an entry from RAM without writing it. + pub fn drop(s: *Store, id: u32) !void { + try s.leaveRam(id); + s.stats.drops += 1; + } + + /// Lets go of a RAM entry: dropped when covered, parked if the tier takes it, else dropped; returns a started park. + pub fn evict(s: *Store, id: u32) !?*Job { + const e = s.index.get(id); + if (e.extra.job != null) return null; + if (s.policy.skip_covered and (e.extra.covered_by or s.index.covered(id, .any))) { + try s.drop(id); + s.stats.covered_drops += 1; + return null; + } + if (s.disk) |d| if (d.accepts(e.pos)) { + const j = try s.beginPark(id) orelse return null; + if (!s.policy.async_io) _ = try s.settle(j); + return if (s.policy.async_io) j else null; + }; + try s.drop(id); + return null; + } + + /// Evicts least recently used RAM entries past the bounded-state budget. + fn trimRam(s: *Store) !void { + var id = s.oldest; + while (s.bounded_total > s.policy.ram_bytes and s.ram_count > 1 and id != none) { + const next = s.index.get(id).extra.newer; + if (s.index.get(id).extra.job == null) _ = try s.evict(id); + id = next; + } + } + + /// Files the NVMe budget deleted leave the index (a RAM entry with that key stays, unwritten again). + fn trimDisk(s: *Store) !void { + const d = s.disk orelse return; + var gone: std.ArrayList(Digest) = .empty; + defer gone.deinit(s.gpa); + try d.trim(&gone); + for (gone.items) |k| if (s.index.lookup(k)) |id| { + s.index.get(id).extra.file = false; + if (!s.index.get(id).ram) s.index.remove(id); + }; + } + + const Chat = struct { leaf: u32, recency: u64 }; + + fn chatLess(_: void, a: Chat, b: Chat) bool { + return a.recency < b.recency or (a.recency == b.recency and a.leaf < b.leaf); + } + + /// A chat is live while a slot maps its leaf's first page. + fn idleLeaf(s: *Store, id: u32) bool { + const e = s.index.get(id); + return e.extra.job == null and s.index.isLeaf(id) and (e.extra.pages.len == 0 or s.pool.slot_refs[e.extra.pages[0]] == 0); + } + + /// Evicts LRU idle chats past the residency budget (leaf parked, unshared older turns dropped); returns how many. + pub fn enforce(s: *Store) !u32 { + const pol = s.policy; + if (pol.chats == 0 and pol.resident_bytes == 0) return 0; + var chats: std.ArrayList(Chat) = .empty; + defer chats.deinit(s.gpa); + var members: std.ArrayList(u32) = .empty; + defer members.deinit(s.gpa); + var id = s.oldest; + while (id != none) : (id = s.index.get(id).extra.newer) { + if (!s.idleLeaf(id)) continue; + members.clearRetainingCapacity(); + try chats.append(s.gpa, .{ .leaf = id, .recency = try s.index.members(id, &members, usedOf) }); + } + std.mem.sort(Chat, chats.items, {}, chatLess); + var idle: u32 = @intCast(chats.items.len); + var gone: u32 = 0; + for (chats.items) |c| { + const over_n = pol.chats > 0 and idle > pol.chats; + const over_b = pol.resident_bytes > 0 and s.pool.entryOnlyBytes() > pol.resident_bytes; + if (!over_n and !over_b) break; + if (!s.index.get(c.leaf).ram or !s.idleLeaf(c.leaf)) continue; + members.clearRetainingCapacity(); + _ = try s.index.members(c.leaf, &members, usedOf); + _ = try s.evict(c.leaf); + // deepest first: a member more RAM entries extend than this chat's leaf is a branch point, as is all above + const own: u32 = if (s.index.get(c.leaf).ram) 1 else 0; + for (members.items) |m| { + if (s.index.get(m).extra.job != null) continue; + if (s.index.extensions(m) > own) break; + if (pol.park_prefixes) _ = try s.evict(m) else try s.drop(m); + } + idle -= 1; + gone += 1; + s.stats.chat_evictions += 1; + } + return gone; + } + + /// RAM entries to let go of (coldest first, never keep) so need pages are free: [] if they are, null if impossible. + pub fn planSpill(s: *Store, need: u32, keep: []const u32, out: *std.ArrayList(u32)) !bool { + out.clearRetainingCapacity(); + var avail = s.pool.available(); + if (avail >= need) return true; + var left: std.AutoHashMapUnmanaged(u32, u16) = .empty; + defer left.deinit(s.gpa); + var id = s.oldest; + while (id != none and avail < need) : (id = s.index.get(id).extra.newer) { + if (std.mem.indexOfScalar(u32, keep, id) != null or s.index.get(id).extra.job != null) continue; + try out.append(s.gpa, id); + for (s.index.get(id).extra.pages) |pg| { + const g = try left.getOrPut(s.gpa, pg); + if (!g.found_existing) g.value_ptr.* = s.pool.refs[pg]; + g.value_ptr.* -= 1; + if (g.value_ptr.* == 0) avail += 1; // split: a page of the scarcer residue may count for less; the check below decides + } + } + return avail >= need; + } + + /// Rebuilds NVMe entries on restart (ranks then intersect them), deltas on their bases' paths; returns how many. + pub fn reconcile(s: *Store) !u32 { + const d = s.disk orelse return 0; + var found: std.ArrayList(Disk.Found) = .empty; + defer { + for (found.items) |f| s.gpa.free(f.ids); + found.deinit(s.gpa); + } + try d.scan(&found); + for (found.items) |f| { + var anchor: u32 = none; + if (f.head.hasBase()) { + const b = s.index.lookup(f.head.base) orelse return error.NoBase; + anchor = s.index.nodeAt(b, f.head.first_page) orelse return error.BadChain; + } + const r = try s.index.insertFrom(f.head.tag, anchor, f.ids, .{ .file = true, .kind = f.head.kind }); + s.clock += 1; + s.index.get(r.id).extra.used = s.clock; + } + return @intCast(found.items.len); + } + + /// Idle resident chats and the pool bytes held by entries alone. + pub fn resident(s: *Store) struct { chats: u32, bytes: u64 } { + var n: u32 = 0; + var id = s.oldest; + while (id != none) : (id = s.index.get(id).extra.newer) if (s.idleLeaf(id)) { + n += 1; + }; + return .{ .chats = n, .bytes = s.pool.entryOnlyBytes() }; + } +}; diff --git a/zig/src/sessions/store_test.zig b/zig/src/sessions/store_test.zig new file mode 100644 index 0000000000..86b58a4f59 --- /dev/null +++ b/zig/src/sessions/store_test.zig @@ -0,0 +1,216 @@ +//! The session store on a host pool and a real NVMe tier: chats, evictions, resumes bit for bit, deltas, two ranks. + +const std = @import("std"); +const testing = std.testing; +const pool_mod = @import("pool.zig"); +const Pool = pool_mod.Pool; +const Slot = pool_mod.Slot; +const Family = pool_mod.Family; +const HostPool = @import("pagestore.zig").HostPool; +const disk_mod = @import("disk.zig"); +const store_mod = @import("store.zig"); +const Store = store_mod.Store; +const Bounded = store_mod.Bounded; + +const page = 64; +const fams = [_]Family{ + .{ .name = "comp.2", .ratio = 2, .row_bytes = 584, .split = true }, + .{ .name = "index_k.2", .ratio = 2, .row_bytes = 132 }, + .{ .name = "comp.20", .ratio = 1, .row_bytes = 584, .split = true }, +}; + +/// Bounded state in host memory: one blob naming its chat and length. +const HostBounded = struct { + gpa: std.mem.Allocator, + bytes: []u8, + + fn make(gpa: std.mem.Allocator, chat: u32, pos: u64) !Bounded { + const h = try gpa.create(HostBounded); + h.* = .{ .gpa = gpa, .bytes = try gpa.alloc(u8, 3000) }; + for (h.bytes, 0..) |*b, i| b.* = @truncate(i +% chat *% 17 +% pos); + return .{ .ptr = h, .vtable = &vt }; + } + const vt: Bounded.VTable = .{ .nbytes = nbytes, .blobs = blobs, .release = release }; + fn nbytes(p: *anyopaque) u64 { + const h: *HostBounded = @ptrCast(@alignCast(p)); + return h.bytes.len; + } + fn blobs(p: *anyopaque, gpa: std.mem.Allocator) anyerror![]disk_mod.Blob { + const h: *HostBounded = @ptrCast(@alignCast(p)); + const out = try gpa.alloc(disk_mod.Blob, 1); + out[0] = .{ .name = try gpa.dupe(u8, "rings"), .bytes = try gpa.dupe(u8, h.bytes) }; + return out; + } + fn release(p: *anyopaque) void { + const h: *HostBounded = @ptrCast(@alignCast(p)); + h.gpa.free(h.bytes); + h.gpa.destroy(h); + } +}; + +/// One rank: its pool, host pages and store, plus what each chat's slot wrote (the reference rows by logical page). +const Rank = struct { + pool: Pool, + host: HostPool, + disk: *disk_mod.Disk, + store: Store, + + fn init(r: *Rank, gpa: std.mem.Allocator, root: []const u8, world: u32, rank: u32, policy: store_mod.Policy) !void { + r.pool = try Pool.init(gpa, .{ .families = &fams, .page = page }, 96 * page, world, rank); + r.host = try HostPool.init(gpa, &r.pool); + r.disk = try disk_mod.Disk.open(gpa, testing.io, .{ .root = root, .compat = @splat(7), .rank = rank, .world = world, .chunk = 64 << 10, .staging = 4, .lanes = 2, .min_tokens = 1 }); + r.store = try Store.init(gpa, testing.io, &r.pool, r.host.store(), r.disk, policy); + } + + fn deinit(r: *Rank) void { + r.store.deinit(); + r.disk.close(); + r.host.deinit(); + r.pool.deinit(); + } + + /// The rows a chat's slot writes at a logical page: a function of (chat, logical page, family) alone. + fn write(r: *Rank, s: *const Slot, chat: u32, from_page: u32) void { + for (s.mapped()[from_page..], from_page..) |pg, k| for (fams, 0..) |f, fi| { + if (f.split and !r.pool.owns(pg)) continue; + const v = r.host.pageOf(@intCast(fi), r.pool.familyPage(f, pg)); + for (v, 0..) |*b, i| b.* = @truncate(i *% 13 +% k *% 101 +% chat *% 7 +% fi); + }; + } + + /// Every stored row of the slot's pages below pos equals what write() put there. + fn check(r: *Rank, s: *const Slot, chat: u32, pos: u64) !void { + const n = pool_mod.pagesFor(pos, page); + for (s.mapped()[0..n], 0..) |pg, k| for (fams, 0..) |f, fi| { + if (f.split and !r.pool.owns(pg)) continue; + const v = r.host.pageOf(@intCast(fi), r.pool.familyPage(f, pg)); + const rows: usize = if (k + 1 == n and pos % page != 0) (pos % page) / f.ratio * f.row_bytes else v.len; + for (v[0..rows], 0..) |b, i| try testing.expectEqual(@as(u8, @truncate(i *% 13 +% k *% 101 +% chat *% 7 +% fi)), b); + }; + } +}; + +fn chatIds(gpa: std.mem.Allocator, chat: u32, n: usize) ![]i32 { + const v = try gpa.alloc(i32, n); + for (v, 0..) |*t, i| t.* = @intCast(chat * 100_000 + i); + return v; +} + +/// Prefill or resume a chat to len(ids) in slot s, save it, free the slot, enforce residency; returns the resume point. +fn turn(gpa: std.mem.Allocator, r: *Rank, s: *Slot, chat: u32, ids: []const i32) ![]const u8 { + var tier: []const u8 = "fresh"; + var from: u32 = 0; + if (try r.store.find(1, ids)) |hit| { + if (hit.ram) { + const b = try r.store.restoreRam(hit.id, s); + try testing.expectEqual(@as(u64, 3000), b.vtable.nbytes(b.ptr)); + tier = "ram"; + } else { + const j = try r.store.beginRestore(hit.id, s); + var got = (try r.store.settle(j)).?; + defer got.deinit(gpa); + try testing.expectEqual(@as(u8, @truncate(chat *% 17 +% hit.pos)), got.blobs[0].bytes[0]); + tier = "disk"; + } + try r.check(s, chat, hit.pos); + from = @intCast(hit.pos / page); + } + try s.ensure(ids.len); + r.write(s, chat, from); + _ = try r.store.save(s, try HostBounded.make(gpa, chat, ids.len), ids, 1, 0); + _ = try s.releaseAll(); + _ = try r.store.enforce(); + try r.store.settleAll(); // the plan's settle point (async parks finish here) + return tier; +} + +/// The scenario on one rank; returns a trace of its decisions (stats and page ids) to compare ranks. +fn scenario(gpa: std.mem.Allocator, root: []const u8, world: u32, rank: u32, async_io: bool, trace: *std.ArrayList(u64)) !void { + var r: Rank = undefined; + try r.init(gpa, root, world, rank, .{ .chats = 1, .async_io = async_io }); + defer r.deinit(); + const s = try r.pool.newSlot(64 * page); + const a = try chatIds(gpa, 1, 20 * page); + defer gpa.free(a); + const b = try chatIds(gpa, 2, 20 * page); + defer gpa.free(b); + const st = &r.store.stats; + try testing.expectEqualStrings("fresh", try turn(gpa, &r, s, 1, a[0 .. 5 * page + 10])); + try testing.expectEqualStrings("ram", try turn(gpa, &r, s, 1, a[0 .. 8 * page + 3])); + try testing.expectEqual(@as(u64, 0), st.parks); + // a second idle chat: chat 1 (least recent) parked whole, its older turn dropped (not written) + try testing.expectEqualStrings("fresh", try turn(gpa, &r, s, 2, b[0 .. 6 * page])); + try testing.expectEqual(@as(u64, 1), st.parks); + try testing.expectEqual(@as(u64, 1), st.drops); + try testing.expectEqual(@as(u32, 1), r.store.resident().chats); + const whole = r.disk.used; + // chat 1 resumes from NVMe bit for bit; chat 2 goes + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 1, a[0 .. 12 * page + 40])); + try testing.expectEqual(@as(u64, 2), st.parks); + const before = r.disk.used; + // chat 2 resumes; chat 1 goes again, written as a delta on its first file (pages 8..12 of 13) + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 2, b[0 .. 9 * page])); + try testing.expectEqual(@as(u64, 3), st.parks); + try testing.expectEqual(@as(u64, 1), st.delta_parks); + // the delta holds 5 of the entry's 13 pages: less than chat 1's first whole file of 9 pages + try testing.expect(r.disk.used - before < whole); + // chat 1's newest turn restores through its two files, bit for bit + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 1, a[0 .. 15 * page])); + try testing.expectEqual(@as(u64, 3), st.hits_disk); + try trace.appendSlice(gpa, &.{ st.parks, st.delta_parks, st.drops, st.covered_drops, st.hits_ram, st.hits_disk, r.pool.freePages(), r.pool.entry_only }); + var id = r.store.oldest; + while (id != store_mod.none) : (id = r.store.entry(id).extra.newer) for (r.store.entry(id).extra.pages) |pg| try trace.append(gpa, pg); +} + +/// After the scenario: a new process on the same tier finds every file and resumes chat 1 through its delta chain. +fn restart(gpa: std.mem.Allocator, root: []const u8, world: u32, rank: u32) !void { + var r: Rank = undefined; + try r.init(gpa, root, world, rank, .{ .chats = 1 }); + defer r.deinit(); + try testing.expectEqual(@as(u32, 4), try r.store.reconcile()); + const s = try r.pool.newSlot(64 * page); + const a = try chatIds(gpa, 1, 20 * page); + defer gpa.free(a); + try testing.expectEqualStrings("disk", try turn(gpa, &r, s, 1, a[0 .. 16 * page])); + try testing.expectEqual(@as(u64, 1), r.store.stats.hits_disk); +} + +fn tmpRoot(gpa: std.mem.Allocator, tmp: *std.testing.TmpDir, sub: []const u8) ![:0]u8 { + try tmp.dir.createDirPath(testing.io, sub); + return tmp.dir.realPathFileAlloc(testing.io, sub, gpa); +} + +test "chats over turns: residency to NVMe, bit-identical resumes, deltas, older turns dropped (sync and async)" { + const gpa = testing.allocator; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + var t0: std.ArrayList(u64) = .empty; + defer t0.deinit(gpa); + var t1: std.ArrayList(u64) = .empty; + defer t1.deinit(gpa); + const r0 = try tmpRoot(gpa, &tmp, "sync"); + defer gpa.free(r0); + try scenario(gpa, r0, 1, 0, false, &t0); + try restart(gpa, r0, 1, 0); + const r1 = try tmpRoot(gpa, &tmp, "async"); + defer gpa.free(r1); + try scenario(gpa, r1, 1, 0, true, &t1); + try testing.expectEqualSlices(u64, t0.items, t1.items); // the tier thread changes when bytes move, not what is decided +} + +test "two split ranks make the same decisions on the same calls and each restores its own rows" { + const gpa = testing.allocator; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + var t0: std.ArrayList(u64) = .empty; + defer t0.deinit(gpa); + var t1: std.ArrayList(u64) = .empty; + defer t1.deinit(gpa); + const root = try tmpRoot(gpa, &tmp, "tier"); + defer gpa.free(root); + try scenario(gpa, root, 2, 0, false, &t0); + try scenario(gpa, root, 2, 1, true, &t1); + try restart(gpa, root, 2, 0); + try restart(gpa, root, 2, 1); + try testing.expectEqualSlices(u64, t0.items, t1.items); +} From a5b8b3570ad444b3859d91c82d9e384abd6aa595 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 14:53:36 +0700 Subject: [PATCH 24/43] Server: a stream that has not sent its first token gets keepalives every TF_SSE_KEEPALIVE_S seconds (default 15). Queueing and a long prefill send nothing after the head, and clients or proxies that time out a silent body leave and retry from zero. Chat, completions and responses streams get an SSE comment (": keepalive"); Anthropic streams get its ping event. A keepalive that cannot be written cancels the request as a client that left. --- zig/src/server/anthropic.zig | 8 +++- zig/src/server/chat_generation.zig | 13 ++++++ zig/src/server/drain_test.zig | 29 ++++++++---- zig/src/server/keepalive_test.zig | 75 ++++++++++++++++++++++++++++++ zig/src/server/openai.zig | 21 ++++++++- zig/src/server/openai_run.zig | 5 +- zig/src/server/responses.zig | 7 ++- zig/src/server/root.zig | 1 + zig/src/server/serve.zig | 7 ++- zig/src/server/server.zig | 2 + zig/src/server/sse.zig | 5 ++ 11 files changed, 158 insertions(+), 15 deletions(-) create mode 100644 zig/src/server/keepalive_test.zig diff --git a/zig/src/server/anthropic.zig b/zig/src/server/anthropic.zig index a8ec8132f5..f1d8e0ec23 100644 --- a/zig/src/server/anthropic.zig +++ b/zig/src/server/anthropic.zig @@ -105,7 +105,7 @@ const Wire = struct { payload: Value = .null, fn out(w: *Wire) openai.Out { - return .{ .ctx = w, .vt = &.{ .open = open, .event = event, .reply = whole } }; + return .{ .ctx = w, .vt = &.{ .open = open, .event = event, .reply = whole, .keepalive = ping } }; } fn open(ctx: *anyopaque, prompt_len: usize) error{Closed}!void { @@ -130,6 +130,12 @@ const Wire = struct { w.status = status; w.payload = payload; } + + /// Anthropic's own keepalive event. + fn ping(ctx: *anyopaque) error{Closed}!void { + const w: *Wire = @ptrCast(@alignCast(ctx)); + return w.reply.send("ping", try w.reply.obj("ping")); + } }; const Call = struct { id: []const u8 = "", name: []const u8 = "", arguments: []const u8 = "", sent: usize = 0, done: bool = false }; diff --git a/zig/src/server/chat_generation.zig b/zig/src/server/chat_generation.zig index 9c49e3f14a..0b4e164c98 100644 --- a/zig/src/server/chat_generation.zig +++ b/zig/src/server/chat_generation.zig @@ -75,6 +75,8 @@ pub const Generation = struct { streaming_done: bool = false, first_ns: ?i96 = null, last_ns: ?i96 = null, + /// when the stream last carried bytes (a keepalive goes out after a silence) + wrote_ns: ?i96 = null, reason: ?[]const u8 = null, // the server ended the reply (a stop string, the length) before the engine said so engine_done: bool = false, consumed: usize = 0, @@ -166,6 +168,7 @@ pub const Generation = struct { g.cancel(); return error.Cancelled; }; + g.wrote_ns = nowNs(g.srv.io); } pub const cancel = chat.cancel; // in chat.zig with its doc @@ -206,6 +209,16 @@ pub const Generation = struct { g.cancel(); return error.Cancelled; } + if (chunk == null and @hasDecl(@TypeOf(gone), "beat")) { + // queueing and a long prefill send nothing: a client that times out a silent stream would leave + const now = nowNs(g.srv.io); + const since = g.wrote_ns orelse now; + if (g.wrote_ns == null) g.wrote_ns = now; + if (gone.beat(now - since) catch { + g.cancel(); + return error.Cancelled; + }) g.wrote_ns = now; + } const tokens = chunk orelse continue; if (g.reason == null) try g.commit(tokens); } diff --git a/zig/src/server/drain_test.zig b/zig/src/server/drain_test.zig index 2bdc9f4a59..411b651cf3 100644 --- a/zig/src/server/drain_test.zig +++ b/zig/src/server/drain_test.zig @@ -10,7 +10,7 @@ const Allocator = std.mem.Allocator; const testing = std.testing; /// Slow enough to stop mid-reply: "@short" ends after a dozen tokens, "@long" only when halted or cancelled. -const Slow = struct { +pub const Slow = struct { gpa: Allocator, io: std.Io, halted: std.atomic.Value(bool) = .init(false), @@ -65,6 +65,15 @@ const Slow = struct { for (r.prompt, prompt) |t, *b| b.* = @truncate(t); const long = std.mem.indexOf(u8, prompt, "@long") != null; const word: []const u32 = if (long) &.{ ' ', 'l', 'a' } else &.{ ' ', 'o', 'k' }; + const hang = std.mem.indexOf(u8, prompt, "@hang") != null; + const wait = std.mem.indexOf(u8, prompt, "@wait") != null; + // a long prefill: nothing for 800 ms ("@wait"), or nothing until stopped ("@hang") + var quiet: u32 = 0; + while (hang or (wait and quiet < 80)) : (quiet += 1) { + if (s.halted.load(.acquire)) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .failed, .message = s.reason } }); + if (s.cancelled.load(.acquire) == id) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .cancelled } }); + std.Io.sleep(s.io, .fromMilliseconds(10), .awake) catch {}; + } sink.event(sink.ctx, id, &.{ .prefilled = 0 }); var n: usize = 0; while (true) : (n += 1) { @@ -93,7 +102,7 @@ const Slow = struct { }; /// Bytes as tokens; a prompt renders as its messages' JSON, so the engine sees the client's marker. -const TestText = struct { +pub const TestText = struct { fn text(t: *@This()) model_text.Text { return .{ .ctx = t, .vtable = &.{ .encode = encode, .decode = decode, .token_id = tokenId, .token_string = tokenString, .vocab_size = vocabSize, .eos_ids = eosIds, .render = render, .template_source = templateSource } }; } @@ -135,16 +144,16 @@ const TestText = struct { } }; -var listen_port: std.atomic.Value(u16) = .init(0); +pub var listen_port: std.atomic.Value(u16) = .init(0); fn onListen(port: u16) void { listen_port.store(port, .release); } -const Reply = struct { status: u16, head: []const u8, body: []const u8 }; +pub const Reply = struct { status: u16, head: []const u8, body: []const u8 }; /// One request on its own connection (closed by us after it), the whole response read. -fn request(a: Allocator, io: std.Io, port: u16, method: []const u8, path: []const u8, body: []const u8) !Reply { +pub fn request(a: Allocator, io: std.Io, port: u16, method: []const u8, path: []const u8, body: []const u8) !Reply { const addr: std.Io.net.IpAddress = .{ .ip4 = .loopback(port) }; const stream = try addr.connect(io, .{ .mode = .stream }); defer stream.close(io); @@ -175,7 +184,7 @@ fn unchunk(a: Allocator, s: []const u8) ![]const u8 { return out.items; } -const Background = struct { +pub const Background = struct { a: Allocator, io: std.Io, port: u16, @@ -191,11 +200,11 @@ const Background = struct { } }; -fn chatBody(a: Allocator, marker: []const u8, stream: bool) ![]const u8 { +pub fn chatBody(a: Allocator, marker: []const u8, stream: bool) ![]const u8 { return std.fmt.allocPrint(a, "{{\"messages\":[{{\"role\":\"user\",\"content\":\"go {s}\"}}],\"stream\":{},\"max_tokens\":4000}}", .{ marker, stream }); } -const Served = struct { +pub const Served = struct { code: u8 = 255, fn run(x: *Served, gpa: Allocator, io: std.Io, setup: serve.Setup) void { x.code = serve.run(gpa, io, .{ .host = "127.0.0.1", .port = 0 }, setup); @@ -203,7 +212,7 @@ const Served = struct { }; /// serve.run on its own thread with `drain_s`, the engine `slow`; the port once it listens. -fn start(gpa: Allocator, io: std.Io, env: *std.process.Environ.Map, slow: *Slow, tt: *TestText, x: *Served) !std.Thread { +pub fn start(gpa: Allocator, io: std.Io, env: *std.process.Environ.Map, slow: *Slow, tt: *TestText, x: *Served) !std.Thread { listen_port.store(0, .release); const t = try std.Thread.spawn(.{}, Served.run, .{ x, gpa, io, serve.Setup{ .engine = slow.engine(), @@ -217,7 +226,7 @@ fn start(gpa: Allocator, io: std.Io, env: *std.process.Environ.Map, slow: *Slow, return t; } -fn waitActive(io: std.Io, slow: *Slow, n: u32) void { +pub fn waitActive(io: std.Io, slow: *Slow, n: u32) void { while (slow.active.load(.acquire) < n) std.Io.sleep(io, .fromMilliseconds(2), .awake) catch {}; } diff --git a/zig/src/server/keepalive_test.zig b/zig/src/server/keepalive_test.zig new file mode 100644 index 0000000000..43435c1d04 --- /dev/null +++ b/zig/src/server/keepalive_test.zig @@ -0,0 +1,75 @@ +//! A silent stream keeps its client: keepalives on both routes before the first token; a failed one ends the request. +const std = @import("std"); +const server = @import("server.zig"); +const d = @import("drain_test.zig"); +const testing = std.testing; + +/// A raw request whose response is read until the connection ends. +fn messagesBody(a: std.mem.Allocator, marker: []const u8) ![]const u8 { + return std.fmt.allocPrint(a, "{{\"model\":\"x\",\"max_tokens\":50,\"stream\":true,\"messages\":[{{\"role\":\"user\",\"content\":\"go {s}\"}}]}}", .{marker}); +} + +/// Sends `body` streamed and hangs up once the response head arrived and the request reached `slow`. +fn leave(a: std.mem.Allocator, io: std.Io, port: u16, body: []const u8, slow: *d.Slow) !void { + const addr: std.Io.net.IpAddress = .{ .ip4 = .loopback(port) }; + const stream = try addr.connect(io, .{ .mode = .stream }); + defer stream.close(io); + const head = try std.fmt.allocPrint(a, "POST /v1/chat/completions HTTP/1.1\r\nHost: x\r\nContent-Type: application/json\r\nContent-Length: {d}\r\n\r\n{s}", .{ body.len, body }); + var wbuf: [4096]u8 = undefined; + var w = std.Io.net.Stream.Writer.init(stream, io, &wbuf); + try w.interface.writeAll(head); + try w.interface.flush(); + var rbuf: [64]u8 = undefined; + var r = std.Io.net.Stream.Reader.init(stream, io, &rbuf); + _ = try r.interface.peek(1); + // the head goes out before the engine has the request: a hang-up before that leaves nothing to cancel + d.waitActive(io, slow, 1); +} + +test "a silent engine's stream gets keepalives before its first token (chat: SSE comments, messages: ping events)" { + const period = comptime std.meta.fieldInfo(server.Config, .sse_keepalive_s).attrs.defaultValue(f64).?; + try testing.expect(period > 0 and period <= 15); + const gpa = testing.allocator; + const io = testing.io; + var tt: d.TestText = .{}; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var env = std.process.Environ.Map.init(gpa); + defer env.deinit(); + try env.put("TENSORFOLD_NO_LIVE", "1"); + try env.put("TF_DRAIN_S", "5"); + try env.put("TF_SSE_KEEPALIVE_S", "0.1"); + var slow: d.Slow = .{ .gpa = gpa, .io = io }; + var x: d.Served = .{}; + const t = try d.start(gpa, io, &env, &slow, &tt, &x); + const port = d.listen_port.load(.acquire); + + const chat = try d.request(a, io, port, "POST", "/v1/chat/completions", try d.chatBody(a, "@wait", true)); + try testing.expectEqual(@as(u16, 200), chat.status); + const first_token = std.mem.indexOf(u8, chat.body, "\" ok\"") orelse { + std.debug.print("{s}\n", .{chat.body[0..@min(chat.body.len, 1500)]}); + return error.NoToken; + }; + try testing.expect(std.mem.count(u8, chat.body[0..first_token], ": keepalive\n\n") >= 3); + try testing.expect(std.mem.indexOf(u8, chat.body, "data: [DONE]") != null); + + const msg = try d.request(a, io, port, "POST", "/v1/messages", try messagesBody(a, "@wait")); + try testing.expectEqual(@as(u16, 200), msg.status); + const start = std.mem.indexOf(u8, msg.body, "event: message_start") orelse return error.NoStart; + const ping = std.mem.indexOf(u8, msg.body, "event: ping\ndata: {\"type\": \"ping\"}") orelse return error.NoPing; + const delta = std.mem.indexOf(u8, msg.body, "event: content_block_delta") orelse return error.NoDelta; + try testing.expect(start < ping and ping < delta); + + // a client that leaves a silent stream: the request is cancelled (the disconnect poll or a failed keepalive) + var idle: u32 = 0; + while (slow.active.load(.acquire) > 0 and idle < 300) : (idle += 1) std.Io.sleep(io, .fromMilliseconds(10), .awake) catch {}; + try leave(a, io, port, try d.chatBody(a, "@hang", true), &slow); + var waited: u32 = 0; + while (slow.cancelled.load(.acquire) == 0 and waited < 300) : (waited += 1) std.Io.sleep(io, .fromMilliseconds(10), .awake) catch {}; + try testing.expect(slow.cancelled.load(.acquire) != 0); + + try std.posix.raise(.TERM); + t.join(); + try testing.expectEqual(@as(u8, 0), x.code); +} diff --git a/zig/src/server/openai.zig b/zig/src/server/openai.zig index 41b0ce7fa1..54e5771b4b 100644 --- a/zig/src/server/openai.zig +++ b/zig/src/server/openai.zig @@ -38,6 +38,8 @@ pub const Out = struct { event: *const fn (ctx: *anyopaque, payload: ?Value) error{Closed}!void, /// A whole JSON reply. reply: *const fn (ctx: *anyopaque, status: u16, payload: Value) void, + /// Bytes that keep a silent stream alive (an SSE comment, or the API's ping); null: none. + keepalive: ?*const fn (ctx: *anyopaque) error{Closed}!void = null, }; }; @@ -58,7 +60,7 @@ const ClientOut = struct { a: Allocator, fn out(c: *ClientOut) Out { - return .{ .ctx = c, .vt = &.{ .open = open, .event = event, .reply = reply } }; + return .{ .ctx = c, .vt = &.{ .open = open, .event = event, .reply = reply, .keepalive = keepalive } }; } fn open(ctx: *anyopaque, _: usize) error{Closed}!void { @@ -75,14 +77,31 @@ const ClientOut = struct { const c: *ClientOut = @ptrCast(@alignCast(ctx)); routes.sendValue(c.conn, c.a, status, payload); } + + fn keepalive(ctx: *anyopaque) error{Closed}!void { + const c: *ClientOut = @ptrCast(@alignCast(ctx)); + return sse.comment(c.conn); + } }; /// Whether the client has gone, read between rounds (Python's socket_cancellation). pub const Gone = struct { conn: *Conn, + /// The stream a keepalive goes to once it has been silent `every_ns` (0: never) + out: ?Out = null, + every_ns: i96 = 0, pub fn check(g: Gone) bool { return g.conn.peerGone(); } + + /// A keepalive if the stream has been silent long enough; a failed write means the client left. + pub fn beat(g: Gone, silent_ns: i96) error{Closed}!bool { + const o = g.out orelse return false; + const send = o.vt.keepalive orelse return false; + if (g.every_ns <= 0 or silent_ns < g.every_ns) return false; + try send(o.ctx); + return true; + } }; /// ``error_body``: OpenAI's error object, with param and code where clients key on them. diff --git a/zig/src/server/openai_run.zig b/zig/src/server/openai_run.zig index 9ea9f740ab..5476dc8674 100644 --- a/zig/src/server/openai_run.zig +++ b/zig/src/server/openai_run.zig @@ -280,8 +280,11 @@ pub const Run = struct { if (!tools and r.is_chat) r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; var sink_state: StreamSink = .{ .run = r, .tools = tools }; handed = true; + var live = gone; + live.out = r.out; + live.every_ns = @intFromFloat(r.srv.config.sse_keepalive_s * std.time.ns_per_s); const sink: chat.Sink = .{ .ctx = &sink_state, .call = StreamSink.call }; - var reply = (if (r.turn) |t| cortheon.loop.stream(t, r.srv, &cx, &r.plan.input, prepared, sink, gone) else chat.generate(r.srv, &cx, prepared, sink, gone)) catch |e| { + var reply = (if (r.turn) |t| cortheon.loop.stream(t, r.srv, &cx, &r.plan.input, prepared, sink, live) else chat.generate(r.srv, &cx, prepared, sink, live)) catch |e| { switch (e) { error.Cancelled => return, error.Refused => if (cx.kind != .other and cx.kind != .server) { diff --git a/zig/src/server/responses.zig b/zig/src/server/responses.zig index 2e096025b5..598d4b336a 100644 --- a/zig/src/server/responses.zig +++ b/zig/src/server/responses.zig @@ -203,7 +203,7 @@ const Wire = struct { payload: Value = .null, fn out(w: *Wire) openai.Out { - return .{ .ctx = w, .vt = &.{ .open = open, .event = event, .reply = whole } }; + return .{ .ctx = w, .vt = &.{ .open = open, .event = event, .reply = whole, .keepalive = keepalive } }; } fn open(ctx: *anyopaque, _: usize) error{Closed}!void { @@ -223,6 +223,11 @@ const Wire = struct { w.status = status; w.payload = payload; } + + fn keepalive(ctx: *anyopaque) error{Closed}!void { + const w: *Wire = @ptrCast(@alignCast(ctx)); + return sse.comment(w.reply.conn); + } }; /// A Response built from chat-completion deltas, in order: reasoning, text, calls; streamed as typed events. diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index af6e3057a5..0db752b8af 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -44,6 +44,7 @@ test { _ = @import("stream_dsml_test.zig"); _ = @import("logprobs_test.zig"); _ = @import("drain_test.zig"); + _ = @import("keepalive_test.zig"); _ = @import("late_system_test.zig"); _ = @import("compact_test.zig"); _ = @import("completion_prompt_test.zig"); diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index e83a9bbf97..4fa490d103 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -135,6 +135,10 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { std.debug.print("tensorfold: TF_DRAIN_S: seconds, a number >= 0\n", .{}); return 1; }; + const keepalive_s = seconds(env(s, "TF_SSE_KEEPALIVE_S"), 15) orelse { + std.debug.print("tensorfold: TF_SSE_KEEPALIVE_S: seconds, a number >= 0\n", .{}); + return 1; + }; const config: server_mod.Config = .{ .served_name = s.served, .model_dir = s.model_dir, @@ -148,6 +152,7 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { .default_sampling = s.sampling, .use_drafts = !args.no_drafts, .seed_salt = salt, + .sse_keepalive_s = keepalive_s, .request_log = env(s, "TENSORFOLD_REQUEST_LOG"), .dashboard = args.dashboard, .family = s.family, @@ -239,7 +244,7 @@ fn drain(srv: *server_mod.Server, io: std.Io, stop: ?EngineStop, drain_s: f64) v while (srv.generating.load(.seq_cst) > 0 and waited < 100) : (waited += 1) std.Io.sleep(io, .fromMilliseconds(20), .awake) catch {}; } -/// A seconds setting (TF_DRAIN_S): unset or blank keeps the default; null: not a number >= 0. +/// A seconds setting (TF_DRAIN_S, TF_SSE_KEEPALIVE_S): unset or blank keeps the default; null: not a number >= 0. fn seconds(raw: ?[]const u8, default: f64) ?f64 { const t = std.mem.trim(u8, raw orelse "", " \t"); if (t.len == 0) return default; diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 79efe48e94..505fbdc344 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -63,6 +63,8 @@ pub const Config = struct { /// BRAVE_SEARCH_API_KEY: Cortheon searches through Brave instead of DuckDuckGo. cortheon_search_key: ?[]const u8 = null, timeouts: http_conn.Timeouts = .{}, + /// TF_SSE_KEEPALIVE_S: a silent stream's keepalive period before its first token (serve.zig; 0: none) + sse_keepalive_s: f64 = 15, /// The checkpoint family's own chat rules (family.zig); null: the chat template path. family: ?family_mod.Family = null, }; diff --git a/zig/src/server/sse.zig b/zig/src/server/sse.zig index de30ac1212..bcdcf56f4c 100644 --- a/zig/src/server/sse.zig +++ b/zig/src/server/sse.zig @@ -19,6 +19,11 @@ pub fn data(conn: *Conn, a: Allocator, payload: json.Value) error{Closed}!void { return conn.writeAll(std.mem.concat(a, u8, &.{ "data: ", text, "\n\n" }) catch return error.Closed); } +/// An SSE comment line, which clients skip: it keeps a stream alive while nothing else is sent. +pub fn comment(conn: *Conn) error{Closed}!void { + return conn.writeAll(": keepalive\n\n"); +} + pub fn done(conn: *Conn) error{Closed}!void { return conn.writeAll("data: [DONE]\n\n"); } From a3e2fb184bc80be3bf9e280d20c0fbe9c3cfdedc Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Wed, 14 Oct 2026 21:06:16 +0700 Subject: [PATCH 25/43] TP NCCL: all_reduce's doc says NCCL adds in its ring's order, so its rank-order sum holds at two ranks --- zig/src/fabric/tp/nccl.zig | 1 + 1 file changed, 1 insertion(+) diff --git a/zig/src/fabric/tp/nccl.zig b/zig/src/fabric/tp/nccl.zig index b641839aad..c9369369ef 100644 --- a/zig/src/fabric/tp/nccl.zig +++ b/zig/src/fabric/tp/nccl.zig @@ -97,6 +97,7 @@ pub const Nccl = struct { if (n.aborted.load(.acquire)) return error.Aborted; } + /// NCCL adds in its ring's order: the rank-order sum `all_reduce` promises holds at two ranks, not past them. pub fn allReduce(ptr: *anyopaque, s: DevicePtr, r: DevicePtr, count: usize, t: DType, op: Op, stream: collective.Stream) Error!void { const n = self_(ptr); try n.live(); From 0d0bdb216313a5ed85feb27e69dbd6f398d7123a Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 11 Oct 2026 18:33:54 +0700 Subject: [PATCH 26/43] Sessions: the page map is rehashed before tombstones fill it The prefix index keeps one node per distinct full page of every saved entry in `by_digest`. Each saved turn adds its new pages, and each entry dropped or evicted removes the pages no other entry passes through, for the life of the server. std's tombstones never clear, so in time they take every free slot, and each save's new pages and each find's walk past the cached prefix then probe the whole map. The map now counts its removals and is rehashed in place once they reach a tenth of its capacity. Test: 3,000 turns of up to 60 fresh pages, the oldest dropped past eight, keep a tenth of the map free after every turn; without the rehash it fails. --- zig/src/sessions/prefix.zig | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/zig/src/sessions/prefix.zig b/zig/src/sessions/prefix.zig index 28c639fbd7..98634bbfd2 100644 --- a/zig/src/sessions/prefix.zig +++ b/zig/src/sessions/prefix.zig @@ -85,6 +85,7 @@ pub fn Index(comptime Extra: type) type { nodes: std.ArrayList(Node) = .empty, free_nodes: std.ArrayList(u32) = .empty, by_digest: std.AutoHashMapUnmanaged(Digest, u32) = .empty, + dead: u32 = 0, // nodes dropped from `by_digest` since it was last rehashed entries: std.ArrayList(Entry) = .empty, free_entries: std.ArrayList(u32) = .empty, by_key: std.AutoHashMapUnmanaged(Digest, u32) = .empty, @@ -259,12 +260,20 @@ pub fn Index(comptime Extra: type) type { if (nd.prev_sib != none) x.nodes.items[nd.prev_sib].next_sib = nd.next_sib else x.nodes.items[nd.parent].first_child = nd.next_sib; if (nd.next_sib != none) x.nodes.items[nd.next_sib].prev_sib = nd.prev_sib; } - _ = x.by_digest.remove(nd.digest); + if (x.by_digest.remove(nd.digest)) x.dropped(); x.free_nodes.append(x.gpa, m) catch {}; m = nd.parent; } } + /// std's tombstones never clear and keys fill 80 %: rehash at a tenth, so a tenth of the map always stays free. + fn dropped(x: *Self) void { + x.dead += 1; + if (@as(u64, x.dead) * 10 < x.by_digest.capacity()) return; + x.by_digest.rehash(std.hash_map.AutoContext(Digest){}); + x.dead = 0; + } + /// The entry's tail matches the prompt there and the entry is shorter (its full pages matched by digest). fn tailOf(e: *const Entry, prompt: []const i32) bool { const at: usize = @intCast(e.pos - e.tail.len); @@ -455,3 +464,23 @@ test "longest strict prefix, ties, tags and coverage" { try std.testing.expectEqual(@as(?u32, null), x.find(7, chat[0..150], d[0..9], .any)); try std.testing.expectEqual(@as(u32, 1 + 6), x.by_digest.count()); // tag 8's root + its 6 full pages } + +test "a long run of turns saved and dropped leaves the page map free slots" { + const gpa = std.testing.allocator; + var x = T.init(gpa, 16); + defer x.deinit(); + var live: std.ArrayList(u32) = .empty; + defer live.deinit(gpa); + var rng = std.Random.DefaultPrng.init(5); + for (0..3000) |i| { + // a chat's turn (its own pages), the oldest dropped once eight are kept: nodes come and go + const ids = try seq(gpa, 16 * 60, @intCast(i)); + defer gpa.free(ids); + try live.append(gpa, try add(&x, 7, ids[0 .. 16 * (1 + rng.random().uintLessThan(usize, 60))], true, i)); + if (live.items.len > 8) x.remove(live.orderedRemove(0)); + const meta = x.by_digest.metadata orelse continue; + var free: usize = 0; + for (meta[0..x.by_digest.capacity()]) |m| free += @intFromBool(m.isFree()); + try std.testing.expect(free * 10 >= x.by_digest.capacity()); + } +} From 2b3cccf2ab2fefb4a9c715ce7a605755328c9fa4 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:10:45 +0700 Subject: [PATCH 27/43] Server tests: the drain test's reply in progress is held until its drain checks are done The listener stays open only while a generation runs. Case 1 streamed a "@short" reply (about 180 ms) and then sent a messages request and a /health request expecting 503; when the reply ended first, the drain finished and those requests found the listener closed. The stub engine's "@hold" reply now waits for the test's release. --- zig/src/server/drain_test.zig | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/zig/src/server/drain_test.zig b/zig/src/server/drain_test.zig index 411b651cf3..09873e92f5 100644 --- a/zig/src/server/drain_test.zig +++ b/zig/src/server/drain_test.zig @@ -9,7 +9,7 @@ const routes = @import("routes.zig"); const Allocator = std.mem.Allocator; const testing = std.testing; -/// Slow enough to stop mid-reply: "@short" ends after a dozen tokens, "@long" only when halted or cancelled. +/// Stops mid-reply: "@short" ends after 12 tokens, "@long" when halted or cancelled, "@hold" waits for `release`. pub const Slow = struct { gpa: Allocator, io: std.Io, @@ -21,6 +21,8 @@ pub const Slow = struct { /// halts before the first close (the halt must come first) halts_at_close: u32 = 0, cancelled: std.atomic.Value(u64) = .init(0), + /// "@hold" replies stay silent until this is set + release: std.atomic.Value(bool) = .init(false), fn engine(s: *Slow) api.Engine { return .{ .ctx = s, .vtable = &.{ .info = info, .submit = submit, .cancel = cancel, .status = status, .memory = memory } }; @@ -67,9 +69,10 @@ pub const Slow = struct { const word: []const u32 = if (long) &.{ ' ', 'l', 'a' } else &.{ ' ', 'o', 'k' }; const hang = std.mem.indexOf(u8, prompt, "@hang") != null; const wait = std.mem.indexOf(u8, prompt, "@wait") != null; + const hold = std.mem.indexOf(u8, prompt, "@hold") != null; // a long prefill: nothing for 800 ms ("@wait"), or nothing until stopped ("@hang") var quiet: u32 = 0; - while (hang or (wait and quiet < 80)) : (quiet += 1) { + while (hang or (wait and quiet < 80) or (hold and !s.release.load(.acquire))) : (quiet += 1) { if (s.halted.load(.acquire)) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .failed, .message = s.reason } }); if (s.cancelled.load(.acquire) == id) return sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .cancelled } }); std.Io.sleep(s.io, .fromMilliseconds(10), .awake) catch {}; @@ -248,8 +251,10 @@ test "SIGTERM drains: 503 + Retry-After and /health draining while a reply finis var x: Served = .{}; const t = try start(gpa, io, &env, &slow, &tt, &x); const port = listen_port.load(.acquire); - var bg: Background = .{ .a = a, .io = io, .port = port, .body = try chatBody(a, "@short", true) }; + // held through the checks below: a reply that ended first would end the drain and close the listener + var bg: Background = .{ .a = a, .io = io, .port = port, .body = try chatBody(a, "@hold", true) }; const bt = try std.Thread.spawn(.{}, Background.run, .{&bg}); + errdefer slow.release.store(true, .release); waitActive(io, &slow, 1); try std.posix.raise(.TERM); // the drain has begun once a new request is refused @@ -268,6 +273,7 @@ test "SIGTERM drains: 503 + Retry-After and /health draining while a reply finis try testing.expectEqual(@as(u16, 503), health.status); try testing.expect(std.mem.indexOf(u8, health.head, "Retry-After: 30") != null); try testing.expect(std.mem.indexOf(u8, health.body, "\"status\": \"draining\"") != null); + slow.release.store(true, .release); bt.join(); t.join(); try testing.expect(!bg.failed); From 62a75ae677057ce5c5a55b07a20581f28e999ab7 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:13:19 +0700 Subject: [PATCH 28/43] TP: keep Linux-only transport host tests off macOS builds --- zig/build/tp.zig | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/zig/build/tp.zig b/zig/build/tp.zig index f5dacc46e3..873127c6af 100644 --- a/zig/build/tp.zig +++ b/zig/build/tp.zig @@ -1,4 +1,4 @@ -//! Tensor parallelism's build: host tests on any OS, and on Linux the `tf-tp-test` GPU runner with the mailbox kernel. +//! Tensor parallelism's build: Linux host tests, and on Linux the `tf-tp-test` GPU runner with the mailbox kernel. const std = @import("std"); @@ -41,9 +41,11 @@ fn kernelModule(b: *std.Build, image: ?std.Build.LazyPath) *std.Build.Module { pub fn add(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, test_step: *std.Build.Step) void { const host = b.graph.host; - const tests = b.addRunArtifact(b.addTest(.{ .root_module = module(b, host, .debug, runtime(b, host, .debug)) })); - b.step("test-tp", "Tensor parallelism's host tests: bootstrap, fate channel, in-process ranks (no GPU)").dependOn(&tests.step); - test_step.dependOn(&tests.step); + if (host.result.os.tag == .linux) { + const tests = b.addRunArtifact(b.addTest(.{ .root_module = module(b, host, .debug, runtime(b, host, .debug)) })); + b.step("test-tp", "Tensor parallelism's host tests: bootstrap, fate channel, in-process ranks (no GPU)").dependOn(&tests.step); + test_step.dependOn(&tests.step); + } if (target.result.os.tag != .linux) return; const nvcc = b.option([]const u8, "tp-nvcc", "nvcc that builds the TP mailbox kernel"); const prebuilt = b.option([]const u8, "tp-fatbin", "absolute path of a prebuilt tp_mailbox.fatbin to embed"); From 2cf61d3cbaa2df71b1c369f8d1f73c423cfec8ae Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:17:20 +0700 Subject: [PATCH 29/43] Sessions: honor zero residency budgets and reserve uncapped requests Use null for unlimited residency, retain portable bookkeeping tests on macOS, and run the Linux NVMe tier tests only on Linux. --- zig/build/sessions.zig | 13 ++++++++++--- zig/src/sessions/pool.zig | 11 ++++++++++- zig/src/sessions/portable_test.zig | 8 ++++++++ zig/src/sessions/store.zig | 12 ++++++------ zig/src/sessions/store_test.zig | 25 +++++++++++++++++++++++++ 5 files changed, 59 insertions(+), 10 deletions(-) create mode 100644 zig/src/sessions/portable_test.zig diff --git a/zig/build/sessions.zig b/zig/build/sessions.zig index a5d8ea6f3c..5dd2331845 100644 --- a/zig/build/sessions.zig +++ b/zig/build/sessions.zig @@ -1,4 +1,4 @@ -//! Session memory's build: the family-neutral `sessions` module (zig/src/sessions) and its host tests on any OS. +//! Session memory: portable bookkeeping tests everywhere, NVMe tier tests on Linux. const std = @import("std"); @@ -7,7 +7,14 @@ pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bui } pub fn add(b: *std.Build, test_step: *std.Build.Step) void { - const run = b.addRunArtifact(b.addTest(.{ .root_module = module(b, b.graph.host, .debug) })); - b.step("test-sessions", "Session memory's host tests: pool, prefix trie, RAM tier, NVMe tier on real files (no GPU)").dependOn(&run.step); + const host = b.graph.host; + const tests = if (host.result.os.tag == .linux) module(b, host, .debug) else b.createModule(.{ + .root_source_file = b.path("zig/src/sessions/portable_test.zig"), + .target = host, + .optimize = .debug, + .link_libc = true, + }); + const run = b.addRunArtifact(b.addTest(.{ .root_module = tests })); + b.step("test-sessions", "Session memory's host tests: pool, prefix trie, RAM tier, NVMe tier on Linux (no GPU)").dependOn(&run.step); test_step.dependOn(&run.step); } diff --git a/zig/src/sessions/pool.zig b/zig/src/sessions/pool.zig index d3e67516af..85ee2c48c8 100644 --- a/zig/src/sessions/pool.zig +++ b/zig/src/sessions/pool.zig @@ -379,7 +379,8 @@ pub const Pool = struct { /// Pages an admission needs for a request (prompt + reply + slack, at most the slot's capacity). pub fn needPages(p: *const Pool, prompt: u64, max_tokens: u64, capacity: u64, slack: u64) u32 { - return pagesFor(@min(prompt + max_tokens + slack, capacity), p.page); + const total = prompt + max_tokens + slack; + return pagesFor(if (capacity == 0) total else @min(total, capacity), p.page); } /// This rank stores the split families' rows of page pg. @@ -473,3 +474,11 @@ test "residues, holders, quotas and the device tables" { try std.testing.expectEqual(@as(u16, 3), p.holders(b.pages[0])); try std.testing.expectEqual(@as(u32, 0), c.pages[2] % 2); } + +test "admission without a capacity still reserves prompt, reply and slack pages" { + var p = try Pool.init(std.testing.allocator, .{ .families = &.{}, .page = 64 }, 1024, 1, 0); + defer p.deinit(); + try std.testing.expectEqual(@as(u32, 3), p.needPages(65, 64, 0, 1)); + try std.testing.expectEqual(@as(u32, 1), p.needPages(65, 64, 64, 1)); + try std.testing.expectEqual(@as(u32, 0), p.needPages(0, 0, 0, 0)); +} diff --git a/zig/src/sessions/portable_test.zig b/zig/src/sessions/portable_test.zig new file mode 100644 index 0000000000..6ba2694b4e --- /dev/null +++ b/zig/src/sessions/portable_test.zig @@ -0,0 +1,8 @@ +//! Portable session bookkeeping tests; the Linux NVMe tier is tested separately. +test { + _ = @import("freeset.zig"); + _ = @import("pool.zig"); + _ = @import("pagestore.zig"); + _ = @import("prefix.zig"); + _ = @import("format.zig"); +} diff --git a/zig/src/sessions/store.zig b/zig/src/sessions/store.zig index 851ce509aa..0bd128dd84 100644 --- a/zig/src/sessions/store.zig +++ b/zig/src/sessions/store.zig @@ -31,9 +31,9 @@ pub const Bounded = struct { pub const Policy = struct { /// bounded state the RAM tier keeps (least recently used entries go past it) ram_bytes: u64 = 512 << 20, - /// idle chats resident in the pool (0: no limit) and pool bytes held by entries alone (0: no limit) - chats: u32 = 0, - resident_bytes: u64 = 0, + /// Idle chats and entry-only pool bytes: null is unlimited; zero retains no idle residency. + chats: ?u32 = null, + resident_bytes: ?u64 = null, /// park an evicted chat's older turns too (default: drop them, the leaf's file covers them) park_prefixes: bool = false, /// drop, never write, an entry a newer entry of its chain covers @@ -479,7 +479,7 @@ pub const Store = struct { /// Evicts LRU idle chats past the residency budget (leaf parked, unshared older turns dropped); returns how many. pub fn enforce(s: *Store) !u32 { const pol = s.policy; - if (pol.chats == 0 and pol.resident_bytes == 0) return 0; + if (pol.chats == null and pol.resident_bytes == null) return 0; var chats: std.ArrayList(Chat) = .empty; defer chats.deinit(s.gpa); var members: std.ArrayList(u32) = .empty; @@ -494,8 +494,8 @@ pub const Store = struct { var idle: u32 = @intCast(chats.items.len); var gone: u32 = 0; for (chats.items) |c| { - const over_n = pol.chats > 0 and idle > pol.chats; - const over_b = pol.resident_bytes > 0 and s.pool.entryOnlyBytes() > pol.resident_bytes; + const over_n = if (pol.chats) |limit| idle > limit else false; + const over_b = if (pol.resident_bytes) |limit| s.pool.entryOnlyBytes() > limit else false; if (!over_n and !over_b) break; if (!s.index.get(c.leaf).ram or !s.idleLeaf(c.leaf)) continue; members.clearRetainingCapacity(); diff --git a/zig/src/sessions/store_test.zig b/zig/src/sessions/store_test.zig index 86b58a4f59..97d018d8a0 100644 --- a/zig/src/sessions/store_test.zig +++ b/zig/src/sessions/store_test.zig @@ -214,3 +214,28 @@ test "two split ranks make the same decisions on the same calls and each restore try restart(gpa, root, 2, 1); try testing.expectEqualSlices(u64, t0.items, t1.items); } + +test "zero chat or byte budgets evict idle residency while null remains unlimited" { + const gpa = testing.allocator; + var tmp = testing.tmpDir(.{}); + defer tmp.cleanup(); + const root = try tmpRoot(gpa, &tmp, "zero-budget"); + defer gpa.free(root); + var r: Rank = undefined; + try r.init(gpa, root, 1, 0, .{}); + defer r.deinit(); + const slot = try r.pool.newSlot(64 * page); + const ids = try chatIds(gpa, 1, 5 * page); + defer gpa.free(ids); + _ = try turn(gpa, &r, slot, 1, ids); + try testing.expectEqual(@as(u32, 1), r.store.resident().chats); + r.store.policy.chats = 0; + try testing.expectEqual(@as(u32, 1), try r.store.enforce()); + try testing.expectEqual(@as(u32, 0), r.store.resident().chats); + r.store.policy.chats = null; + _ = try turn(gpa, &r, slot, 1, ids); + try testing.expect(r.store.resident().bytes > 0); + r.store.policy.resident_bytes = 0; + try testing.expectEqual(@as(u32, 1), try r.store.enforce()); + try testing.expectEqual(@as(u64, 0), r.store.resident().bytes); +} From 93677846ef0b922bad3a42c5b305464e6af9f03e Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:22:29 +0700 Subject: [PATCH 30/43] server: keep shared halt and drain on the existing template path --- zig/src/server/anthropic.zig | 3 +- zig/src/server/chat.zig | 361 ++++++++++++++++++++--- zig/src/server/chat_generation.zig | 452 ----------------------------- zig/src/server/family.zig | 135 --------- zig/src/server/openai.zig | 344 ++++++++++++++++++++-- zig/src/server/openai_run.zig | 346 ---------------------- zig/src/server/prompt.zig | 19 +- zig/src/server/root.zig | 1 - zig/src/server/serve.zig | 3 - zig/src/server/server.zig | 10 +- zig/src/server/tokens.zig | 22 +- 11 files changed, 640 insertions(+), 1056 deletions(-) delete mode 100644 zig/src/server/chat_generation.zig delete mode 100644 zig/src/server/family.zig delete mode 100644 zig/src/server/openai_run.zig diff --git a/zig/src/server/anthropic.zig b/zig/src/server/anthropic.zig index f1d8e0ec23..9ffdfa15f4 100644 --- a/zig/src/server/anthropic.zig +++ b/zig/src/server/anthropic.zig @@ -213,8 +213,7 @@ const Reply = struct { fn text(r: *Reply, kind: []const u8, value: []const u8) E!void { if (value.len == 0) return; if (r.active != null) { - // held until the open tool_use block closes: copied, ``value`` is a slice of a text that grows by then - r.pending.append(r.a, .{ .kind = kind, .text = r.a.dupe(u8, value) catch return error.Closed }) catch return error.Closed; + r.pending.append(r.a, .{ .kind = kind, .text = value }) catch return error.Closed; return; } const key = if (std.mem.eql(u8, kind, "thinking")) "thinking" else "text"; diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 3183ea6df8..ce1dffbbb4 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -12,12 +12,8 @@ const log = @import("log.zig"); const ids = @import("ids.zig"); const clock = @import("clock.zig"); const Server = @import("server.zig").Server; -const family_mod = @import("family.zig"); const chunk_plan = @import("chunk_plan.zig"); const Mailbox = @import("mailbox.zig").Mailbox; -const chat_generation = @import("chat_generation.zig"); -const Shown = chat_generation.Shown; -const Generation = chat_generation.Generation; const Value = json.Value; const Cx = errors.Cx; const Allocator = std.mem.Allocator; @@ -39,10 +35,6 @@ pub const Input = struct { logprobs: ?u8 = null, /// ``/v1/completions``' ``prompt_logprobs``: rows for the prompt with this many alternatives; null: not asked. prompt_logprobs: ?u8 = null, - /// The request body (a family reads fields of its own from it, such as ``response_format``). - body: Value = .null, - /// ``parallel_tool_calls: false``: the reply keeps its first call (a family applies it while streaming). - single_call: bool = false, }; /// A streamed piece: content text (a string) or a delta object (reasoning or tool calls). @@ -70,8 +62,6 @@ pub const Reply = struct { logprobs: ?Value = null, /// ``choices[0].prompt_logprobs`` (vLLM's form) for a request that asked; null otherwise. prompt_logprobs: ?Value = null, - /// A family's parsed calls (OpenAI tool-call objects); null: the route parses the content itself. - calls: ?[]Value = null, }; pub const Failure = error{ Refused, Cancelled, Failed, OutOfMemory }; @@ -98,9 +88,6 @@ pub const Prepared = struct { drafts: bool, stops: fields_mod.Stops, preparing: bool, - images_omitted: u32 = 0, - /// the prompt's prepared images (a family's), held until the reply ends - images: ?family_mod.Images = null, }; /// Render ``input`` and run every check that can refuse it, before anything reaches the client or the engine. @@ -120,10 +107,8 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa if (gone.check()) return error.Cancelled; var thinking = flag(f, "enable_thinking") orelse srv.config.enable_thinking; if (input.prompt != null) thinking = false; - const asked: ?[]const u8 = if (f.get("reasoning_effort")) |e| (if (e == .string) e.string else null) else null; - const effort: ?[]const u8 = if (srv.family) |fam| asked orelse srv.config.reasoning_effort orelse fam.defaultEffort() else srv.effortFor(asked); + const effort = srv.effortFor(if (f.get("reasoning_effort")) |e| (if (e == .string) e.string else null) else null); const rendered = try prompt_mod.prepare(srv, cx, input, thinking, effort); - errdefer if (rendered.images) |im| im.release(); if (gone.check()) return error.Cancelled; if (rendered.ids.len == 0) return cx.refuse("rendered prompt is empty"); const window: i64 = srv.info.context_window; @@ -172,7 +157,7 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa } if (input.logprobs) |count| request.logprobs = try @import("logprobs.zig").admit(srv, cx, thinking, request.think_budget, count); request.prompt_logprobs = input.prompt_logprobs; - return .{ .input = input, .request = request, .prompt_len = rendered.ids.len, .received = received, .thinking = thinking, .effort = effort, .sampled = sampling != null, .drafts = drafts, .stops = stops_opt, .preparing = preparing, .images_omitted = rendered.images_omitted, .images = rendered.images }; + return .{ .input = input, .request = request, .prompt_len = rendered.ids.len, .received = received, .thinking = thinking, .effort = effort, .sampled = sampling != null, .drafts = drafts, .stops = stops_opt, .preparing = preparing }; } /// Submit a prepared request and collect its reply; ``sink`` hears the stream (null: not streamed). @@ -187,7 +172,6 @@ pub fn generate(srv: *Server, cx: *Cx, prepared: Prepared, sink: ?Sink, gone: an const stops_opt = prepared.stops; var preparing = prepared.preparing; defer release(srv, preparing); - defer if (prepared.images) |im| im.release(); // after the engine's finished event (the loop below waits for it) if (request.background) { while (true) { const now = nowNs(io); @@ -199,19 +183,16 @@ pub fn generate(srv: *Server, cx: *Cx, prepared: Prepared, sink: ?Sink, gone: an } if (gone.check()) return error.Cancelled; var stop_hook: StopHook = .{ .srv = srv, .stops = .{ .strings = stops_opt.strings } }; - // a family matches stop strings in the answer (``push``); the generic path in the raw text on the engine's thread - if (stops_opt.strings.len > 0 and srv.family == null) request.stop = .{ .ctx = &stop_hook, .check = StopHook.check }; + if (stops_opt.strings.len > 0) request.stop = .{ .ctx = &stop_hook, .check = StopHook.check }; var box: Mailbox = .{ .io = io, .gpa = srv.gpa }; defer box.deinit(); const id = srv.next_id.fetchAdd(1, .monotonic); if (srv.keepalive) |k| k.begin(); // the GPU is busy: the idle ticker holds its commits defer if (srv.keepalive) |k| k.end(); - srv.engine.submit(id, &request, .{ .ctx = &box, .event = Mailbox.onEvent }) catch |e| { - return switch (e) { - error.Busy => cx.fail(.capacity, "the engine is busy; retry shortly", .{}), - error.Closed => cx.fail(.other, "the scheduler is closed", .{}), - error.InvalidSpans => cx.fail(.request, "the prompt span layout is invalid", .{}), - }; + srv.engine.submit(id, &request, .{ .ctx = &box, .event = Mailbox.onEvent }) catch |e| return switch (e) { + error.Busy => cx.fail(.capacity, "the engine is busy; retry shortly", .{}), + error.Closed => cx.fail(.other, "the scheduler is closed", .{}), + error.InvalidSpans => cx.fail(.request, "the prompt span layout is invalid", .{}), }; release(srv, preparing); // a background request waits only while a foreground one prepares preparing = false; @@ -219,11 +200,9 @@ pub fn generate(srv: *Server, cx: *Cx, prepared: Prepared, sink: ?Sink, gone: an 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); errdefer if (!gen.engine_done) gen.cancel(); // the engine writes to the mailbox until it says finished const result: Failure!Reply = blk: { - if (srv.family) |fam| { - gen.reader = fam.reader(a, thinking, input.tools, input.single_call) catch |e| break :blk e; - } else if (sink != null and input.tools.len > 0) gen.calls = tool_stream.Streamer.init(a, input.tools) catch |e| break :blk e; + if (sink != null and input.tools.len > 0) gen.calls = tool_stream.Streamer.init(a, input.tools) catch |e| break :blk e; gen.loop(gone) catch |e| break :blk e; - break :blk gen.finish(cx, prepared.prompt_len, received, thinking, effort, prepared.sampled, prepared.drafts, prepared.images_omitted); + break :blk gen.finish(cx, prepared.prompt_len, received, thinking, effort, prepared.sampled, prepared.drafts); }; if (result) |_| {} else |e| logEnded(input.id, e, cx.message, prepared.prompt_len, gen.collected.items.len, seconds(nowNs(io) - received)); return result; @@ -251,6 +230,312 @@ pub fn release(srv: *Server, preparing: bool) void { if (preparing) _ = srv.preparing.fetchSub(1, .acq_rel); } +/// The text a stream has sent: Python's ``streamed = visible``, grown by its delta while it only grows. +const Shown = struct { + text: std.ArrayList(u8) = .empty, + chars: usize = 0, + at: ?[*]const u8 = null, // where the last shown text lay; the decode buffer only grows, so that prefix holds + extends: bool = false, + + /// Python's ``now[len(shown):]``; ``stable``: ``now`` lies in the append-only decode buffer. + fn after(s: *Shown, now: []const u8, stable: bool) []const u8 { + const sent = s.text.items; + s.extends = (stable and s.at == now.ptr and now.len >= sent.len) or std.mem.startsWith(u8, now, sent); + if (!s.extends) return reply_text.afterChars(now, s.chars); + if (stable) s.at = now.ptr; // the buffer moved but kept its bytes: the next check is free again + return now[sent.len..]; + } + + fn set(s: *Shown, a: Allocator, now: []const u8, delta: []const u8) Allocator.Error!void { + if (s.extends) { + try s.text.appendSlice(a, delta); + s.chars += reply_text.charCount(delta); + } else { + s.text = .empty; + try s.text.appendSlice(a, now); + s.chars = reply_text.charCount(now); + } + s.at = now.ptr; + } +}; + +/// The token loop and the text it streams. +const Generation = struct { + srv: *Server, + a: Allocator, + box: *Mailbox, + id: api.Id, + reply_id: []const u8, // the id the client gets, which the done line prints + sink: ?Sink, + thinking: bool, + stops: reply_text.Stops, + ignore_eos: bool, + max_tokens: u32, + tools: []const Value, + calls: ?tool_stream.Streamer = null, + collected: std.ArrayList(u32) = .empty, + visible: reply_text.Incremental = .{}, + hidden: std.ArrayList(u8) = .empty, // reused for the answer without its call blocks + streamed: Shown = .{}, // what content streamed, kept apart from the decode buffer that grows under slices + streamed_reasoning: Shown = .{}, + streaming_done: bool = false, + first_ns: ?i96 = null, + last_ns: ?i96 = null, + wrote_ns: ?i96 = null, + reason: ?[]const u8 = null, // the server ended the reply (a stop string, the length) before the engine said so + engine_done: bool = false, + consumed: usize = 0, + chunk_index: usize = 0, + lp: ?@import("logprobs.zig").Stream = null, // the request asked for logprobs: its rows, copied with each round + + fn eos(g: *const Generation, t: u32) bool { + return !g.ignore_eos and std.mem.indexOfScalar(u32, g.srv.eos, t) != null; + } + + /// What closes a call the model's end token left open (``tool_parse.closeCall``); nothing when a stop string, the length or ``ignore_eos`` ended the reply instead. + fn closeCall(g: *const Generation, text: []const u8) Allocator.Error![]const u8 { + const t = g.collected.items; + if (g.tools.len == 0 or t.len == 0 or !g.eos(t[t.len - 1])) return ""; + return tool_parse.closeCall(g.a, text, g.tools); + } + + /// The engine's stop check, here: the newest tokens' text holds a stop string. + fn stopHit(g: *Generation) Allocator.Error!bool { + if (g.stops.strings.len == 0) return false; + const t = g.collected.items; + const tail = t[t.len -| g.stops.tail()..]; + const text = try g.srv.text.decode(g.a, tail); + for (g.stops.strings) |s| if (std.mem.indexOf(u8, text, s) != null) return true; + return false; + } + + /// Commits a round's tokens as ``LaneStream.commit`` would, then streams what they add. + fn commit(g: *Generation, chunk: []const u32) Failure!void { + var landed: usize = 0; + for (chunk) |t| { + if (g.reason != null) break; + try g.collected.append(g.a, t); + landed += 1; + if (g.eos(t)) { + g.reason = "stop"; + } else if (try g.stopHit()) { + g.reason = "stop"; + g.srv.engine.cancel(g.id); // the engine ends EOS and length itself; a stop string is the server's + } else if (g.collected.items.len >= g.max_tokens) g.reason = "length"; + } + if (landed == 0) return; + const arrived = nowNs(g.srv.io); + if (g.first_ns == null) g.first_ns = arrived; + g.last_ns = arrived; + const sink = g.sink orelse return; + if (g.streaming_done) return; + var fresh: std.ArrayList(u32) = .empty; + for (chunk[0..landed]) |t| { + if (g.eos(t)) { + g.streaming_done = true; + break; + } + try fresh.append(g.a, t); + } + const all = try g.visible.extend(g.a, g.srv.text, decodeText, fresh.items); + const text = g.stops.visible(all, true); + if (g.visible.pending()) return; // a character still split across tokens + var answer = text; + if (g.thinking) { + const split = try reply_text.splitThinking(g.a, text, false, g.srv.markers); + const piece = g.streamed_reasoning.after(split.reasoning, true); + if (piece.len > 0) { + try g.streamed_reasoning.set(g.a, split.reasoning, piece); + try g.emit(sink, try deltaOf(g.a, "reasoning_content", piece)); + } + answer = split.answer; + } + const shown = if (g.calls != null) try reply_text.hideInto(g.a, &g.hidden, answer, false) else answer; + const vis = reply_text.heldBack(reply_text.streamingVisible(shown), g.calls != null); + const copied = @intFromPtr(vis.ptr) >= @intFromPtr(g.hidden.items.ptr) and @intFromPtr(vis.ptr) <= @intFromPtr(g.hidden.items.ptr) + g.hidden.items.len; + const delta = g.streamed.after(vis, !copied); + if (delta.len > 0) { + try g.streamed.set(g.a, vis, delta); + try g.emit(sink, if (g.lp) |*l| (try l.delta(g.srv, g.a, g.thinking, g.ignore_eos, g.collected.items, delta, g.streamed.text.items.len, false)).? else .{ .string = delta }); + } + if (g.calls) |*c| { + var out: std.ArrayList(Value) = .empty; + try c.feed(answer, &out); // never the reasoning: a call it mentions is not made + for (out.items) |d| try g.emit(sink, d); + } + } + + fn emit(g: *Generation, sink: Sink, delta: Value) Failure!void { + sink.call(sink.ctx, delta) catch { + g.cancel(); + return error.Cancelled; + }; + g.wrote_ns = nowNs(g.srv.io); + } + + /// Ends the request and waits for the engine; one it ended unfinished counts as a disconnect, as the Mac scheduler counts it. + fn cancel(g: *Generation) void { + if (!g.engine_done) g.srv.engine.cancel(g.id); + g.drain(); + if (g.box.reason == .cancelled and g.reason == null) g.srv.metrics.disconnected(g.srv.io); + } + + /// Waits for the engine's own end, so the request it holds may be freed. + fn drain(g: *Generation) void { + const m = g.box; + m.mutex.lockUncancelable(m.io); + defer m.mutex.unlock(m.io); + while (!m.finished) m.cond.waitUncancelable(m.io, &m.mutex); + g.engine_done = true; + } + + /// Takes rounds until the reply ends; a client that leaves cancels it. + fn loop(g: *Generation, gone: anytype) Failure!void { + const m = g.box; + while (true) { + m.mutex.lockUncancelable(m.io); + var chunk: ?[]u32 = null; + var ended = false; + if (g.chunk_index < m.chunks.items.len) { + const end = m.chunks.items[g.chunk_index]; + g.chunk_index += 1; + chunk = g.a.dupe(u32, m.tokens.items[g.consumed..end]) catch null; + if (g.lp) |*l| l.rows.appendSlice(g.a, m.rows.items[g.consumed..end]) catch {}; + g.consumed = end; + } else if (m.finished) { + ended = true; + } else { + m.cond.waitTimeout(m.io, &m.mutex, .{ .duration = .{ .raw = .fromMilliseconds(50), .clock = .awake } }) catch {}; + } + m.mutex.unlock(m.io); + if (ended) { + g.engine_done = true; + return; + } + if (gone.check()) { + g.cancel(); + return error.Cancelled; + } + if (chunk == null and @hasDecl(@TypeOf(gone), "beat")) { + const now = nowNs(g.srv.io); + const since = g.wrote_ns orelse now; + if (g.wrote_ns == null) g.wrote_ns = now; + if (gone.beat(now - since) catch { + g.cancel(); + return error.Cancelled; + }) g.wrote_ns = now; + } + const tokens = chunk orelse continue; + if (g.reason == null) try g.commit(tokens); + } + } + + fn finish(g: *Generation, cx: *Cx, prompt_len: usize, received: i96, thinking: bool, effort: ?[]const u8, exact: bool, drafts: bool) Failure!Reply { + const a = g.a; + const m = g.box; + if (!g.engine_done) g.drain(); + const reason: []const u8 = g.reason orelse switch (m.reason) { + .stop => "stop", + .length => "length", + .cancelled => "cancelled", + .failed => return cx.other(if (m.message.len > 0) try a.dupe(u8, m.message) else "the reply failed"), + }; + const finished_ns = nowNs(g.srv.io); + const content_tokens = if (g.ignore_eos) g.collected.items else reply_text.stripTrailing(g.collected.items, g.srv.eos); + const raw = try g.srv.text.decode(a, content_tokens); + const visible = g.stops.visible(raw, false); + // closed before the think split, so a call ending an unclosed think block is the answer + const close = try g.closeCall(visible); + const text = if (close.len > 0) try std.mem.concat(a, u8, &.{ visible, close }) else visible; + var content: []const u8 = text; + var reasoning: ?[]const u8 = null; + if (g.thinking) { + const split = try reply_text.splitThinking(a, text, true, g.srv.markers); + const r = reply_text.pyStrip(split.reasoning); + reasoning = if (r.len > 0) r else null; + content = split.answer; + } else { + const h = reply_text.parseHarmony(text); + content = h.content; + reasoning = h.reasoning; + } + if (g.sink) |sink| { + const thought = g.streamed_reasoning.text.items; + if (reasoning) |r| if (std.mem.startsWith(u8, r, thought) and r.len > thought.len) + try g.emit(sink, try deltaOf(a, "reasoning_content", r[thought.len..])); + const shown = if (g.calls != null) try reply_text.finishedProse(a, content) else content; + const sent = g.streamed.text.items; + const rest = if (std.mem.startsWith(u8, shown, sent) and shown.len > sent.len) shown[sent.len..] else ""; + if (g.lp) |*l| (if (try l.delta(g.srv, a, g.thinking, g.ignore_eos, g.collected.items, rest, shown.len, true)) |d| try g.emit(sink, d)) else if (rest.len > 0) try g.emit(sink, .{ .string = rest }); + if (g.calls) |*c| if (close.len > 0) { + // the streamer reads the closers as markup the model wrote, so the call ends with the same deltas + var out: std.ArrayList(Value) = .empty; + try c.feed(if (g.thinking) content else text, &out); + for (out.items) |d| try g.emit(sink, d); + }; + } + const prefilled = m.prefilled_ns; + const decode_s: f64 = if (prefilled) |p| @max(0, seconds(finished_ns - p)) else 0; + const decode_tokens = g.collected.items.len -| 1; + const runtime = try json.newObject(a); + try runtime.put(a, "enable_thinking", .{ .bool = thinking }); + try runtime.put(a, "reasoning_effort", if (!thinking) .{ .string = "none" } else if (effort) |e| .{ .string = e } else .null); + try runtime.put(a, "engine", .{ .string = g.srv.info.name }); + try runtime.put(a, "tokens_per_second", .{ .float = if (decode_s > 0) @as(f64, @floatFromInt(decode_tokens)) / decode_s else 0 }); + try runtime.put(a, "seconds", .{ .float = @max(0, decode_s) }); // the reply's decode seconds, as tokens_per_second divides them + try runtime.put(a, "prefill_seconds", if (m.stats.prefill_seconds) |p| .{ .float = p } else .null); + const widths = try a.alloc(Value, m.widths.len); + for (m.widths, widths) |w, *slot| slot.* = try json.intValue(a, w); + try runtime.put(a, "prefill_widths", .{ .array = widths }); + const raised = try a.alloc(Value, m.raised.len); + for (m.raised, raised) |r, *slot| slot.* = .{ .bool = r }; + try runtime.put(a, "prefill_raised", .{ .array = raised }); + try runtime.put(a, "time_to_first_token", if (g.first_ns) |t| .{ .float = seconds(t - received) } else .null); + try runtime.put(a, "sampling", .{ .string = if (exact) "exact" else "greedy" }); + try runtime.put(a, "drafts", .{ .bool = drafts }); + const sha = @import("tokens.zig").tokenSha(g.collected.items); + try runtime.put(a, "token_sha", .{ .string = try a.dupe(u8, &sha) }); + try runtime.put(a, "min_rows", try json.intValue(a, m.stats.min_rows)); + if (m.stats.loop_period) |period| { + const loop_field = try json.newObject(a); + try loop_field.put(a, "period", try json.intValue(a, period)); + try runtime.put(a, "loop", .{ .object = loop_field }); + } + const s = m.stats; + const spec = try json.newObject(a); + try spec.put(a, "rounds", try json.intValue(a, s.rounds)); + try spec.put(a, "drafted", try json.intValue(a, s.drafted)); + try spec.put(a, "accepted", try json.intValue(a, s.accepted)); + try spec.put(a, "acceptance_rate", .{ .float = if (s.drafted > 0) @as(f64, @floatFromInt(s.accepted)) / @as(f64, @floatFromInt(s.drafted)) else 0 }); + try spec.put(a, "tokens_per_round", .{ .float = if (s.rounds > 0) @as(f64, @floatFromInt(g.collected.items.len)) / @as(f64, @floatFromInt(s.rounds)) else 0 }); + if (m.telemetry.len > 0) if ((try json.parse(a, m.telemetry)) == .ok) try spec.put(a, "proposer", (try json.parse(a, m.telemetry)).ok); + const think_end: ?u32 = if (thinking) g.srv.text.tokenId(g.srv.markers.close) else null; + const reply: Reply = .{ + .content = content, + .stop_sequence = g.stops.matched(raw), + .reasoning = reasoning, + .tool_calls_streamed = g.calls != null and g.calls.?.streamed, + .finish_reason = reason, + .prompt_tokens = prompt_len, + .cached_tokens = m.cached orelse 0, + .completion_tokens = g.collected.items.len, + .reasoning_tokens = reply_text.reasoningCount(g.collected.items, think_end), + .runtime = .{ .object = runtime }, + .speculative = .{ .object = spec }, + .thinking = thinking, + .effort = effort, + .logprobs = if (g.lp) |l| try @import("logprobs.zig").value(g.srv, cx, a, g.thinking, content_tokens, l.rows.items) else null, + .prompt_logprobs = if (m.prompt_rows) |r| try @import("logprobs.zig").promptValue(g.srv, a, r) else null, + }; + if (std.mem.eql(u8, reason, "length") and thinking and reply_text.pyStrip(content).len == 0) + log.line("warning: a reply reached max_tokens while still thinking, so its content is empty and its text is all in reasoning_content; raise max_tokens, or send chat_template_kwargs {{\"enable_thinking\": false}} (server: --no-thinking)", .{}); + var cycle_text: [32]u8 = undefined; + const cycle = if (s.loop_period) |period| std.fmt.bufPrint(&cycle_text, " loop=period:{d}", .{period}) catch "" else ""; + log.line("done {s} prompt={d} cached={d} thinking={s} effort={s} tokens={d} sha={s} finish={s}{s} rounds={d} accepted={d}/{d}", .{ g.reply_id, prompt_len, reply.cached_tokens, if (thinking) "True" else "False", if (thinking) effort orelse "none" else "none", g.collected.items.len, sha, reason, cycle, s.rounds, s.accepted, s.drafted }); + return reply; + } +}; + /// The request's stop strings as the engine checks them after each token: the newest tokens' text holds one. const StopHook = struct { srv: *Server, @@ -267,21 +552,7 @@ const StopHook = struct { } }; -/// What closes a call the model's end token left open (``tool_parse.closeCall``); nothing when a stop string, the length or ``ignore_eos`` ended the reply instead. -pub fn closeCall(g: *const Generation, text: []const u8) Allocator.Error![]const u8 { - const t = g.collected.items; - if (g.tools.len == 0 or t.len == 0 or !g.eos(t[t.len - 1])) return ""; - return tool_parse.closeCall(g.a, text, g.tools); -} - -/// Ends the request and waits for the engine; one it ended unfinished counts as a disconnect, as the Mac scheduler counts it. -pub fn cancel(g: *Generation) void { - if (!g.engine_done) g.srv.engine.cancel(g.id); - g.drain(); - if (g.box.reason == .cancelled and g.reason == null) g.srv.metrics.disconnected(g.srv.io); -} - -pub fn decodeText(t: anytype, a: Allocator, tokens: []const u32) ![]u8 { +fn decodeText(t: anytype, a: Allocator, tokens: []const u32) ![]u8 { return t.decode(a, tokens); } diff --git a/zig/src/server/chat_generation.zig b/zig/src/server/chat_generation.zig deleted file mode 100644 index 0b4e164c98..0000000000 --- a/zig/src/server/chat_generation.zig +++ /dev/null @@ -1,452 +0,0 @@ -//! A submitted reply's collection: the engine's events in, streamed deltas and the finished ``Reply`` out. -const std = @import("std"); -const api = @import("engine_api"); -const json = @import("json.zig"); -const errors = @import("errors.zig"); -const reply_text = @import("reply_text.zig"); -const tool_stream = @import("tool_stream.zig"); -const tool_parse = @import("tool_parse.zig"); -const log = @import("log.zig"); -const clock = @import("clock.zig"); -const Server = @import("server.zig").Server; -const family_mod = @import("family.zig"); -const chunk_plan = @import("chunk_plan.zig"); -const Value = json.Value; -const Cx = errors.Cx; -const Allocator = std.mem.Allocator; -const chat = @import("chat.zig"); -const Sink = chat.Sink; -const Reply = chat.Reply; -const Failure = chat.Failure; -const Mailbox = @import("mailbox.zig").Mailbox; -const decodeText = chat.decodeText; -const deltaOf = chat.deltaOf; -const nowNs = clock.nowNs; -const seconds = clock.seconds; - -/// The text a stream has sent: Python's ``streamed = visible``, grown by its delta while it only grows. -pub const Shown = struct { - text: std.ArrayList(u8) = .empty, - chars: usize = 0, - at: ?[*]const u8 = null, // where the last shown text lay; the decode buffer only grows, so that prefix holds - extends: bool = false, - - /// Python's ``now[len(shown):]``; ``stable``: ``now`` lies in the append-only decode buffer. - fn after(s: *Shown, now: []const u8, stable: bool) []const u8 { - const sent = s.text.items; - s.extends = (stable and s.at == now.ptr and now.len >= sent.len) or std.mem.startsWith(u8, now, sent); - if (!s.extends) return reply_text.afterChars(now, s.chars); - if (stable) s.at = now.ptr; // the buffer moved but kept its bytes: the next check is free again - return now[sent.len..]; - } - - fn set(s: *Shown, a: Allocator, now: []const u8, delta: []const u8) Allocator.Error!void { - if (s.extends) { - try s.text.appendSlice(a, delta); - s.chars += reply_text.charCount(delta); - } else { - s.text = .empty; - try s.text.appendSlice(a, now); - s.chars = reply_text.charCount(now); - } - s.at = now.ptr; - } -}; - -/// The token loop and the text it streams. -pub const Generation = struct { - srv: *Server, - a: Allocator, - box: *Mailbox, - id: api.Id, - reply_id: []const u8, // the id the client gets, which the done line prints - sink: ?Sink, - thinking: bool, - stops: reply_text.Stops, - ignore_eos: bool, - max_tokens: u32, - tools: []const Value, - calls: ?tool_stream.Streamer = null, - collected: std.ArrayList(u32) = .empty, - visible: reply_text.Incremental = .{}, - hidden: std.ArrayList(u8) = .empty, // reused for the answer without its call blocks - streamed: Shown = .{}, // what content streamed, kept apart from the decode buffer that grows under slices - streamed_reasoning: Shown = .{}, - streaming_done: bool = false, - first_ns: ?i96 = null, - last_ns: ?i96 = null, - /// when the stream last carried bytes (a keepalive goes out after a silence) - wrote_ns: ?i96 = null, - reason: ?[]const u8 = null, // the server ended the reply (a stop string, the length) before the engine said so - engine_done: bool = false, - consumed: usize = 0, - chunk_index: usize = 0, - lp: ?@import("logprobs.zig").Stream = null, // the request asked for logprobs: its rows, copied with each round - // a family's reply: its reader, the answer so far, how much of it streamed, and a stop string's hit - reader: ?family_mod.Reader = null, - answer: std.ArrayList(u8) = .empty, - answer_sent: usize = 0, - deltas: std.ArrayList(Value) = .empty, - stop_hit: bool = false, - stop_at: usize = 0, - - pub fn eos(g: *const Generation, t: u32) bool { - return !g.ignore_eos and std.mem.indexOfScalar(u32, g.srv.eos, t) != null; - } - - const closeCall = chat.closeCall; // in chat.zig with its doc - - /// The engine's stop check, here: the newest tokens' text holds a stop string. - fn stopHit(g: *Generation) Allocator.Error!bool { - if (g.stops.strings.len == 0) return false; - const t = g.collected.items; - const tail = t[t.len -| g.stops.tail()..]; - const text = try g.srv.text.decode(g.a, tail); - for (g.stops.strings) |s| if (std.mem.indexOf(u8, text, s) != null) return true; - return false; - } - - /// Commits a round's tokens as ``LaneStream.commit`` would, then streams what they add. - fn commit(g: *Generation, chunk: []const u32) Failure!void { - if (g.reader) |r| return g.commitFamily(r, chunk); - var landed: usize = 0; - for (chunk) |t| { - if (g.reason != null) break; - try g.collected.append(g.a, t); - landed += 1; - if (g.eos(t)) { - g.reason = "stop"; - } else if (try g.stopHit()) { - g.reason = "stop"; - g.srv.engine.cancel(g.id); // the engine ends EOS and length itself; a stop string is the server's - } else if (g.collected.items.len >= g.max_tokens) g.reason = "length"; - } - if (landed == 0) return; - const arrived = nowNs(g.srv.io); - if (g.first_ns == null) g.first_ns = arrived; - g.last_ns = arrived; - const sink = g.sink orelse return; - if (g.streaming_done) return; - var fresh: std.ArrayList(u32) = .empty; - for (chunk[0..landed]) |t| { - if (g.eos(t)) { - g.streaming_done = true; - break; - } - try fresh.append(g.a, t); - } - const all = try g.visible.extend(g.a, g.srv.text, decodeText, fresh.items); - const text = g.stops.visible(all, true); - if (g.visible.pending()) return; // a character still split across tokens - var answer = text; - if (g.thinking) { - const split = try reply_text.splitThinking(g.a, text, false, g.srv.markers); - const piece = g.streamed_reasoning.after(split.reasoning, true); - if (piece.len > 0) { - try g.streamed_reasoning.set(g.a, split.reasoning, piece); - try g.emit(sink, try deltaOf(g.a, "reasoning_content", piece)); - } - answer = split.answer; - } - const shown = if (g.calls != null) try reply_text.hideInto(g.a, &g.hidden, answer, false) else answer; - const vis = reply_text.heldBack(reply_text.streamingVisible(shown), g.calls != null); - const copied = @intFromPtr(vis.ptr) >= @intFromPtr(g.hidden.items.ptr) and @intFromPtr(vis.ptr) <= @intFromPtr(g.hidden.items.ptr) + g.hidden.items.len; - const delta = g.streamed.after(vis, !copied); - if (delta.len > 0) { - try g.streamed.set(g.a, vis, delta); - try g.emit(sink, if (g.lp) |*l| (try l.delta(g.srv, g.a, g.thinking, g.ignore_eos, g.collected.items, delta, g.streamed.text.items.len, false)).? else .{ .string = delta }); - } - if (g.calls) |*c| { - var out: std.ArrayList(Value) = .empty; - try c.feed(answer, &out); // never the reasoning: a call it mentions is not made - for (out.items) |d| try g.emit(sink, d); - } - } - - fn emit(g: *Generation, sink: Sink, delta: Value) Failure!void { - sink.call(sink.ctx, delta) catch { - g.cancel(); - return error.Cancelled; - }; - g.wrote_ns = nowNs(g.srv.io); - } - - pub const cancel = chat.cancel; // in chat.zig with its doc - - /// Waits for the engine's own end, so the request it holds may be freed. - pub fn drain(g: *Generation) void { - const m = g.box; - m.mutex.lockUncancelable(m.io); - defer m.mutex.unlock(m.io); - while (!m.finished) m.cond.waitUncancelable(m.io, &m.mutex); - g.engine_done = true; - } - - /// Takes rounds until the reply ends; a client that leaves cancels it. - pub fn loop(g: *Generation, gone: anytype) Failure!void { - const m = g.box; - while (true) { - m.mutex.lockUncancelable(m.io); - var chunk: ?[]u32 = null; - var ended = false; - if (g.chunk_index < m.chunks.items.len) { - const end = m.chunks.items[g.chunk_index]; - g.chunk_index += 1; - chunk = g.a.dupe(u32, m.tokens.items[g.consumed..end]) catch null; - if (g.lp) |*l| l.rows.appendSlice(g.a, m.rows.items[g.consumed..end]) catch {}; - g.consumed = end; - } else if (m.finished) { - ended = true; - } else { - m.cond.waitTimeout(m.io, &m.mutex, .{ .duration = .{ .raw = .fromMilliseconds(50), .clock = .awake } }) catch {}; - } - m.mutex.unlock(m.io); - if (ended) { - g.engine_done = true; - return; - } - if (gone.check()) { - g.cancel(); - return error.Cancelled; - } - if (chunk == null and @hasDecl(@TypeOf(gone), "beat")) { - // queueing and a long prefill send nothing: a client that times out a silent stream would leave - const now = nowNs(g.srv.io); - const since = g.wrote_ns orelse now; - if (g.wrote_ns == null) g.wrote_ns = now; - if (gone.beat(now - since) catch { - g.cancel(); - return error.Cancelled; - }) g.wrote_ns = now; - } - const tokens = chunk orelse continue; - if (g.reason == null) try g.commit(tokens); - } - } - - /// A family's round: committed tokens through the family's reader, its deltas pushed (an EOS id ends, unshown). - fn commitFamily(g: *Generation, r: family_mod.Reader, chunk: []const u32) Failure!void { - var landed: usize = 0; - const start = g.collected.items.len; - for (chunk) |t| { - if (g.reason != null) break; - try g.collected.append(g.a, t); - landed += 1; - if (g.eos(t)) { - g.reason = "stop"; - } else if (g.collected.items.len >= g.max_tokens) g.reason = "length"; - } - if (landed == 0) return; - const arrived = nowNs(g.srv.io); - if (g.first_ns == null) g.first_ns = arrived; - g.last_ns = arrived; - for (g.collected.items[start..]) |t| if (!g.isEos(t)) try r.push(&.{t}); - if (r.pending()) return; // a character still split across tokens - g.deltas.clearRetainingCapacity(); - try r.feed(false, &g.deltas); - try g.pushFamily(g.deltas.items, false); - } - - /// Any end-of-sentence id, ``ignore_eos`` or not: never text. - fn isEos(g: *const Generation, t: u32) bool { - return std.mem.indexOfScalar(u32, g.srv.eos, t) != null; - } - - /// ``push``: reasoning and calls go out as they come; the answer is held where a stop may begin, cut where it does. - fn pushFamily(g: *Generation, deltas: []const Value, finished: bool) Failure!void { - for (deltas) |d| switch (d) { - .string => |t| try g.answer.appendSlice(g.a, t), - else => if (g.sink) |sink| if (!g.streaming_done or finished) try g.emit(sink, d), - }; - var full = g.answer.items; - const stops = g.stops.strings; - if (stops.len > 0 and !g.stop_hit) { - var longest: usize = 0; - for (stops) |st| longest = @max(longest, st.len); - const from = (g.answer_sent + 1) -| longest; - var cut: ?usize = null; - for (stops) |st| if (std.mem.indexOfPos(u8, full, @min(from, full.len), st)) |at| { - if (cut == null or at < cut.?) cut = at; - }; - if (cut) |c| { - g.answer.shrinkRetainingCapacity(c); - full = g.answer.items; - g.stop_hit = true; - g.stop_at = g.collected.items.len; - if (g.reason == null) { - g.reason = "stop"; - g.srv.engine.cancel(g.id); // the engine ends EOS and length itself; a stop string is the server's - } - } - } - var upto = full.len; - if (!finished and !g.stop_hit) { - var hold: usize = 0; - for (stops) |st| { - var k = @min(st.len - 1, full.len); - while (k > 0) : (k -= 1) if (std.mem.endsWith(u8, full, st[0..k])) break; - hold = @max(hold, k); - } - upto -= hold; - } - if (upto > g.answer_sent) { - if (g.sink) |sink| try g.emit(sink, .{ .string = full[g.answer_sent..upto] }); - g.answer_sent = upto; - } - } - - /// A family's finished reply: the stream's last deltas, then the parts (``app._run``'s end). - fn finishFamily(g: *Generation, r: family_mod.Reader) Failure!struct { content: []const u8, reasoning: ?[]const u8, calls: []Value, reason: []const u8, used: usize } { - const streaming = g.sink != null; - g.deltas.clearRetainingCapacity(); - try r.feed(true, &g.deltas); - try g.pushFamily(g.deltas.items, true); - const parsed = try r.parse(); - var content: []const u8 = undefined; - var calls: []Value = undefined; - if (streaming) { - calls = try r.calls(); - content = g.answer.items[0..g.answer_sent]; - } else { - calls = parsed.calls; - content = parsed.content; - for (g.stops.strings) |st| if (std.mem.indexOf(u8, content, st)) |at| { - content = content[0..at]; - if (!g.stop_hit) { - g.stop_hit = true; - g.stop_at = g.collected.items.len; - } - }; - } - const t = g.collected.items; - const reason: []const u8 = if (calls.len > 0) "tool_calls" else if (g.stop_hit or (!g.ignore_eos and t.len > 0 and g.isEos(t[t.len - 1]))) "stop" else "length"; - if (g.srv.family) |fam| fam.remember(parsed.reasoning, content, calls); - // tools offered, no call, but tool markup: log up to 300 bytes of what the parser saw (UTF-8 cut, escaped) - if (calls.len == 0 and g.tools.len > 0) if (r.markup() catch null) |m| - std.log.warn("tool markup with no call: request {s}, finish {s}: {s}; reply bytes {d}..: {f}", .{ g.reply_id, reason, m.reason, m.at, std.zig.fmtString(m.window) }); - return .{ .content = content, .reasoning = if (parsed.reasoning.len > 0) parsed.reasoning else null, .calls = calls, .reason = reason, .used = if (g.stop_hit and g.stop_at > 0) g.stop_at else t.len }; - } - - pub fn finish(g: *Generation, cx: *Cx, prompt_len: usize, received: i96, thinking: bool, effort: ?[]const u8, exact: bool, drafts: bool, images_omitted: u32) Failure!Reply { - const a = g.a; - const m = g.box; - if (!g.engine_done) g.drain(); - const reason_generic: []const u8 = g.reason orelse switch (m.reason) { - .stop => "stop", - .length => "length", - .cancelled => "cancelled", - .failed => return cx.other(if (m.message.len > 0) try a.dupe(u8, m.message) else "the reply failed"), - }; - var reason = reason_generic; - var raw: []const u8 = ""; - var content: []const u8 = undefined; - var reasoning: ?[]const u8 = null; - var calls: ?[]Value = null; - var used = g.collected.items.len; - var logprobs: ?Value = null; - if (g.reader) |r| { - const f = try g.finishFamily(r); - content = f.content; - reasoning = f.reasoning; - calls = f.calls; - used = f.used; - if (!(m.reason == .cancelled and g.reason == null)) reason = f.reason; - } else { - const content_tokens = if (g.ignore_eos) g.collected.items else reply_text.stripTrailing(g.collected.items, g.srv.eos); - raw = try g.srv.text.decode(a, content_tokens); - if (g.lp) |l| logprobs = try @import("logprobs.zig").value(g.srv, cx, a, g.thinking, content_tokens, l.rows.items); - const visible = g.stops.visible(raw, false); - // closed before the think split, so a call ending an unclosed think block is the answer - const close = try g.closeCall(visible); - const text = if (close.len > 0) try std.mem.concat(a, u8, &.{ visible, close }) else visible; - content = text; - if (g.thinking) { - const split = try reply_text.splitThinking(a, text, true, g.srv.markers); - const r = reply_text.pyStrip(split.reasoning); - reasoning = if (r.len > 0) r else null; - content = split.answer; - } else { - const h = reply_text.parseHarmony(text); - content = h.content; - reasoning = h.reasoning; - } - if (g.sink) |sink| { - const thought = g.streamed_reasoning.text.items; - if (reasoning) |r| if (std.mem.startsWith(u8, r, thought) and r.len > thought.len) - try g.emit(sink, try deltaOf(a, "reasoning_content", r[thought.len..])); - const shown = if (g.calls != null) try reply_text.finishedProse(a, content) else content; - const sent = g.streamed.text.items; - const rest = if (std.mem.startsWith(u8, shown, sent) and shown.len > sent.len) shown[sent.len..] else ""; - if (g.lp) |*l| (if (try l.delta(g.srv, a, g.thinking, g.ignore_eos, g.collected.items, rest, shown.len, true)) |d| try g.emit(sink, d)) else if (rest.len > 0) try g.emit(sink, .{ .string = rest }); - if (g.calls) |*c| if (close.len > 0) { - // the streamer reads the closers as markup the model wrote, so the call ends with the same deltas - var out: std.ArrayList(Value) = .empty; - try c.feed(if (g.thinking) content else text, &out); - for (out.items) |d| try g.emit(sink, d); - }; - } - } - const finished_ns = nowNs(g.srv.io); - const prefilled = m.prefilled_ns; - const decode_s: f64 = if (prefilled) |p| @max(0, seconds(finished_ns - p)) else 0; - const decode_tokens = g.collected.items.len -| 1; - const runtime = try json.newObject(a); - try runtime.put(a, "enable_thinking", .{ .bool = thinking }); - try runtime.put(a, "reasoning_effort", if (!thinking) .{ .string = "none" } else if (effort) |e| .{ .string = e } else .null); - try runtime.put(a, "engine", .{ .string = g.srv.info.name }); - try runtime.put(a, "tokens_per_second", .{ .float = if (decode_s > 0) @as(f64, @floatFromInt(decode_tokens)) / decode_s else 0 }); - try runtime.put(a, "seconds", .{ .float = @max(0, decode_s) }); // the reply's decode seconds, as tokens_per_second divides them - try runtime.put(a, "prefill_seconds", if (m.stats.prefill_seconds) |p| .{ .float = p } else .null); - const widths = try a.alloc(Value, m.widths.len); - for (m.widths, widths) |w, *slot| slot.* = try json.intValue(a, w); - try runtime.put(a, "prefill_widths", .{ .array = widths }); - const raised = try a.alloc(Value, m.raised.len); - for (m.raised, raised) |r, *slot| slot.* = .{ .bool = r }; - try runtime.put(a, "prefill_raised", .{ .array = raised }); - try runtime.put(a, "time_to_first_token", if (g.first_ns) |t| .{ .float = seconds(t - received) } else .null); - try runtime.put(a, "sampling", .{ .string = if (exact) "exact" else "greedy" }); - try runtime.put(a, "drafts", .{ .bool = drafts }); - const sha = @import("tokens.zig").tokenSha(g.collected.items); - try runtime.put(a, "token_sha", .{ .string = try a.dupe(u8, &sha) }); - if (images_omitted > 0) try runtime.put(a, "images_omitted", try json.intValue(a, images_omitted)); - try runtime.put(a, "min_rows", try json.intValue(a, m.stats.min_rows)); - if (m.stats.loop_period) |period| { - const loop_field = try json.newObject(a); - try loop_field.put(a, "period", try json.intValue(a, period)); - try runtime.put(a, "loop", .{ .object = loop_field }); - } - const s = m.stats; - const spec = try json.newObject(a); - try spec.put(a, "rounds", try json.intValue(a, s.rounds)); - try spec.put(a, "drafted", try json.intValue(a, s.drafted)); - try spec.put(a, "accepted", try json.intValue(a, s.accepted)); - try spec.put(a, "acceptance_rate", .{ .float = if (s.drafted > 0) @as(f64, @floatFromInt(s.accepted)) / @as(f64, @floatFromInt(s.drafted)) else 0 }); - try spec.put(a, "tokens_per_round", .{ .float = if (s.rounds > 0) @as(f64, @floatFromInt(g.collected.items.len)) / @as(f64, @floatFromInt(s.rounds)) else 0 }); - if (m.telemetry.len > 0) if ((try json.parse(a, m.telemetry)) == .ok) try spec.put(a, "proposer", (try json.parse(a, m.telemetry)).ok); - const think_end: ?u32 = if (thinking) g.srv.text.tokenId(g.srv.markers.close) else null; - const reply: Reply = .{ - .content = content, - .stop_sequence = g.stops.matched(raw), - .reasoning = reasoning, - .tool_calls_streamed = g.calls != null and g.calls.?.streamed, - .finish_reason = reason, - .prompt_tokens = prompt_len, - .cached_tokens = m.cached orelse 0, - .completion_tokens = used, - .reasoning_tokens = reply_text.reasoningCount(g.collected.items, think_end), - .runtime = .{ .object = runtime }, - .speculative = .{ .object = spec }, - .thinking = thinking, - .effort = effort, - .logprobs = logprobs, - .prompt_logprobs = if (m.prompt_rows) |r| try @import("logprobs.zig").promptValue(g.srv, a, r) else null, - .calls = calls, - }; - if (std.mem.eql(u8, reason, "length") and thinking and reply_text.pyStrip(content).len == 0) - log.line("warning: a reply reached max_tokens while still thinking, so its content is empty and its text is all in reasoning_content; raise max_tokens, or send chat_template_kwargs {{\"enable_thinking\": false}} (server: --no-thinking)", .{}); - var cycle_text: [32]u8 = undefined; - const cycle = if (s.loop_period) |period| std.fmt.bufPrint(&cycle_text, " loop=period:{d}", .{period}) catch "" else ""; - log.line("done {s} prompt={d} cached={d} thinking={s} effort={s} tokens={d} sha={s} finish={s}{s} rounds={d} accepted={d}/{d}", .{ g.reply_id, prompt_len, reply.cached_tokens, if (thinking) "True" else "False", if (thinking) effort orelse "none" else "none", g.collected.items.len, sha, reason, cycle, s.rounds, s.accepted, s.drafted }); - return reply; - } -}; diff --git a/zig/src/server/family.zig b/zig/src/server/family.zig deleted file mode 100644 index 9e24c8dab4..0000000000 --- a/zig/src/server/family.zig +++ /dev/null @@ -1,135 +0,0 @@ -//! A model family's own chat rules (prompt encoding, effort budgets, call markup) in place of the chat template. -const std = @import("std"); -const json = @import("json.zig"); -const errors = @import("errors.zig"); -const Value = json.Value; -const Cx = errors.Cx; -const Allocator = std.mem.Allocator; - -/// A request's thinking switch and effort as it asked for them (null: the server's default). -pub const Thinking = struct { enable: ?bool = null, effort: ?[]const u8 = null }; - -/// What a chat prompt is rendered from: the request's own messages and tools, as the client sent them. -pub const Prompt = struct { - messages: Value, - tools: []const Value, - thinking: bool, - /// The effort label (``Family.thinking``'s, else the default). - effort: []const u8, - generation: bool = true, - /// The request body, for fields the family reads itself (``response_format``). - body: Value = .null, -}; - -pub const Rendered = struct { ids: []const u32, images_omitted: u32 = 0, images: ?Images = null }; - -/// A prompt's prepared images: ``held`` lives until ``release``; ``token`` for /tokenize. -pub const Images = struct { - held: *const anyopaque, - token: u32, - ctx: *anyopaque, - release_fn: *const fn (ctx: *anyopaque) void, - - pub fn release(im: Images) void { - im.release_fn(im.ctx); - } -}; - -/// A finished reply's parts, calls as OpenAI tool-call objects. -pub const Parsed = struct { reasoning: []const u8, content: []const u8, calls: []Value }; - -/// One reply as its tokens arrive: the text so far, and the deltas it adds. -pub const Reader = struct { - ctx: *anyopaque, - vt: *const VTable, - - pub const VTable = struct { - /// Tokens committed (end-of-sentence ids already left out). - push: *const fn (ctx: *anyopaque, ids: []const u32) Allocator.Error!void, - /// A character is split across tokens: nothing new is readable yet. - pending: *const fn (ctx: *anyopaque) bool, - /// Deltas since the last feed (content a string, reasoning and calls delta objects); ``finished`` flushes. - feed: *const fn (ctx: *anyopaque, finished: bool, out: *std.ArrayList(Value)) Allocator.Error!void, - /// The whole reply parsed (after the last ``push``; a split character at the end reads as U+FFFD). - parse: *const fn (ctx: *anyopaque) Allocator.Error!Parsed, - /// The calls the stream sent, in order, as OpenAI tool-call objects. - calls: *const fn (ctx: *anyopaque) Allocator.Error![]Value, - /// The finished reply's tool markup that yielded no call (a reason, at most 300 bytes) for the log; null: none - markup: ?*const fn (ctx: *anyopaque) Allocator.Error!?Markup = null, - }; - - pub fn push(r: Reader, ids: []const u32) Allocator.Error!void { - return r.vt.push(r.ctx, ids); - } - pub fn pending(r: Reader) bool { - return r.vt.pending(r.ctx); - } - pub fn feed(r: Reader, finished: bool, out: *std.ArrayList(Value)) Allocator.Error!void { - return r.vt.feed(r.ctx, finished, out); - } - pub fn parse(r: Reader) Allocator.Error!Parsed { - return r.vt.parse(r.ctx); - } - pub fn calls(r: Reader) Allocator.Error![]Value { - return r.vt.calls(r.ctx); - } - pub fn markup(r: Reader) Allocator.Error!?Markup { - const f = r.vt.markup orelse return null; - return f(r.ctx); - } -}; - -/// A reply's tool markup that the parser turned into no call (``Reader.markup``). -pub const Markup = struct { at: usize, reason: []const u8, window: []const u8 }; - -/// Where a reply's reasoning goes: ``reasoning_content`` (OpenAI's), ``reasoning`` (vLLM's), or both. -pub const ReasoningFields = enum { both, reasoning_content, reasoning }; - -pub const Family = struct { - ctx: *anyopaque, - vt: *const VTable, - - pub const VTable = struct { - name: []const u8, - /// ``chat_template_kwargs`` and ``reasoning_effort`` read; a bad value refuses the request. - thinking: *const fn (ctx: *anyopaque, cx: *Cx, body: Value) errors.Refused!Thinking, - /// The effort a request that names none gets. - default_effort: *const fn (ctx: *anyopaque) []const u8, - /// The tools the prompt offers: ``tool_choice`` none offers none; otherwise all of them. - tools: *const fn (ctx: *anyopaque, cx: *Cx, body: Value) errors.Refused![]const Value, - render: *const fn (ctx: *anyopaque, cx: *Cx, p: Prompt) errors.Refused!Rendered, - reader: *const fn (ctx: *anyopaque, a: Allocator, thinking: bool, tools: []const Value, single: bool) Allocator.Error!Reader, - /// A finished reply's reasoning, kept for a later request whose history comes back without it. - remember: *const fn (ctx: *anyopaque, reasoning: []const u8, content: []const u8, calls: []const Value) void, - reasoning_fields: ReasoningFields = .both, - /// Whether the request's ``min_p`` reaches the sampler (some families ignore it). - reads_min_p: bool = true, - }; - - pub fn thinking(f: Family, cx: *Cx, body: Value) errors.Refused!Thinking { - return f.vt.thinking(f.ctx, cx, body); - } - pub fn defaultEffort(f: Family) []const u8 { - return f.vt.default_effort(f.ctx); - } - pub fn tools(f: Family, cx: *Cx, body: Value) errors.Refused![]const Value { - return f.vt.tools(f.ctx, cx, body); - } - pub fn render(f: Family, cx: *Cx, p: Prompt) errors.Refused!Rendered { - return f.vt.render(f.ctx, cx, p); - } - pub fn reader(f: Family, a: Allocator, think: bool, offered: []const Value, single: bool) Allocator.Error!Reader { - return f.vt.reader(f.ctx, a, think, offered, single); - } - pub fn remember(f: Family, reasoning: []const u8, content: []const u8, calls: []const Value) void { - f.vt.remember(f.ctx, reasoning, content, calls); - } - - /// A reasoning delta under the family's field names. - pub fn reasoningDelta(f: Family, a: Allocator, text: []const u8) Allocator.Error!Value { - const o = try json.newObject(a); - if (f.vt.reasoning_fields != .reasoning) try o.put(a, "reasoning_content", .{ .string = text }); - if (f.vt.reasoning_fields != .reasoning_content) try o.put(a, "reasoning", .{ .string = text }); - return .{ .object = o }; - } -}; diff --git a/zig/src/server/openai.zig b/zig/src/server/openai.zig index 54e5771b4b..de0a66c2a6 100644 --- a/zig/src/server/openai.zig +++ b/zig/src/server/openai.zig @@ -20,8 +20,6 @@ const warm = @import("warm.zig"); const cortheon = @import("cortheon/cortheon.zig"); const Server = @import("server.zig").Server; const Conn = @import("http_conn.zig").Conn; -const openai_run = @import("openai_run.zig"); -const Run = openai_run.Run; const Value = json.Value; const Cx = errors.Cx; const Allocator = std.mem.Allocator; @@ -129,7 +127,7 @@ pub fn replyModel(srv: *const Server, body: Value) []const u8 { return srv.config.served_name; } -pub const Plan = struct { +const Plan = struct { input: chat.Input, stream: bool, separate_usage: bool, @@ -146,18 +144,11 @@ fn plan(srv: *Server, cx: *Cx, is_chat: bool, raw: Value) errors.Refused!Plan { const exact_bytes = asked and srv.info.logprobs and (srv.text.tokenBytes(cx.a, 0) catch return error.OutOfMemory) != null; const logprobs = if (is_chat) try fields.probabilityOptions(cx, body, exact_bytes) else try fields.legacyLogprobs(cx, body, exact_bytes); if (srv.config.request_log) |path| request_log.append(cx.a, path, body); - var input: chat.Input = .{ .fields = undefined, .body = body }; - if (srv.family) |fam| if (is_chat) { - // a family renders the client's messages and tools itself (its encoding reads OpenAI's shapes) - const m = body.get("messages") orelse Value.null; - if (m != .array or m.array.len == 0) return cx.refuse("messages must be a non-empty list"); - input.messages = m; - input.tools = try fam.tools(cx, body); - }; - if (is_chat and srv.family == null) { + var input: chat.Input = .{ .fields = undefined }; + if (is_chat) { input.messages = try messages.normalize(cx, body.get("messages"), "system", srv.needs_user_after_tool); input.tools = try tool_specs.active(cx, body.get("tools"), body.get("tool_choice")); - } else if (!is_chat) if (body.get("messages")) |m| if (m == .array and m.array.len > 0) { + } else if (body.get("messages")) |m| if (m == .array and m.array.len > 0) { input.messages = try messages.normalize(cx, m, "system", srv.needs_user_after_tool); }; if (!is_chat and input.messages.array.len == 0) input.prompt = try legacyPrompt(srv, cx, body.get("prompt")); @@ -175,12 +166,9 @@ fn plan(srv: *Server, cx: *Cx, is_chat: bool, raw: Value) errors.Refused!Plan { const sampling = try json.newObject(cx.a); for ([_][]const u8{ "temperature", "top_p", "top_k", "min_p", "seed", "priority", "draft", "thinking_budget", "ignore_eos", "stop" } ++ grammar.fields) |k| if (body.get(k)) |v| try sampling.put(cx.a, k, v); try grammar.refusal(cx, body); - if (srv.family == null and input.tools.len > 0 and try tool_specs.choiceRequiresCall(cx.a, body.get("tool_choice"))) try sampling.put(cx.a, "tool_call_required", .{ .bool = true }); - if (srv.family == null) if (body.field("parallel_tool_calls")) |p| if (p == .bool and !p.bool) try sampling.put(cx.a, "tool_call_single", .{ .bool = true }); - const thinking: fields.Thinking = if (srv.family) |fam| blk: { - const t = try fam.thinking(cx, body); - break :blk .{ .enable = t.enable, .effort = t.effort }; - } else try fields.thinkingFields(cx, body, srv.effort_levels); + if (input.tools.len > 0 and try tool_specs.choiceRequiresCall(cx.a, body.get("tool_choice"))) try sampling.put(cx.a, "tool_call_required", .{ .bool = true }); + if (body.field("parallel_tool_calls")) |p| if (p == .bool and !p.bool) try sampling.put(cx.a, "tool_call_single", .{ .bool = true }); + const thinking = try fields.thinkingFields(cx, body, srv.effort_levels); if (thinking.effort) |e| try sampling.put(cx.a, "reasoning_effort", .{ .string = e }); if (thinking.enable) |on| try sampling.put(cx.a, "enable_thinking", .{ .bool = on }); input.fields = .{ .object = sampling }; @@ -198,7 +186,6 @@ fn plan(srv: *Server, cx: *Cx, is_chat: bool, raw: Value) errors.Refused!Plan { if (body.field("parallel_tool_calls")) |p| { if (p != .bool) return cx.refuse("parallel_tool_calls must be a boolean"); policy.single = !p.bool; - input.single_call = policy.single; } const options = body.get("stream_options"); return .{ @@ -266,20 +253,329 @@ pub fn run(srv: *Server, a: Allocator, out: Out, gone: Gone, is_chat: bool, raw: return; }; p.input.id = id; - // Cortheon reads calls in the template path's format: a family with its own chat rules runs without it - const turn = if (srv.family == null) cortheon.begin(srv, a, raw, id, is_chat, &p.input) else null; + const turn = cortheon.begin(srv, a, raw, id, is_chat, &p.input); var r: Run = .{ .srv = srv, .a = a, .out = out, .is_chat = is_chat, .plan = p, .id = id, .created = std.Io.Clock.real.now(srv.io).toSeconds(), .turn = turn }; if (p.stream) r.stream(gone, field) else r.whole(gone, field); } /// Why a request got an error reply, logged before it: its access line shows only the status. -pub fn logRefused(id: []const u8, message: []const u8) void { +fn logRefused(id: []const u8, message: []const u8) void { var buf: [1400]u8 = undefined; log.line("{s}", .{log.refused(&buf, id, message)}); } -pub fn errorOther(a: Allocator, message: []const u8) Allocator.Error!Value { +fn errorOther(a: Allocator, message: []const u8) Allocator.Error!Value { const o = try json.newObject(a); try o.put(a, "message", .{ .string = message }); return .{ .object = o }; } + +const Run = struct { + srv: *Server, + a: Allocator, + out: Out, + is_chat: bool, + plan: Plan, + id: []const u8, + created: i64, + streamed_prose: bool = false, + prose_sent: std.ArrayList(u8) = .empty, // the content a tool stream sent + text_offset: usize = 0, // characters a completion stream's logprob rows have covered + turn: ?*cortheon.Turn = null, // --cortheon's work on this request + + fn chunk(r: *Run, delta: ?Value, finish: ?[]const u8) Allocator.Error!Value { + const a = r.a; + const o = try json.newObject(a); + try o.put(a, "id", .{ .string = r.id }); + try o.put(a, "object", .{ .string = if (r.is_chat) "chat.completion.chunk" else "text_completion" }); + try o.put(a, "created", try json.intValue(a, r.created)); + try o.put(a, "model", .{ .string = r.plan.named }); + const choice = try json.newObject(a); + try choice.put(a, "index", .{ .int = "0" }); + if (r.is_chat) { + const d: Value = if (delta) |v| (if (v == .object) v else if (v.string.len > 0) try chat.deltaOf(a, "content", v.string) else .{ .object = try json.newObject(a) }) else .{ .object = try json.newObject(a) }; + try choice.put(a, "delta", d); + // a content delta's rows (logprobs.Stream) are the choice's, as OpenAI streams them + if (d == .object) if (d.object.fetchOrderedRemove("logprobs")) |kv| try choice.put(a, "logprobs", kv.value); + try choice.put(a, "finish_reason", if (finish) |f| .{ .string = f } else .null); + } else { + const d = delta orelse Value{ .string = "" }; + const rows = if (d == .object) d.object.get("logprobs") else null; + try choice.put(a, "text", if (d == .object) d.object.get("content") orelse .{ .string = "" } else d); + try choice.put(a, "finish_reason", if (finish) |f| .{ .string = f } else .null); + try choice.put(a, "logprobs", if (rows) |v| try @import("logprobs.zig").legacy(a, v, &r.text_offset) else .null); + } + const list = try a.alloc(Value, 1); + list[0] = .{ .object = choice }; + try o.put(a, "choices", .{ .array = list }); + return .{ .object = o }; + } + + fn emit(r: *Run, payload: Value) error{Closed}!void { + return r.out.vt.event(r.out.ctx, payload); + } + + fn usage(r: *Run, reply: *const chat.Reply) Allocator.Error!Value { + const a = r.a; + const o = try json.newObject(a); + try o.put(a, "prompt_tokens", try json.intValue(a, reply.prompt_tokens)); + try o.put(a, "completion_tokens", try json.intValue(a, reply.completion_tokens)); + try o.put(a, "total_tokens", try json.intValue(a, reply.prompt_tokens + reply.completion_tokens)); + const cached = try json.newObject(a); + try cached.put(a, "cached_tokens", try json.intValue(a, reply.cached_tokens)); + try o.put(a, "prompt_tokens_details", .{ .object = cached }); + const thought = try json.newObject(a); + try thought.put(a, "reasoning_tokens", try json.intValue(a, reply.reasoning_tokens)); + try o.put(a, "completion_tokens_details", .{ .object = thought }); + return .{ .object = o }; + } + + /// ``response_extras``: exact_mode, the matched stop, the runtime and the draft counters. + fn extras(r: *Run, o: *json.Object, reply: *const chat.Reply) Allocator.Error!void { + try o.put(r.a, "exact_mode", .{ .string = "exact" }); + if (reply.stop_sequence) |s| try o.put(r.a, "stop_sequence", .{ .string = s }); + try o.put(r.a, "tensorfold", reply.runtime); + try o.put(r.a, "speculative", reply.speculative); + if (r.turn) |t| try o.put(r.a, "cortheon", try t.summary(r.a)); + } + + /// ``ToolCallPolicy.finish``: the reply's calls parsed out of its content. + fn attachCalls(r: *Run, reply: *chat.Reply) Allocator.Error!?[]Value { + const tools = r.plan.input.tools; + if (tools.len == 0) return null; + const parsed = try tool_parse.parse(r.a, reply.content, tools, r.plan.policy.maxCalls()); + const calls = (if (r.turn != null) cortheon.clientCalls(parsed.calls) else parsed.calls) orelse { + if (r.plan.policy.single or parsed.calls != null) reply.content = try r.plan.policy.parsedContent(r.a, parsed.content); + return null; + }; + reply.content = try r.plan.policy.parsedContent(r.a, parsed.content); + reply.finish_reason = "tool_calls"; + if (r.plan.policy.single) reply.tool_calls_streamed = false; + return if (r.plan.policy.single and calls.len > 1) calls[0..1] else calls; + } + + fn fail(r: *Run, cx: *const Cx, field: []const u8) void { + const a = r.a; + const kind_ok = cx.kind != .other and cx.kind != .server; + const body = (if (kind_ok) errorBody(a, cx, field) else errorOther(a, cx.message)) catch return; + r.out.vt.reply(r.out.ctx, if (kind_ok) 400 else 500, wrapError(a, body) catch return); + } + + /// A reply that ended before anything was sent: a refusal's 400, a failure's 500, nothing when the client left. + fn unsent(r: *Run, cx: *Cx, e: chat.Failure, field: []const u8) void { + switch (e) { + error.Cancelled => {}, + error.OutOfMemory => r.fail(&.{ .a = r.a, .kind = .server, .message = "out of memory" }, field), + error.Failed => r.fail(&.{ .a = r.a, .kind = .server, .message = "the reply failed" }, field), + error.Refused => { + if (cx.kind == .other) cx.kind = .server; + if (cx.kind != .server) logRefused(r.id, cx.message); // a 500 is a failure, not a refusal + r.fail(cx, field); + }, + } + } + + /// A context-window refusal on a stream: the role chunk when there are no tools, then the error event and [DONE]. + fn streamedContext(r: *Run, cx: *const Cx, tools: bool, field: []const u8) void { + const a = r.a; + // a context-window refusal: the head opens with no prompt behind it, so it carries no count + r.out.vt.open(r.out.ctx, 0) catch return; + if (!tools and r.is_chat) r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; + logRefused(r.id, cx.message); + const body = errorBody(a, cx, field) catch return; + r.emit(wrapError(a, body) catch return) catch return; + r.out.vt.event(r.out.ctx, null) catch {}; + } + + fn whole(r: *Run, gone: Gone, field: []const u8) void { + const a = r.a; + var cx: Cx = .{ .a = a }; + var reply = (if (r.turn) |t| cortheon.loop.whole(t, r.srv, &cx, &r.plan.input, gone) else chat.run(r.srv, &cx, r.plan.input, null, gone)) catch |e| return r.unsent(&cx, e, field); + const calls = r.attachCalls(&reply) catch return; + const o = json.newObject(a) catch return; + r.wholeBody(o, &reply, calls) catch return; + r.out.vt.reply(r.out.ctx, 200, .{ .object = o }); + r.warmNext(&reply, calls); + } + + /// The next turn's prompt prefilled in the background, from the reply as its client got it (warm.zig). + fn warmNext(r: *Run, reply: *const chat.Reply, calls: ?[]Value) void { + if (!r.is_chat) return; + const message = json.newObject(r.a) catch return; + message.put(r.a, "role", .{ .string = "assistant" }) catch return; + message.put(r.a, "content", if (calls != null and reply.content.len == 0) .null else .{ .string = reply.content }) catch return; + if (calls) |c| message.put(r.a, "tool_calls", .{ .array = c }) catch return; + if (r.turn) |t| t.remember(r.a, r.plan.input, .{ .object = message }); + warm.reply(r.srv, r.plan.input, .{ .object = message }, reply.thinking, reply.effort, reply.prompt_tokens); + } + + fn wholeBody(r: *Run, o: *json.Object, reply: *const chat.Reply, calls: ?[]Value) Allocator.Error!void { + const a = r.a; + try o.put(a, "id", .{ .string = r.id }); + try o.put(a, "object", .{ .string = if (r.is_chat) "chat.completion" else "text_completion" }); + try o.put(a, "created", try json.intValue(a, r.created)); + try o.put(a, "model", .{ .string = r.plan.named }); + const choice = try json.newObject(a); + try choice.put(a, "index", .{ .int = "0" }); + if (r.is_chat) { + const message = try json.newObject(a); + try message.put(a, "role", .{ .string = "assistant" }); + try message.put(a, "content", if (calls != null and reply.content.len == 0) .null else .{ .string = reply.content }); + if (reply.reasoning) |t| if (t.len > 0) try message.put(a, "reasoning_content", .{ .string = t }); + if (calls) |c| try message.put(a, "tool_calls", .{ .array = c }); + try choice.put(a, "message", .{ .object = message }); + try choice.put(a, "finish_reason", .{ .string = reply.finish_reason }); + if (reply.logprobs) |lp| try choice.put(a, "logprobs", lp); + } else { + try choice.put(a, "text", .{ .string = reply.content }); + try choice.put(a, "finish_reason", .{ .string = reply.finish_reason }); + var from: usize = 0; + try choice.put(a, "logprobs", if (reply.logprobs) |lp| try @import("logprobs.zig").legacy(a, lp, &from) else .null); + if (reply.prompt_logprobs) |p| try choice.put(a, "prompt_logprobs", p); + } + const list = try a.alloc(Value, 1); + list[0] = .{ .object = choice }; + try o.put(a, "choices", .{ .array = list }); + try o.put(a, "usage", try r.usage(reply)); + try r.extras(o, reply); + } + + const StreamSink = struct { + run: *Run, + tools: bool, + + fn call(ctx: *anyopaque, delta: Value) error{Closed}!void { + const s: *StreamSink = @ptrCast(@alignCast(ctx)); + const r = s.run; + if (!s.tools) { + // completions stream text only, with its logprob rows when asked + if (!r.is_chat and delta != .string and (delta != .object or delta.object.get("content") == null)) return; + return r.emit(r.chunk(delta, null) catch return error.Closed); + } + const filtered = r.policyDelta(delta) catch return error.Closed; + const d = filtered orelse return; + try r.prose(); + r.prose_sent.appendSlice(r.a, if (d == .string) d.string else d.strField("content") orelse "") catch return error.Closed; + return r.emit(r.chunk(d, null) catch return error.Closed); + } + }; + + /// The one-call policy on a streamed delta: text filtered, tool_calls dropped; null when nothing is left. + fn policyDelta(r: *Run, delta: Value) Allocator.Error!?Value { + if (!r.plan.policy.single) return if (delta.truthy()) delta else null; + if (delta == .string) { + const t = try r.plan.policy.text(r.a, delta.string); + return if (t.len > 0) Value{ .string = t } else null; + } + const o = try json.newObject(r.a); + var content: ?[]const u8 = null; + for (delta.object.keys(), delta.object.values()) |k, v| { + if (std.mem.eql(u8, k, "tool_calls")) continue; + if (std.mem.eql(u8, k, "content") and v == .string) { + content = v.string; + continue; + } + try o.put(r.a, k, v); + } + if (content) |c| { + const t = try r.plan.policy.text(r.a, c); + if (t.len > 0) try o.put(r.a, "content", .{ .string = t }); + } + return if (o.count() > 0) Value{ .object = o } else null; + } + + /// The assistant role before the first prose of a tool-using stream. + fn prose(r: *Run) error{Closed}!void { + if (r.streamed_prose) return; + r.streamed_prose = true; + const role = roleDelta(r.a) catch return error.Closed; + return r.emit(r.chunk(role, null) catch return error.Closed); + } + + fn stream(r: *Run, gone: Gone, field: []const u8) void { + const a = r.a; + const tools = r.plan.input.tools.len > 0; + var cx: Cx = .{ .a = a }; + // a context-window refusal is reported inside the stream; every other refusal is still a 400 before it opens + var compaction: ?compact.Stamp = null; + const prepared = if (r.srv.config.compact_at == null) chat.prepare(r.srv, &cx, r.plan.input, gone) catch |e| { + if (e == error.Refused and cx.kind == .context_length) return r.streamedContext(&cx, tools, field); + return r.unsent(&cx, e, field); + } else blk: { + const ready = compact.prepare(r.srv, &cx, r.plan.input, gone) catch |e| { + if (e == error.Refused and cx.kind == .context_length) return r.streamedContext(&cx, tools, field); + return r.unsent(&cx, e, field); + }; + compaction = ready.stamp; + 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); + r.out.vt.open(r.out.ctx, prepared.prompt_len) catch return; + if (!tools and r.is_chat) r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; + var sink_state: StreamSink = .{ .run = r, .tools = tools }; + handed = true; + var live = gone; + live.out = r.out; + live.every_ns = @intFromFloat(r.srv.config.sse_keepalive_s * std.time.ns_per_s); + const sink: chat.Sink = .{ .ctx = &sink_state, .call = StreamSink.call }; + var reply = (if (r.turn) |t| cortheon.loop.stream(t, r.srv, &cx, &r.plan.input, prepared, sink, gone) else chat.generate(r.srv, &cx, prepared, sink, live)) catch |e| { + switch (e) { + error.Cancelled => return, + error.Refused => if (cx.kind != .other and cx.kind != .server) { + logRefused(r.id, cx.message); + const body = errorBody(a, &cx, field) catch return; + r.emit(wrapError(a, body) catch return) catch return; + r.out.vt.event(r.out.ctx, null) catch {}; + return; + }, + else => {}, + } + const message = if (e == error.Refused) cx.message else if (e == error.OutOfMemory) "out of memory" else "the reply failed"; + log.line("stream error: {s}", .{message}); + const err = json.newObject(a) catch return; + err.put(a, "message", .{ .string = message }) catch return; + err.put(a, "type", .{ .string = "server_error" }) catch return; + r.emit(wrapError(a, .{ .object = err }) catch return) catch return; + r.out.vt.event(r.out.ctx, null) catch {}; + return; + }; + var calls: ?[]Value = null; + if (tools) { + calls = r.attachCalls(&reply) catch return; + const tail = r.plan.policy.flush(); + if (tail.len > 0) { + r.prose() catch return; + r.prose_sent.appendSlice(a, tail) catch return; + r.emit(r.chunk(.{ .string = tail }, null) catch return) catch return; + } + if (calls != null and !reply.tool_calls_streamed) { + r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; + const deltas = tool_parse.deltas(a, calls.?) catch return; + for (deltas) |d| r.emit(r.chunk(d, null) catch return) catch return; + } else if (reply.content.len > 0 and !r.streamed_prose) { + r.emit(r.chunk(.{ .string = reply.content }, null) catch return) catch return; + } else if (r.plan.policy.kept(reply.content, r.prose_sent.items)) |rest| { + r.emit(r.chunk(.{ .string = rest }, null) catch return) catch return; + } + } + const last = r.chunk(.{ .string = "" }, if (reply.finish_reason.len > 0) reply.finish_reason else "length") catch return; + if (compaction) |s| compact.stamp(&cx, &reply, s) catch return; + r.extras(last.object, &reply) catch return; + const use = r.usage(&reply) catch return; + if (!r.plan.separate_usage) last.object.put(a, "usage", use) catch return; + r.emit(last) catch return; + if (r.plan.separate_usage) { + const u = r.chunk(null, null) catch return; + u.object.put(a, "choices", .{ .array = &.{} }) catch return; + u.object.put(a, "usage", use) catch return; + r.emit(u) catch return; + } + r.out.vt.event(r.out.ctx, null) catch {}; + r.warmNext(&reply, calls); + } +}; + +fn roleDelta(a: Allocator) Allocator.Error!Value { + return chat.deltaOf(a, "role", "assistant"); +} diff --git a/zig/src/server/openai_run.zig b/zig/src/server/openai_run.zig deleted file mode 100644 index 5476dc8674..0000000000 --- a/zig/src/server/openai_run.zig +++ /dev/null @@ -1,346 +0,0 @@ -//! One OpenAI request's run: the reply generated, then sent whole or streamed as chunks. -const std = @import("std"); -const json = @import("json.zig"); -const errors = @import("errors.zig"); -const chat = @import("chat.zig"); -const compact = @import("compact.zig"); -const tool_parse = @import("tool_parse.zig"); -const tool_stream = @import("tool_stream.zig"); -const ids = @import("ids.zig"); -const log = @import("log.zig"); -const request_log = @import("request_log.zig"); -const sse = @import("sse.zig"); -const warm = @import("warm.zig"); -const cortheon = @import("cortheon/cortheon.zig"); -const Server = @import("server.zig").Server; -const Value = json.Value; -const Cx = errors.Cx; -const Allocator = std.mem.Allocator; -const openai = @import("openai.zig"); -const Out = openai.Out; -const Gone = openai.Gone; -const errorBody = openai.errorBody; -const wrapError = openai.wrapError; -const replyModel = openai.replyModel; -const Plan = openai.Plan; -const logRefused = openai.logRefused; -const errorOther = openai.errorOther; - -pub const Run = struct { - srv: *Server, - a: Allocator, - out: Out, - is_chat: bool, - plan: Plan, - id: []const u8, - created: i64, - streamed_prose: bool = false, - prose_sent: std.ArrayList(u8) = .empty, // the content a tool stream sent - text_offset: usize = 0, // characters a completion stream's logprob rows have covered - turn: ?*cortheon.Turn = null, // --cortheon's work on this request - - fn chunk(r: *Run, delta: ?Value, finish: ?[]const u8) Allocator.Error!Value { - const a = r.a; - const o = try json.newObject(a); - try o.put(a, "id", .{ .string = r.id }); - try o.put(a, "object", .{ .string = if (r.is_chat) "chat.completion.chunk" else "text_completion" }); - try o.put(a, "created", try json.intValue(a, r.created)); - try o.put(a, "model", .{ .string = r.plan.named }); - const choice = try json.newObject(a); - try choice.put(a, "index", .{ .int = "0" }); - if (r.is_chat) { - const d: Value = if (delta) |v| (if (v == .object) v else if (v.string.len > 0) try chat.deltaOf(a, "content", v.string) else .{ .object = try json.newObject(a) }) else .{ .object = try json.newObject(a) }; - try choice.put(a, "delta", d); - // a content delta's rows (logprobs.Stream) are the choice's, as OpenAI streams them - if (d == .object) if (d.object.fetchOrderedRemove("logprobs")) |kv| try choice.put(a, "logprobs", kv.value); - try choice.put(a, "finish_reason", if (finish) |f| .{ .string = f } else .null); - } else { - const d = delta orelse Value{ .string = "" }; - const rows = if (d == .object) d.object.get("logprobs") else null; - try choice.put(a, "text", if (d == .object) d.object.get("content") orelse .{ .string = "" } else d); - try choice.put(a, "finish_reason", if (finish) |f| .{ .string = f } else .null); - try choice.put(a, "logprobs", if (rows) |v| try @import("logprobs.zig").legacy(a, v, &r.text_offset) else .null); - } - const list = try a.alloc(Value, 1); - list[0] = .{ .object = choice }; - try o.put(a, "choices", .{ .array = list }); - return .{ .object = o }; - } - - fn emit(r: *Run, payload: Value) error{Closed}!void { - return r.out.vt.event(r.out.ctx, payload); - } - - fn usage(r: *Run, reply: *const chat.Reply) Allocator.Error!Value { - const a = r.a; - const o = try json.newObject(a); - try o.put(a, "prompt_tokens", try json.intValue(a, reply.prompt_tokens)); - try o.put(a, "completion_tokens", try json.intValue(a, reply.completion_tokens)); - try o.put(a, "total_tokens", try json.intValue(a, reply.prompt_tokens + reply.completion_tokens)); - const cached = try json.newObject(a); - try cached.put(a, "cached_tokens", try json.intValue(a, reply.cached_tokens)); - try o.put(a, "prompt_tokens_details", .{ .object = cached }); - const thought = try json.newObject(a); - try thought.put(a, "reasoning_tokens", try json.intValue(a, reply.reasoning_tokens)); - try o.put(a, "completion_tokens_details", .{ .object = thought }); - return .{ .object = o }; - } - - /// ``response_extras``: exact_mode, the matched stop, the runtime and the draft counters. - fn extras(r: *Run, o: *json.Object, reply: *const chat.Reply) Allocator.Error!void { - try o.put(r.a, "exact_mode", .{ .string = "exact" }); - if (reply.stop_sequence) |s| try o.put(r.a, "stop_sequence", .{ .string = s }); - try o.put(r.a, "tensorfold", reply.runtime); - try o.put(r.a, "speculative", reply.speculative); - if (r.turn) |t| try o.put(r.a, "cortheon", try t.summary(r.a)); - } - - /// ``ToolCallPolicy.finish``: the reply's calls parsed out of its content. - fn attachCalls(r: *Run, reply: *chat.Reply) Allocator.Error!?[]Value { - if (reply.calls) |calls| return if (calls.len > 0) calls else null; // the family parsed them (and applied the one-call policy) - const tools = r.plan.input.tools; - if (tools.len == 0) return null; - const parsed = try tool_parse.parse(r.a, reply.content, tools, r.plan.policy.maxCalls()); - const calls = (if (r.turn != null) cortheon.clientCalls(parsed.calls) else parsed.calls) orelse { - if (r.plan.policy.single or parsed.calls != null) reply.content = try r.plan.policy.parsedContent(r.a, parsed.content); - return null; - }; - reply.content = try r.plan.policy.parsedContent(r.a, parsed.content); - reply.finish_reason = "tool_calls"; - if (r.plan.policy.single) reply.tool_calls_streamed = false; - return if (r.plan.policy.single and calls.len > 1) calls[0..1] else calls; - } - - fn fail(r: *Run, cx: *const Cx, field: []const u8) void { - const a = r.a; - const kind_ok = cx.kind != .other and cx.kind != .server; - const body = (if (kind_ok) errorBody(a, cx, field) else errorOther(a, cx.message)) catch return; - r.out.vt.reply(r.out.ctx, if (kind_ok) 400 else 500, wrapError(a, body) catch return); - } - - /// A reply that ended before anything was sent: a refusal's 400, a failure's 500, nothing when the client left. - fn unsent(r: *Run, cx: *Cx, e: chat.Failure, field: []const u8) void { - switch (e) { - error.Cancelled => {}, - error.OutOfMemory => r.fail(&.{ .a = r.a, .kind = .server, .message = "out of memory" }, field), - error.Failed => r.fail(&.{ .a = r.a, .kind = .server, .message = "the reply failed" }, field), - error.Refused => { - if (cx.kind == .other) cx.kind = .server; - if (cx.kind != .server) logRefused(r.id, cx.message); // a 500 is a failure, not a refusal - r.fail(cx, field); - }, - } - } - - /// A context-window refusal on a stream: the role chunk when there are no tools, then the error event and [DONE]. - fn streamedContext(r: *Run, cx: *const Cx, tools: bool, field: []const u8) void { - const a = r.a; - // a context-window refusal: the head opens with no prompt behind it, so it carries no count - r.out.vt.open(r.out.ctx, 0) catch return; - if (!tools and r.is_chat) r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; - logRefused(r.id, cx.message); - const body = errorBody(a, cx, field) catch return; - r.emit(wrapError(a, body) catch return) catch return; - r.out.vt.event(r.out.ctx, null) catch {}; - } - - pub fn whole(r: *Run, gone: Gone, field: []const u8) void { - const a = r.a; - var cx: Cx = .{ .a = a }; - var reply = (if (r.turn) |t| cortheon.loop.whole(t, r.srv, &cx, &r.plan.input, gone) else chat.run(r.srv, &cx, r.plan.input, null, gone)) catch |e| return r.unsent(&cx, e, field); - const calls = r.attachCalls(&reply) catch return; - const o = json.newObject(a) catch return; - r.wholeBody(o, &reply, calls) catch return; - r.out.vt.reply(r.out.ctx, 200, .{ .object = o }); - r.warmNext(&reply, calls); - } - - /// The next turn's prompt prefilled in the background, from the reply as its client got it (warm.zig). - fn warmNext(r: *Run, reply: *const chat.Reply, calls: ?[]Value) void { - if (!r.is_chat) return; - const message = json.newObject(r.a) catch return; - message.put(r.a, "role", .{ .string = "assistant" }) catch return; - message.put(r.a, "content", if (calls != null and reply.content.len == 0) .null else .{ .string = reply.content }) catch return; - if (calls) |c| message.put(r.a, "tool_calls", .{ .array = c }) catch return; - if (r.turn) |t| t.remember(r.a, r.plan.input, .{ .object = message }); - warm.reply(r.srv, r.plan.input, .{ .object = message }, reply.thinking, reply.effort, reply.prompt_tokens); - } - - fn wholeBody(r: *Run, o: *json.Object, reply: *const chat.Reply, calls: ?[]Value) Allocator.Error!void { - const a = r.a; - try o.put(a, "id", .{ .string = r.id }); - try o.put(a, "object", .{ .string = if (r.is_chat) "chat.completion" else "text_completion" }); - try o.put(a, "created", try json.intValue(a, r.created)); - try o.put(a, "model", .{ .string = r.plan.named }); - const choice = try json.newObject(a); - try choice.put(a, "index", .{ .int = "0" }); - if (r.is_chat) { - const message = try json.newObject(a); - try message.put(a, "role", .{ .string = "assistant" }); - // a family's reply: ``content or None`` and no prose beside calls; the template path keeps a call's prose - const no_content = if (r.srv.family != null) calls != null or reply.content.len == 0 else calls != null and reply.content.len == 0; - try message.put(a, "content", if (no_content) .null else .{ .string = reply.content }); - if (reply.reasoning) |t| if (t.len > 0) { - if (r.srv.family) |fam| { - const d = try fam.reasoningDelta(a, t); - for (d.object.keys(), d.object.values()) |k, v| try message.put(a, k, v); - } else try message.put(a, "reasoning_content", .{ .string = t }); - }; - if (calls) |c| try message.put(a, "tool_calls", .{ .array = c }); - try choice.put(a, "message", .{ .object = message }); - try choice.put(a, "finish_reason", .{ .string = reply.finish_reason }); - if (reply.logprobs) |lp| try choice.put(a, "logprobs", lp); - } else { - try choice.put(a, "text", .{ .string = reply.content }); - try choice.put(a, "finish_reason", .{ .string = reply.finish_reason }); - var from: usize = 0; - try choice.put(a, "logprobs", if (reply.logprobs) |lp| try @import("logprobs.zig").legacy(a, lp, &from) else .null); - if (reply.prompt_logprobs) |p| try choice.put(a, "prompt_logprobs", p); - } - const list = try a.alloc(Value, 1); - list[0] = .{ .object = choice }; - try o.put(a, "choices", .{ .array = list }); - try o.put(a, "usage", try r.usage(reply)); - try r.extras(o, reply); - } - - const StreamSink = struct { - run: *Run, - tools: bool, - - fn call(ctx: *anyopaque, delta: Value) error{Closed}!void { - const s: *StreamSink = @ptrCast(@alignCast(ctx)); - const r = s.run; - if (!s.tools) { - // completions stream text only, with its logprob rows when asked - if (!r.is_chat and delta != .string and (delta != .object or delta.object.get("content") == null)) return; - return r.emit(r.chunk(delta, null) catch return error.Closed); - } - const filtered = r.policyDelta(delta) catch return error.Closed; - const d = filtered orelse return; - try r.prose(); - r.prose_sent.appendSlice(r.a, if (d == .string) d.string else d.strField("content") orelse "") catch return error.Closed; - return r.emit(r.chunk(d, null) catch return error.Closed); - } - }; - - /// The one-call policy on a streamed delta: text filtered, tool_calls dropped; null when nothing is left. - fn policyDelta(r: *Run, delta: Value) Allocator.Error!?Value { - if (!r.plan.policy.single) return if (delta.truthy()) delta else null; - if (delta == .string) { - const t = try r.plan.policy.text(r.a, delta.string); - return if (t.len > 0) Value{ .string = t } else null; - } - const o = try json.newObject(r.a); - var content: ?[]const u8 = null; - for (delta.object.keys(), delta.object.values()) |k, v| { - if (std.mem.eql(u8, k, "tool_calls")) continue; - if (std.mem.eql(u8, k, "content") and v == .string) { - content = v.string; - continue; - } - try o.put(r.a, k, v); - } - if (content) |c| { - const t = try r.plan.policy.text(r.a, c); - if (t.len > 0) try o.put(r.a, "content", .{ .string = t }); - } - return if (o.count() > 0) Value{ .object = o } else null; - } - - /// The assistant role before the first prose of a tool-using stream. - fn prose(r: *Run) error{Closed}!void { - if (r.streamed_prose) return; - r.streamed_prose = true; - const role = roleDelta(r.a) catch return error.Closed; - return r.emit(r.chunk(role, null) catch return error.Closed); - } - - pub fn stream(r: *Run, gone: Gone, field: []const u8) void { - const a = r.a; - // a family streams calls and prose itself, so its stream is the plain one - const tools = r.plan.input.tools.len > 0 and r.srv.family == null; - var cx: Cx = .{ .a = a }; - // a context-window refusal is reported inside the stream - var compaction: ?compact.Stamp = null; - const prepared = if (r.srv.config.compact_at == null) chat.prepare(r.srv, &cx, r.plan.input, gone) catch |e| { - if (e == error.Refused and cx.kind == .context_length) return r.streamedContext(&cx, tools, field); - return r.unsent(&cx, e, field); - } else blk: { - const ready = compact.prepare(r.srv, &cx, r.plan.input, gone) catch |e| { - if (e == error.Refused and cx.kind == .context_length) return r.streamedContext(&cx, tools, field); - return r.unsent(&cx, e, field); - }; - compaction = ready.stamp; - 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); - r.out.vt.open(r.out.ctx, prepared.prompt_len) catch return; - if (!tools and r.is_chat) r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; - var sink_state: StreamSink = .{ .run = r, .tools = tools }; - handed = true; - var live = gone; - live.out = r.out; - live.every_ns = @intFromFloat(r.srv.config.sse_keepalive_s * std.time.ns_per_s); - const sink: chat.Sink = .{ .ctx = &sink_state, .call = StreamSink.call }; - var reply = (if (r.turn) |t| cortheon.loop.stream(t, r.srv, &cx, &r.plan.input, prepared, sink, live) else chat.generate(r.srv, &cx, prepared, sink, live)) catch |e| { - switch (e) { - error.Cancelled => return, - error.Refused => if (cx.kind != .other and cx.kind != .server) { - logRefused(r.id, cx.message); - const body = errorBody(a, &cx, field) catch return; - r.emit(wrapError(a, body) catch return) catch return; - r.out.vt.event(r.out.ctx, null) catch {}; - return; - }, - else => {}, - } - const message = if (e == error.Refused) cx.message else if (e == error.OutOfMemory) "out of memory" else "the reply failed"; - log.line("stream error: {s}", .{message}); - const err = json.newObject(a) catch return; - err.put(a, "message", .{ .string = message }) catch return; - err.put(a, "type", .{ .string = "server_error" }) catch return; - r.emit(wrapError(a, .{ .object = err }) catch return) catch return; - r.out.vt.event(r.out.ctx, null) catch {}; - return; - }; - var calls: ?[]Value = null; - if (tools) { - calls = r.attachCalls(&reply) catch return; - const tail = r.plan.policy.flush(); - if (tail.len > 0) { - r.prose() catch return; - r.prose_sent.appendSlice(a, tail) catch return; - r.emit(r.chunk(.{ .string = tail }, null) catch return) catch return; - } - if (calls != null and !reply.tool_calls_streamed) { - r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; - const deltas = tool_parse.deltas(a, calls.?) catch return; - for (deltas) |d| r.emit(r.chunk(d, null) catch return) catch return; - } else if (reply.content.len > 0 and !r.streamed_prose) { - r.emit(r.chunk(.{ .string = reply.content }, null) catch return) catch return; - } else if (r.plan.policy.kept(reply.content, r.prose_sent.items)) |rest| { - r.emit(r.chunk(.{ .string = rest }, null) catch return) catch return; - } - } - const last = r.chunk(.{ .string = "" }, if (reply.finish_reason.len > 0) reply.finish_reason else "length") catch return; - if (compaction) |s| compact.stamp(&cx, &reply, s) catch return; - r.extras(last.object, &reply) catch return; - const use = r.usage(&reply) catch return; - if (!r.plan.separate_usage) last.object.put(a, "usage", use) catch return; - r.emit(last) catch return; - if (r.plan.separate_usage) { - const u = r.chunk(null, null) catch return; - u.object.put(a, "choices", .{ .array = &.{} }) catch return; - u.object.put(a, "usage", use) catch return; - r.emit(u) catch return; - } - r.out.vt.event(r.out.ctx, null) catch {}; - r.warmNext(&reply, calls); - } -}; - -fn roleDelta(a: Allocator) Allocator.Error!Value { - return chat.deltaOf(a, "role", "assistant"); -} diff --git a/zig/src/server/prompt.zig b/zig/src/server/prompt.zig index f3fa2c9931..e68363bfe7 100644 --- a/zig/src/server/prompt.zig +++ b/zig/src/server/prompt.zig @@ -6,11 +6,10 @@ const messages_mod = @import("messages.zig"); const model_text = @import("model_text.zig"); const chat = @import("chat.zig"); const Server = @import("server.zig").Server; -const family_mod = @import("family.zig"); const Value = json.Value; const Cx = errors.Cx; -pub const Rendered = struct { ids: []const u32, history_len: usize = 0, rewind_len: usize = 0, images_omitted: u32 = 0, images: ?family_mod.Images = null }; +pub const Rendered = struct { ids: []const u32, history_len: usize = 0, rewind_len: usize = 0 }; /// Where lessons decide to act on a reply: the row before the generation prompt's ``, else before its last. pub fn decideOf(srv: *const Server, ids: []const u32) u32 { @@ -28,11 +27,6 @@ pub const isTitle = messages_mod.isTitleRequest; /// ``render_prompt_ids``: messages normalized for this template, rendered, and a dangling ```` closed. pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, thinking: bool, effort: ?[]const u8, generation: bool) errors.Refused![]const u32 { - if (srv.family) |fam| { - const got = try fam.render(cx, .{ .messages = messages, .tools = tools, .thinking = thinking, .effort = effort orelse fam.defaultEffort(), .generation = generation }); - if (got.images) |im| im.release(); // ids only: the images are not held past the render - return got.ids; - } const normalized = try messages_mod.toolArguments(cx, try messages_mod.normalize(cx, messages, srv.late_system, srv.needs_user_after_tool)); var problem: []const u8 = ""; const options: model_text.RenderOptions = .{ @@ -102,17 +96,6 @@ pub fn prepare(srv: *Server, cx: *Cx, input: chat.Input, thinking: bool, effort: } }, .ids => |ids| return .{ .ids = ids }, }; - if (srv.family) |fam| { - const p: family_mod.Prompt = .{ .messages = input.messages, .tools = input.tools, .thinking = thinking, .effort = effort orelse fam.defaultEffort(), .body = input.body }; - const got = try fam.render(cx, p); - var history = p; - history.generation = false; - errdefer if (got.images) |im| im.release(); - const before = try fam.render(cx, history); - if (before.images) |im| im.release(); - const n = before.ids.len; - return .{ .ids = got.ids, .history_len = if (n > 0 and n < got.ids.len and std.mem.eql(u32, got.ids[0..n], before.ids)) n else 0, .images_omitted = got.images_omitted, .images = got.images }; - } const prompt = try renderIds(srv, cx, input.messages, input.tools, thinking, effort, true); const history = try renderIds(srv, cx, input.messages, input.tools, thinking, effort, false); var history_len: usize = 0; diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 0db752b8af..15dacbdf3b 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -12,7 +12,6 @@ pub const server = @import("server.zig"); pub const serve = @import("serve.zig"); pub const hub = @import("hub.zig"); pub const tool_specs = @import("tool_specs.zig"); -pub const family = @import("family.zig"); pub const Server = server.Server; pub const Config = server.Config; diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index 4fa490d103..75a9cc285f 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -86,8 +86,6 @@ pub const Setup = struct { started: i96 = 0, /// Called once the socket listens, with its port (tests read it). on_listen: ?*const fn (port: u16) void = null, - /// The checkpoint family's own chat rules (family.zig), or null. - family: ?@import("family.zig").Family = null, /// Run by `run` so the engine stops before the server is freed; null: the caller closes it after `run`. stop: ?EngineStop = null, }; @@ -155,7 +153,6 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { .sse_keepalive_s = keepalive_s, .request_log = env(s, "TENSORFOLD_REQUEST_LOG"), .dashboard = args.dashboard, - .family = s.family, .compact_at = if (args.compact_auto) .{ .auto = {} } else if (args.compact_fraction) |f| .{ .fraction = f } else null, .compact_keep = args.compact_keep, .compact_memory = args.compact_memory, diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 505fbdc344..1fa8336a39 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -23,7 +23,6 @@ const live = @import("live.zig"); const slide_graph = @import("slide_graph.zig"); const slide_lesson = @import("slide_lesson.zig"); const cortheon = @import("cortheon/cortheon.zig"); -const family_mod = @import("family.zig"); const Value = json.Value; const Cx = errors.Cx; const Allocator = std.mem.Allocator; @@ -65,8 +64,6 @@ pub const Config = struct { timeouts: http_conn.Timeouts = .{}, /// TF_SSE_KEEPALIVE_S: a silent stream's keepalive period before its first token (serve.zig; 0: none) sse_keepalive_s: f64 = 15, - /// The checkpoint family's own chat rules (family.zig); null: the chat template path. - family: ?family_mod.Family = null, }; pub const Server = struct { @@ -87,7 +84,6 @@ pub const Server = struct { think_close_end: ?u32 = null, metrics: metrics.Metrics, store: responses.Store, - family: ?family_mod.Family = null, next_id: std.atomic.Value(u64) = .init(1), /// Connections open now; a stop waits for them before freeing what they read. open_connections: std.atomic.Value(u32) = .init(0), @@ -112,7 +108,7 @@ pub const Server = struct { /// Reads what the template and tokenizer decide once: late system role, think markers, efforts, the forced close. pub fn init(gpa: Allocator, io: std.Io, engine: api.Engine, text: model_text.Text, config: Config, keys: ?*auth.Store) !*Server { const srv = try gpa.create(Server); - srv.* = .{ .gpa = gpa, .io = io, .engine = engine, .info = engine.info(), .text = text, .config = config, .keys = keys, .metrics = .{ .gpa = gpa }, .store = .{ .gpa = gpa }, .arena = .init(gpa), .slide = .init(gpa, io, config.slide_graph), .teacher = .init(gpa), .family = config.family }; + srv.* = .{ .gpa = gpa, .io = io, .engine = engine, .info = engine.info(), .text = text, .config = config, .keys = keys, .metrics = .{ .gpa = gpa }, .store = .{ .gpa = gpa }, .arena = .init(gpa), .slide = .init(gpa, io, config.slide_graph), .teacher = .init(gpa) }; errdefer srv.deinit(); // a later step's failure stops the keepalive before the caller closes the engine if (config.cortheon) srv.cortheon = try cortheon.State.init(gpa, io, "TensorFold-Cortheon (+https://github.com/ashhart/TensorFold)", config.cortheon_search_key); if (config.keep_warm_s > 0) if (engine.keepaliveTarget()) |target| { @@ -166,10 +162,6 @@ pub const Server = struct { const defaults = if (srv.config.default_sampling) |d| try fields.parseNumbers(cx, d) else null; if (defaults) |d| try merge(options, cx, d); const asked = try fields.parseNumbers(cx, f); - // a family whose Python server ignores the request's min_p keeps the default's - if (srv.family) |fam| if (!fam.vt.reads_min_p) { - _ = asked.object.orderedRemove("min_p"); - }; try merge(options, cx, asked); const temp = number(options.get("temperature")) orelse 0; if (temp <= 0) return null; diff --git a/zig/src/server/tokens.zig b/zig/src/server/tokens.zig index 52bcb87092..3a2e351921 100644 --- a/zig/src/server/tokens.zig +++ b/zig/src/server/tokens.zig @@ -27,27 +27,7 @@ fn tokenize(srv: *Server, cx: *Cx, body: Value) errors.Refused!Value { if (body != .object) return cx.refuse("the request body must be a JSON object"); const strings = try flag(cx, body, "return_token_strs", false); var ids: []const u32 = undefined; - if (body.has("messages") and srv.family != null) { - // rendered exactly as a chat request with the same body - const fam = srv.family.?; - const generation = try flag(cx, body, "add_generation_prompt", true); - const m = body.get("messages").?; - if (m != .array or m.array.len == 0) return cx.refuse("messages must be a non-empty list"); - const t = try fam.thinking(cx, body); - const rendered = fam.render(cx, .{ .messages = m, .tools = try fam.tools(cx, body), .thinking = t.enable orelse srv.config.enable_thinking, .effort = t.effort orelse srv.config.reasoning_effort orelse fam.defaultEffort(), .generation = generation, .body = body }) catch |e| switch (e) { - error.Refused => return cx.refuse(cx.message), - else => |x| return x, - }; - ids = rendered.ids; - if (rendered.images) |im| { // /tokenize shows the image token at each image position (expand(virtual=False)) - defer im.release(); - const out = try cx.a.dupe(u32, rendered.ids); - for (out) |*id| if (id.* >= 1 << 24) { - id.* = im.token; - }; - ids = out; - } - } else if (body.has("messages")) { + if (body.has("messages")) { const generation = try flag(cx, body, "add_generation_prompt", true); const msgs = try messages.normalize(cx, body.get("messages"), "system", srv.needs_user_after_tool); const tools = tool_specs.active(cx, body.get("tools"), body.get("tool_choice")) catch |e| switch (e) { From 864fe253277eef5909680c5d5ed42e1aacce6874 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Fri, 16 Oct 2026 11:26:38 +0700 Subject: [PATCH 31/43] Graph cache: effective settings have one digest for ranks to compare at start; RUNBOOK says what a mismatch does An off cache never calls Agree, so a rank left on waits at its first miss for a vote that never comes. Settings.effective applies TF_GRAPH_CACHE, clamps min_keep and maps every off cache (0, off, OFF, a max of 0) to one value; Cache.init runs with it, and Settings.digest gives ranks those bytes to compare before building a cache. Hits still make no collective. --- RUNBOOK.md | 6 +- zig/src/cuda/graph_cache.zig | 34 +++++- zig/src/cuda/graph_cache_ranks_test.zig | 138 ++++++++++++++++++++++++ 3 files changed, 171 insertions(+), 7 deletions(-) create mode 100644 zig/src/cuda/graph_cache_ranks_test.zig diff --git a/RUNBOOK.md b/RUNBOOK.md index 525cdf49e5..4fd649e9be 100644 --- a/RUNBOOK.md +++ b/RUNBOOK.md @@ -157,8 +157,10 @@ budget free for the next states. `TF_HEAT_HIGH` and `TF_HEAT_LOW`, in degrees Celsius and set together, make each prompt chunk wait while the hottest thermal zone (under `/sys/class/thermal`, or `TF_HEAT_ROOT`) is above the high band, until it is at or under the low one. The server logs `heat_wait_s`, a cancel ends the wait, and the reply's tokens do not change. -`TF_GRAPH_CACHE=0` (or `off`) runs every step eagerly in any family that installs the keyed CUDA graph cache, without a -rebuild; set it alike on every tensor-parallel rank. +`TF_GRAPH_CACHE=0` (or `off`, any case) runs every step eagerly in any family that installs the keyed CUDA graph cache, +without a rebuild. Set it alike on every tensor-parallel rank: an off rank never joins the cache's per-miss vote, so a +rank left on would wait at its first miss, and only the off rank logs the switch. Ranks that share a cache compare +`Settings.digest` of their effective settings at start, before building it, and refuse to start when they differ. Archives whose names contain `host-only` are CPU verification artifacts and cannot serve CUDA inference. ## Flash Next and paired Metal serving diff --git a/zig/src/cuda/graph_cache.zig b/zig/src/cuda/graph_cache.zig index 9c62e4fc7d..1916b22b57 100644 --- a/zig/src/cuda/graph_cache.zig +++ b/zig/src/cuda/graph_cache.zig @@ -12,7 +12,7 @@ pub const Body = struct { run: *const fn (ctx: *anyopaque, stream: Stream) anyerror!void, }; -/// Every rank's verdict AND-ed on a miss (hits never call it), so ranks capture, evict and fall back together. +/// Every rank's verdict AND-ed on a miss, never a hit; an off cache never votes, so ranks match `Settings.digest`. pub const Agree = struct { ctx: *anyopaque, all: *const fn (ctx: *anyopaque, ok: bool) anyerror!bool, @@ -39,6 +39,30 @@ pub const Settings = struct { if (std.mem.eql(u8, v, "0") or std.ascii.eqlIgnoreCase(v, "off")) out.on = false; return out; } + + /// What a cache runs with: the operator's value, min_keep under max (0 with hold), and every off cache one value. + pub fn effective(s: Settings, value: ?[]const u8) Settings { + var out = s.operator(value); + if (!out.on or out.max == 0) return .{ .on = false, .max = 0, .min_keep = 0, .hold = false }; + out.min_keep = if (out.hold) 0 else @min(out.min_keep, out.max); + return out; + } + + /// `effective` under this process's TF_GRAPH_CACHE. + pub fn fromEnv(s: Settings) Settings { + return s.effective(if (std.c.getenv("TF_GRAPH_CACHE")) |v| std.mem.span(v) else null); + } + + /// The effective settings as bytes for ranks to compare at start (a family's bootstrap agreement), never per step. + pub fn digest(s: Settings) [10]u8 { + const e = s.effective(null); + var out: [10]u8 = undefined; + out[0] = @intFromBool(e.on); + out[1] = @intFromBool(e.hold); + std.mem.writeInt(u32, out[2..6], e.max, .little); + std.mem.writeInt(u32, out[6..10], e.min_keep, .little); + return out; + } }; pub const Outcome = enum { replayed, captured, eager }; @@ -60,11 +84,10 @@ pub fn Cache(comptime Key: type) type { failed: bool = false, // a capture failed on some rank: eager from here on stats: struct { replayed: u64 = 0, captured: u64 = 0, eager: u64 = 0, evicted: u64 = 0, dropped: u64 = 0, floor_hits: u64 = 0 } = .{}, - /// The cache with `settings`, unless TF_GRAPH_CACHE turns it off (set it alike on every rank, as `settings`). + /// The cache with `settings.fromEnv()`: every rank sharing an `agree` must hold an equal `Settings.digest`. pub fn init(gpa: std.mem.Allocator, engine: Engine, settings: Settings) Self { - var s = settings.operator(if (std.c.getenv("TF_GRAPH_CACHE")) |v| std.mem.span(v) else null); - if (!quiet and settings.on and !s.on) std.log.info("graph cache: off (TF_GRAPH_CACHE), every step eager", .{}); - s.min_keep = @min(s.min_keep, s.max); // the floor never holds more than the configured max + const s = settings.fromEnv(); + if (!quiet and settings.on and settings.max > 0 and !s.on) std.log.info("graph cache: off (TF_GRAPH_CACHE), every step eager", .{}); return .{ .gpa = gpa, .engine = engine, .settings = s, .cap = s.max }; } @@ -421,4 +444,5 @@ test "graph cache: fingerprint" { test { _ = @import("graph_cache_regression_test.zig"); + _ = @import("graph_cache_ranks_test.zig"); } diff --git a/zig/src/cuda/graph_cache_ranks_test.zig b/zig/src/cuda/graph_cache_ranks_test.zig new file mode 100644 index 0000000000..608ebe87a1 --- /dev/null +++ b/zig/src/cuda/graph_cache_ranks_test.zig @@ -0,0 +1,138 @@ +//! Two in-process ranks with fake engines: equal effective settings vote on misses only; a mismatch is a stall. + +const std = @import("std"); +const cache = @import("graph_cache.zig"); +const Stream = @import("stream.zig").Stream; +const no_stream: Stream = .{ .d = undefined, .handle = null }; +const equal = std.testing.expectEqual; +const Settings = cache.Settings; + +const Engine = struct { + captures: usize = 0, + live: usize = 0, + handles: [8]u8 = undefined, + + fn engine(e: *Engine) cache.Engine { + return .{ .ctx = e, .capture = capture, .launch = launch, .free = free }; + } + fn capture(p: *anyopaque, _: Stream, _: cache.Body) !*anyopaque { + const e: *Engine = @ptrCast(@alignCast(p)); + e.captures += 1; + e.live += 1; + return &e.handles[e.captures % e.handles.len]; + } + fn launch(_: *anyopaque, _: *anyopaque, _: Stream) !void {} + fn free(p: *anyopaque, _: *anyopaque) void { + const e: *Engine = @ptrCast(@alignCast(p)); + e.live -= 1; + } + fn body(e: *Engine) cache.Body { + return .{ .ctx = e, .run = run }; + } + fn run(_: *anyopaque, _: Stream) !void {} +}; + +/// The two ranks' votes in call order; a rank waits at most `wait_ns` for the other's vote of the same round. +const Board = struct { + votes: [2][16]bool = undefined, + cast: [2]std.atomic.Value(u32) = .{ .init(0), .init(0) }, + wait_ns: u64, + + fn vote(b: *Board, rank: u1, ok: bool) !bool { + const r = b.cast[rank].load(.monotonic); + if (r == b.votes[rank].len) return error.TooManyVotes; + b.votes[rank][r] = ok; + b.cast[rank].store(r + 1, .release); + const other = 1 - rank; + const deadline = now() + b.wait_ns; + while (b.cast[other].load(.acquire) <= r) { + if (now() > deadline) return error.PeerNeverVoted; + std.Thread.yield() catch {}; + } + return ok and b.votes[other][r]; + } + + fn now() u64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); + } +}; + +const Rank = struct { + board: *Board, + id: u1, + engine: Engine = .{}, + outcomes: [8]cache.Outcome = undefined, + err: ?anyerror = null, + + fn all(ctx: *anyopaque, ok: bool) anyerror!bool { + const r: *Rank = @ptrCast(@alignCast(ctx)); + return r.board.vote(r.id, ok); + } + + /// One rank's steps over `keys` with `settings` exactly as its process resolved them. + fn steps(r: *Rank, settings: Settings, ks: []const u32) void { + var c = cache.Cache(u32).init(std.testing.allocator, r.engine.engine(), settings); + defer c.deinit(); + c.agree = .{ .ctx = r, .all = all }; + for (ks, 0..) |k, i| r.outcomes[i] = c.run(k, no_stream, 7, r.engine.body()) catch |e| { + r.err = e; + return; + }; + } +}; + +const keys = [_]u32{ 1, 2, 1, 2, 1 }; + +test "graph cache ranks: equal enabled settings capture together and vote on each miss, never on a hit" { + var board: Board = .{ .wait_ns = 5 * std.time.ns_per_s }; + var ranks = [2]Rank{ .{ .board = &board, .id = 0 }, .{ .board = &board, .id = 1 } }; + const s = (Settings{ .max = 4 }).effective(null); + const t = try std.Thread.spawn(.{}, Rank.steps, .{ &ranks[1], s, &keys }); + ranks[0].steps(s, &keys); + t.join(); + const want = [_]cache.Outcome{ .captured, .captured, .replayed, .replayed, .replayed }; + for (&ranks) |*r| { + try equal(@as(?anyerror, null), r.err); + try std.testing.expectEqualSlices(cache.Outcome, &want, r.outcomes[0..keys.len]); + try equal(@as(u32, 4), board.cast[r.id].load(.acquire)); // two misses: room and capture each, hits none + try equal(@as(usize, 0), r.engine.live); + } +} + +test "graph cache ranks: TF_GRAPH_CACHE 0, off and OFF are one effective setting, and an off cache never votes" { + const d: Settings = .{}; + const off = d.effective("0").digest(); + try std.testing.expectEqualSlices(u8, &off, &d.effective("off").digest()); + try std.testing.expectEqualSlices(u8, &off, &d.effective("OFF").digest()); + try std.testing.expectEqualSlices(u8, &off, &(Settings{ .max = 0, .hold = true }).digest()); + try std.testing.expect(!std.mem.eql(u8, &off, &d.effective(null).digest())); + const held: Settings = .{ .hold = true, .min_keep = 2 }; + try std.testing.expectEqualSlices(u8, &held.digest(), &(Settings{ .hold = true, .min_keep = 30 }).digest()); + try std.testing.expectEqualSlices(u8, &d.effective("1").digest(), &d.effective(null).digest()); + var board: Board = .{ .wait_ns = 0 }; + var ranks = [2]Rank{ .{ .board = &board, .id = 0 }, .{ .board = &board, .id = 1 } }; + ranks[0].steps(d.effective("0"), &keys); + ranks[1].steps(d.effective("OFF"), &keys); + for (&ranks) |*r| { + try equal(@as(?anyerror, null), r.err); + for (r.outcomes[0..keys.len]) |o| try equal(cache.Outcome.eager, o); + try equal(@as(u32, 0), board.cast[r.id].load(.acquire)); + try equal(@as(usize, 0), r.engine.captures); + } +} + +test "graph cache ranks: one rank switched off leaves the other waiting at its first miss; digests show it first" { + const d: Settings = .{}; + try std.testing.expect(!std.mem.eql(u8, &d.effective(null).digest(), &d.effective("off").digest())); + var board: Board = .{ .wait_ns = 50 * std.time.ns_per_ms }; + var ranks = [2]Rank{ .{ .board = &board, .id = 0 }, .{ .board = &board, .id = 1 } }; + ranks[1].steps(d.effective("off"), &keys); + try equal(@as(?anyerror, null), ranks[1].err); + try equal(@as(u32, 0), board.cast[1].load(.acquire)); // the off rank ran every step eagerly without a vote + ranks[0].steps(d.effective(null), &keys); + try equal(@as(?anyerror, error.PeerNeverVoted), ranks[0].err); // in production this wait has no deadline + try equal(@as(u32, 1), board.cast[0].load(.acquire)); + try equal(@as(usize, 0), ranks[0].engine.live); +} From e2504f6d40ad3825e60146c8e456f17bc3f6a3bf Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Fri, 16 Oct 2026 11:28:48 +0700 Subject: [PATCH 32/43] TP: bootstrap compares the graph cache's effective settings, so a rank switched off refuses to start instead of stalling A family that installs the keyed graph cache sets Config.graph_cache before Session.open; the settings digest then holds their effective digest under TF_GRAPH_CACHE as read, so 0, off and OFF agree and on against off stops every rank at start with the existing mismatch message. Families without the cache leave it null and are not compared. Hits add no collective. --- RUNBOOK.md | 2 ++ zig/src/fabric/tp/bootstrap.zig | 49 +++++++++++++++++++++++++++++++++ zig/src/fabric/tp/config.zig | 32 +++++++++++++++++++++ zig/src/fabric/tp/tp.zig | 2 +- 4 files changed, 84 insertions(+), 1 deletion(-) diff --git a/RUNBOOK.md b/RUNBOOK.md index 4fd649e9be..abda311a50 100644 --- a/RUNBOOK.md +++ b/RUNBOOK.md @@ -161,6 +161,8 @@ The server logs `heat_wait_s`, a cancel ends the wait, and the reply's tokens do without a rebuild. Set it alike on every tensor-parallel rank: an off rank never joins the cache's per-miss vote, so a rank left on would wait at its first miss, and only the off rank logs the switch. Ranks that share a cache compare `Settings.digest` of their effective settings at start, before building it, and refuse to start when they differ. +A tensor-parallel family passes its cache settings in the TP config; the bootstrap then compares them and refuses to +start on every rank, naming TF_GRAPH_CACHE among the settings that must match. Archives whose names contain `host-only` are CPU verification artifacts and cannot serve CUDA inference. ## Flash Next and paired Metal serving diff --git a/zig/src/fabric/tp/bootstrap.zig b/zig/src/fabric/tp/bootstrap.zig index 1486e3a8f7..e1a6911efb 100644 --- a/zig/src/fabric/tp/bootstrap.zig +++ b/zig/src/fabric/tp/bootstrap.zig @@ -260,3 +260,52 @@ test "settings that differ stop every rank" { t.join(); try std.testing.expectEqual(@as(?anyerror, error.Mismatch), err); } + +/// Rank 1 agrees on `theirs.digest()` against rank 0's `ours.digest()`; returns both ranks' results. +fn agreeDigests(ours: Config, theirs: Config) !struct { ?anyerror, ?anyerror } { + const l = try testListener(); + defer sock.close(l.fd); + const T = struct { + fn go(port: u16, cfg: Config, err: *?anyerror) void { + var c = cfg; + c.world = 2; + c.rank = 1; + c.link_port = port; + c.connect_timeout_ns = 5 * std.time.ns_per_s; + var b = Bootstrap.open(c) catch |e| { + err.* = e; + return; + }; + defer b.close(); + b.agree(&c.digest(), "settings") catch |e| { + err.* = e; + }; + } + }; + var err1: ?anyerror = null; + const t = try std.Thread.spawn(.{}, T.go, .{ l.port, theirs, &err1 }); + var b = try Bootstrap.acceptOn(l.fd, 2, 5 * std.time.ns_per_s); + defer b.close(); + var err0: ?anyerror = null; + b.agree(&ours.digest(), "settings") catch |e| { + err0 = e; + }; + t.join(); + return .{ err0, err1 }; +} + +test "a graph cache switched off on one rank stops both at start; 0 and OFF on the two ranks start" { + const on: Config = .{ .graph_cache = .{} }; + const off: Config = .{ .graph_cache = .{}, .graph_cache_env = "off" }; + const r = try agreeDigests(on, off); + try std.testing.expectEqual(@as(?anyerror, error.Mismatch), r[0]); + try std.testing.expectEqual(@as(?anyerror, error.Mismatch), r[1]); + const none = try agreeDigests(.{}, .{ .graph_cache_env = "off" }); // a family without the cache is not refused + try std.testing.expectEqual(@as(?anyerror, null), none[0]); + try std.testing.expectEqual(@as(?anyerror, null), none[1]); + const same = try agreeDigests(.{ .graph_cache = .{}, .graph_cache_env = "0" }, .{ .graph_cache = .{}, .graph_cache_env = "OFF" }); + try std.testing.expectEqual(@as(?anyerror, null), same[0]); + try std.testing.expectEqual(@as(?anyerror, null), same[1]); + const family = try agreeDigests(.{ .graph_cache = .{ .max = 8 } }, .{ .graph_cache = .{ .max = 16 } }); + try std.testing.expectEqual(@as(?anyerror, error.Mismatch), family[0]); +} diff --git a/zig/src/fabric/tp/config.zig b/zig/src/fabric/tp/config.zig index 442f82ed0e..6f96091c1c 100644 --- a/zig/src/fabric/tp/config.zig +++ b/zig/src/fabric/tp/config.zig @@ -1,6 +1,7 @@ //! A rank's TP settings from the environment: rank, master, port, backend, fail-fast knobs. const std = @import("std"); const log = @import("log.zig"); +const graph_cache = @import("cuda").graph_cache; pub const Backend = enum { /// NCCL for everything (the default). @@ -64,6 +65,10 @@ pub const Config = struct { /// TF_TP_PLAN_LINK=rdma: the follower polls the plan socket this long before blocking (TF_TP_PLAN_SPIN_US). plan_spin_ns: u64 = 0, plan_pin: ?PlanPin = null, + /// The keyed graph cache settings a family builds with, set before `Session.open` (null: the family has none). + graph_cache: ?graph_cache.Settings = null, + /// TF_GRAPH_CACHE as read (`Cache.init` applies the same value); only its effect on `graph_cache` is compared. + graph_cache_env: ?[]const u8 = null, /// Reads every setting from `get` (the environment in `fromEnv`); unset names keep the defaults. pub fn parse(get: *const fn ([]const u8) ?[]const u8) Error!Config { @@ -131,6 +136,7 @@ pub const Config = struct { c.plan_pin = .{ .cpu = try int(u32, get, "TF_TP_PLAN_PIN", 0, 0, 4095) }; } } + c.graph_cache_env = get("TF_GRAPH_CACHE"); return c; } @@ -148,6 +154,7 @@ pub const Config = struct { h.update(std.mem.asBytes(&c.failfast)); h.update(std.mem.asBytes(&c.roce_wire)); h.update(std.mem.asBytes(&c.roce_fallback)); + if (c.graph_cache) |g| h.update(&g.effective(c.graph_cache_env).digest()); var out: [32]u8 = undefined; h.final(&out); return out; @@ -254,3 +261,28 @@ test "the plan link's spin and pin, the RoCE grid cap" { try std.testing.expectError(error.BadSetting, Config.parse(&Fake.get)); Fake.pairs = &.{}; } + +test "the digest holds the graph cache's effective settings: 0, off and OFF agree, on and off differ" { + Fake.pairs = &.{.{ "TF_GRAPH_CACHE", "0" }}; + var zero = try Config.parse(&Fake.get); + Fake.pairs = &.{.{ "TF_GRAPH_CACHE", "OFF" }}; + var upper = try Config.parse(&Fake.get); + Fake.pairs = &.{.{ "TF_GRAPH_CACHE", "1" }}; + var one = try Config.parse(&Fake.get); + Fake.pairs = &.{}; + var unset = try Config.parse(&Fake.get); + try std.testing.expectEqualSlices(u8, &zero.digest(), &unset.digest()); // no cache installed: not compared + for ([_]*Config{ &zero, &upper, &one, &unset }) |c| c.graph_cache = .{}; + var max0 = unset; + max0.graph_cache = .{ .max = 0 }; + try std.testing.expectEqualSlices(u8, &zero.digest(), &upper.digest()); + try std.testing.expectEqualSlices(u8, &zero.digest(), &max0.digest()); + try std.testing.expectEqualSlices(u8, &one.digest(), &unset.digest()); + try std.testing.expect(!std.mem.eql(u8, &zero.digest(), &unset.digest())); + var family = unset; + family.graph_cache = .{ .max = 8, .hold = true }; + try std.testing.expect(!std.mem.eql(u8, &family.digest(), &unset.digest())); + var family_off = upper; + family_off.graph_cache = family.graph_cache; + try std.testing.expectEqualSlices(u8, &zero.digest(), &family_off.digest()); +} diff --git a/zig/src/fabric/tp/tp.zig b/zig/src/fabric/tp/tp.zig index 6cc5ceb6ec..2c9eeb85ec 100644 --- a/zig/src/fabric/tp/tp.zig +++ b/zig/src/fabric/tp/tp.zig @@ -41,7 +41,7 @@ pub const Session = struct { pub fn open(s: *Session, d: *const cuda.Driver, cfg: Config, kernel_image: []const u8) !void { var b = try Bootstrap.open(cfg); defer b.close(); - try b.agree(&cfg.digest(), "settings (TF_TP_WORLD, TF_COMM_BACKEND, TF_TP_MAILBOX_MAX_KB, TF_TP_ROCE_CHUNKED_KB, TF_TP_FAILFAST, TF_TP_ROCE_WIRE / _FALLBACK)"); + try b.agree(&cfg.digest(), "settings (TF_TP_WORLD, TF_COMM_BACKEND, TF_TP_MAILBOX_MAX_KB, TF_TP_ROCE_CHUNKED_KB, TF_TP_FAILFAST, TF_TP_ROCE_WIRE / _FALLBACK, TF_GRAPH_CACHE)"); var lib = try cuda.nccl.Library.openPath(cfg.nccl_lib); errdefer lib.close(); const v = try lib.version(); From dcdc30c9e79046ccfd7605bed34dccb04b5c259d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Schr=C3=B6der?= <1188270+cwschroeder@users.noreply.github.com> Date: Fri, 16 Oct 2026 10:01:36 +0200 Subject: [PATCH 33/43] Keep the Linux NVMe lane test out of portable host tests --- zig/src/sessions/lanes.zig | 1 + 1 file changed, 1 insertion(+) diff --git a/zig/src/sessions/lanes.zig b/zig/src/sessions/lanes.zig index 38fad8cd89..2b005aca9a 100644 --- a/zig/src/sessions/lanes.zig +++ b/zig/src/sessions/lanes.zig @@ -293,6 +293,7 @@ pub fn preadAll(fd: linux.fd_t, out: []u8, off: u64) !void { } test "lanes write, read back and check chunks through bounded staging" { + if (@import("builtin").os.tag != .linux) return error.SkipZigTest; const gpa = std.testing.allocator; const io = std.testing.io; var st = try Staging.init(gpa, io, 3, 2 * block); From 368a0d384d8558478dfa5b14523127e8e37cb554 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Schr=C3=B6der?= <1188270+cwschroeder@users.noreply.github.com> Date: Fri, 16 Oct 2026 10:01:36 +0200 Subject: [PATCH 34/43] Keep the Linux NVMe lane test out of portable host tests --- zig/src/sessions/lanes.zig | 1 + 1 file changed, 1 insertion(+) diff --git a/zig/src/sessions/lanes.zig b/zig/src/sessions/lanes.zig index 38fad8cd89..2b005aca9a 100644 --- a/zig/src/sessions/lanes.zig +++ b/zig/src/sessions/lanes.zig @@ -293,6 +293,7 @@ pub fn preadAll(fd: linux.fd_t, out: []u8, off: u64) !void { } test "lanes write, read back and check chunks through bounded staging" { + if (@import("builtin").os.tag != .linux) return error.SkipZigTest; const gpa = std.testing.allocator; const io = std.testing.io; var st = try Staging.init(gpa, io, 3, 2 * block); From 7a6b32f6f1f3d24dc20d376cb81fdabaa33dac70 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Fri, 16 Oct 2026 19:39:43 +0700 Subject: [PATCH 35/43] server: main's drain stop closes the engine without the closed lane-host halt Opened.halt came with the closed #662 and is absent from this stack, so the native build failed to compile at main.zig. --- zig/src/server/main.zig | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/zig/src/server/main.zig b/zig/src/server/main.zig index 2b0de6f365..434329759d 100644 --- a/zig/src/server/main.zig +++ b/zig/src/server/main.zig @@ -80,7 +80,7 @@ pub fn main(init: std.process.Init) !u8 { var closer: Closer = .{ .opened = up.engine }; defer closer.closeOnce(); return serve.run(gpa, io, args, .{ - .stop = .{ .ctx = &closer, .halt = if (up.engine.halt != null) Closer.halt else null, .close = Closer.close }, + .stop = .{ .ctx = &closer, .close = Closer.close }, .engine = up.engine.engine, .text = up.text.text(), .served = hub.servedName(args.name, args.model, dir), @@ -105,11 +105,6 @@ const Closer = struct { fn close(ctx: *anyopaque) void { closeOnce(@ptrCast(@alignCast(ctx))); } - - fn halt(ctx: *anyopaque, reason: []const u8) void { - const c: *Closer = @ptrCast(@alignCast(ctx)); - if (c.opened.halt) |h| h(c.opened.ctx, reason); - } }; /// HfText.load with every failure named in `problem`, so main tells the text's failure from the engine's. From af0bc5378687c188a2461e080db4b9db72ec9972 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Fri, 16 Oct 2026 23:06:58 +0700 Subject: [PATCH 36/43] Server tests: the drain test imports the named json module Main builds zig/src/server/json.zig as the `json` module, so a relative import puts the file in two modules and `zig build test` / `test-server-cpu` fail to compile. Co-Authored-By: Claude Opus 5.5 --- zig/src/server/drain_test.zig | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zig/src/server/drain_test.zig b/zig/src/server/drain_test.zig index 09873e92f5..71b9514baf 100644 --- a/zig/src/server/drain_test.zig +++ b/zig/src/server/drain_test.zig @@ -1,7 +1,7 @@ //! A stop's drain over HTTP: serve.run with a real SIGTERM, a slow stub engine, bytes as tokens. const std = @import("std"); const api = @import("engine_api"); -const json = @import("json.zig"); +const json = @import("json"); const model_text = @import("model_text.zig"); const serve = @import("serve.zig"); const cli = @import("cli.zig"); From 8900c7cc59e617fd0ad267a40bffabeccffab098 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:16:48 +0700 Subject: [PATCH 37/43] Sessions: a failed prefix insert leaves the index as it was newNode and place reserve the maps, the slot lists and the free lists before they change anything. A failed insertion no longer consumes a free-list ID or leaves an appended slot undefined (deinit read that slot's flag). The free lists keep room for every slot, so remove and prune never allocate or drop an ID. A failed insert or insertFrom prunes the path nodes it made. A failed setRam in insert drops the entry instead of leaving it indexed, and setRam returns the token blocks it took. prefix_fail_test.zig fails each allocation of a RAM insert and of a disk insert into an empty index in turn. Each time, the index must keep its shape and slot accounting, the same insert must then succeed, and deinit must leak nothing. The tests run in the Linux and the portable (macOS) session roots. Spark receipt: lane job doug-637-prefix-01 on dgx-01 (aarch64, Zig 0.17.0), 2026-10-17T02:08:30Z, VERDICT PASS, candidate tree 1322aeeff84f7dd3af075e524d42b79ddd41fcd2. zig fmt --check clean; lean 517/517 against merge-base 46a08c18, 0 new; zig build 16/16; test test-kernel-contract test-golden fuzz-cortheon 99/99 steps, 1499/1507 (8 skipped); test-sessions 3/3. With only the new tests on 9266f3f1, both fail at the node-slot accounting check. Co-Authored-By: Claude Opus 5.5 --- zig/src/sessions/portable_test.zig | 1 + zig/src/sessions/prefix.zig | 50 +++++++--- zig/src/sessions/prefix_fail_test.zig | 127 ++++++++++++++++++++++++++ zig/src/sessions/sessions.zig | 1 + 4 files changed, 167 insertions(+), 12 deletions(-) create mode 100644 zig/src/sessions/prefix_fail_test.zig diff --git a/zig/src/sessions/portable_test.zig b/zig/src/sessions/portable_test.zig index 6ba2694b4e..2e70055c68 100644 --- a/zig/src/sessions/portable_test.zig +++ b/zig/src/sessions/portable_test.zig @@ -4,5 +4,6 @@ test { _ = @import("pool.zig"); _ = @import("pagestore.zig"); _ = @import("prefix.zig"); + _ = @import("prefix_fail_test.zig"); _ = @import("format.zig"); } diff --git a/zig/src/sessions/prefix.zig b/zig/src/sessions/prefix.zig index 98634bbfd2..80c5838c60 100644 --- a/zig/src/sessions/prefix.zig +++ b/zig/src/sessions/prefix.zig @@ -122,13 +122,19 @@ pub fn Index(comptime Extra: type) type { return x.by_key.get(key); } + /// Reserves first, so a failure changes nothing; free_nodes keeps room for every slot (prune never allocates). fn newNode(x: *Self, digest: Digest, parent: u32, depth: u32) !u32 { - const id: u32 = if (x.free_nodes.pop()) |i| i else blk: { - try x.nodes.append(x.gpa, undefined); + try x.by_digest.ensureUnusedCapacity(x.gpa, 1); + if (x.free_nodes.items.len == 0) { + try x.nodes.ensureUnusedCapacity(x.gpa, 1); + try x.free_nodes.ensureTotalCapacity(x.gpa, x.nodes.items.len + 1); + } + const id: u32 = x.free_nodes.pop() orelse blk: { + x.nodes.appendAssumeCapacity(undefined); break :blk @intCast(x.nodes.items.len - 1); }; x.nodes.items[id] = .{ .digest = digest, .parent = parent, .depth = depth }; - try x.by_digest.put(x.gpa, digest, id); + x.by_digest.putAssumeCapacityNoClobber(digest, id); if (parent != none) { const p = &x.nodes.items[parent]; x.nodes.items[id].next_sib = p.first_child; @@ -143,18 +149,22 @@ pub fn Index(comptime Extra: type) type { return x.gpa.alloc(i32, x.page); } - /// Adds an entry for ids (chain: its full pages' digests), or returns the one with the same key. + /// Adds an entry for ids (chain: its full pages' digests) or returns the one with its key; an error adds none. pub fn insert(x: *Self, tag: u32, ids: []const i32, digests: []const Digest, ram: bool, extra: Extra) !Placed { const full: u32 = @intCast(ids.len / x.page); std.debug.assert(digests.len >= full); const r = root(tag); var n = x.by_digest.get(r) orelse try x.newNode(r, none, 0); + errdefer x.prune(n); // the path's new nodes, which no entry holds yet for (0..full) |j| { const d = digests[j]; n = x.by_digest.get(d) orelse try x.newNode(d, n, @intCast(j + 1)); } const r2 = try x.place(tag, n, ids.len, ids[full * x.page ..], extra); - if (r2.new and ram) try x.setRam(r2.id, true, ids); + if (r2.new and ram) x.setRam(r2.id, true, ids) catch |err| { + _ = x.drop(r2.id); // the errdefer above prunes its path + return err; + }; return r2; } @@ -164,6 +174,7 @@ pub fn Index(comptime Extra: type) type { const r = root(tag); break :blk x.by_digest.get(r) orelse try x.newNode(r, none, 0); }; + errdefer x.prune(n); // the path's new nodes, which no entry holds yet const d0 = x.nodes.items[n].depth; const full: u32 = @intCast(suffix.len / x.page); for (0..full) |j| { @@ -183,13 +194,17 @@ pub fn Index(comptime Extra: type) type { fn place(x: *Self, tag: u32, n: u32, pos: u64, tail_ids: []const i32, extra: Extra) !Placed { const key = entryKey(x.nodes.items[n].digest, pos, tail_ids); if (x.by_key.get(key)) |e| return .{ .id = e, .new = false }; - const tail = try x.gpa.dupe(i32, tail_ids); - errdefer x.gpa.free(tail); - const id: u32 = if (x.free_entries.pop()) |i| i else blk: { - try x.entries.append(x.gpa, undefined); + try x.by_key.ensureUnusedCapacity(x.gpa, 1); + if (x.free_entries.items.len == 0) { + try x.entries.ensureUnusedCapacity(x.gpa, 1); + try x.free_entries.ensureTotalCapacity(x.gpa, x.entries.items.len + 1); + } + const tail = try x.gpa.dupe(i32, tail_ids); // the last step that can fail + const id: u32 = x.free_entries.pop() orelse blk: { + x.entries.appendAssumeCapacity(undefined); break :blk @intCast(x.entries.items.len - 1); }; - try x.by_key.put(x.gpa, key, id); + x.by_key.putAssumeCapacityNoClobber(key, id); const nd = &x.nodes.items[n]; x.entries.items[id] = .{ .key = key, .tag = tag, .pos = pos, .node = n, .tail = tail, .ram = false, .next_at = nd.first_entry, .extra = extra }; if (nd.first_entry != none) x.entries.items[nd.first_entry].prev_at = id; @@ -214,6 +229,7 @@ pub fn Index(comptime Extra: type) type { if (e.ram == ram) return; if (ram) { const all = ids orelse return error.IdsNeeded; + errdefer x.dropTokens(e.node); // the blocks this call took: tokens stay only where sub_ram > 0 var n = e.node; while (n != none and x.nodes.items[n].depth > 0) : (n = x.nodes.items[n].parent) { const nd = &x.nodes.items[n]; @@ -243,6 +259,11 @@ pub fn Index(comptime Extra: type) type { /// Removes an entry and the nodes no entry needs any more. pub fn remove(x: *Self, id: u32) void { + x.prune(x.drop(id)); + } + + /// Unlinks an entry and frees its slot; returns its node, whose path may now hold no entry. + fn drop(x: *Self, id: u32) u32 { x.setRam(id, false, null) catch unreachable; const e = &x.entries.items[id]; const n = e.node; @@ -251,8 +272,13 @@ pub fn Index(comptime Extra: type) type { _ = x.by_key.remove(e.key); x.gpa.free(e.tail); e.live = false; - x.free_entries.append(x.gpa, id) catch {}; + x.free_entries.appendAssumeCapacity(id); // place keeps room for every slot x.addPath(n, -1, 0); + return n; + } + + /// Frees the nodes from n up that no entry holds (sub_all 0). + fn prune(x: *Self, n: u32) void { var m = n; while (m != none and x.nodes.items[m].sub_all == 0) { const nd = x.nodes.items[m]; @@ -261,7 +287,7 @@ pub fn Index(comptime Extra: type) type { if (nd.next_sib != none) x.nodes.items[nd.next_sib].prev_sib = nd.prev_sib; } if (x.by_digest.remove(nd.digest)) x.dropped(); - x.free_nodes.append(x.gpa, m) catch {}; + x.free_nodes.appendAssumeCapacity(m); // newNode keeps room for every slot m = nd.parent; } } diff --git a/zig/src/sessions/prefix_fail_test.zig b/zig/src/sessions/prefix_fail_test.zig new file mode 100644 index 0000000000..dd0d8ec489 --- /dev/null +++ b/zig/src/sessions/prefix_fail_test.zig @@ -0,0 +1,127 @@ +//! prefix.Index under allocation failure: an insert that fails changes nothing, and the same insert then succeeds. + +const std = @import("std"); +const prefix = @import("prefix.zig"); +const testing = std.testing; +const T = prefix.Index(u64); +const Digest = prefix.Digest; +const page = 16; + +/// What a failed insert must leave as it was; spare token blocks may stay in the reuse pool. +const Shape = struct { keys: u32, digests: u32, ram: u32, held: u32, sub_all: u32, sub_ram: u32 }; + +fn shape(x: *T, tag: u32) Shape { + var s: Shape = .{ .keys = x.by_key.count(), .digests = x.by_digest.count(), .ram = 0, .held = 0, .sub_all = 0, .sub_ram = 0 }; + var e = x.by_key.valueIterator(); + while (e.next()) |id| s.ram += @intFromBool(x.get(id.*).ram); + var n = x.by_digest.valueIterator(); + while (n.next()) |id| s.held += @intFromBool(x.node(id.*).tokens != null); + if (x.by_digest.get(prefix.root(tag))) |r| { + s.sub_all = x.node(r).sub_all; + s.sub_ram = x.node(r).sub_ram; + } + return s; +} + +/// Every slot is live in its map or on its free list, and a node keeps its page's tokens exactly while sub_ram > 0. +fn consistent(x: *T) !void { + try testing.expectEqual(x.entries.items.len - x.free_entries.items.len, x.by_key.count()); + try testing.expectEqual(x.nodes.items.len - x.free_nodes.items.len, x.by_digest.count()); + var n = x.by_digest.valueIterator(); + while (n.next()) |id| { + const nd = x.node(id.*); + if (nd.depth > 0) try testing.expectEqual(nd.sub_ram > 0, nd.tokens != null); + } +} + +fn seq(gpa: std.mem.Allocator, n: usize, salt: i32) ![]i32 { + const v = try gpa.alloc(i32, n); + for (v, 0..) |*t, i| t.* = @as(i32, @intCast(i % 997)) * 3 + salt; + return v; +} + +fn add(x: *T, tag: u32, ids: []const i32, ram: bool, used: u64) !u32 { + var d: [64]Digest = undefined; + prefix.chain(tag, ids, page, d[0 .. ids.len / page]); + return (try x.insert(tag, ids, d[0 .. ids.len / page], ram, used)).id; +} + +/// The (k + 1)-th allocation from now fails, and every one after it until `heal`. +fn failAfter(fa: *testing.FailingAllocator, k: usize) void { + fa.fail_index = fa.alloc_index + k; +} + +fn heal(fa: *testing.FailingAllocator) void { + fa.fail_index = std.math.maxInt(usize); +} + +test "a RAM insert failing at any allocation leaves the index as it was, and the same insert then succeeds" { + const chat = try seq(testing.allocator, 160, 0); + defer testing.allocator.free(chat); + var d: [16]Digest = undefined; + prefix.chain(7, chat[0..150], page, d[0..9]); + var k: usize = 0; + while (k < 256) : (k += 1) { + var fa = testing.FailingAllocator.init(testing.allocator, .{}); + var x = T.init(fa.allocator(), page); + defer x.deinit(); + // what the insert meets: entries sharing its first pages, a disk entry, another tag, and free slots to reuse + _ = try add(&x, 7, chat[0..40], true, 1); + _ = try add(&x, 7, chat[0..72], false, 2); + _ = try add(&x, 8, chat[0..33], true, 3); + x.remove(try add(&x, 9, chat[0..48], true, 4)); + try consistent(&x); + const before = shape(&x, 7); + failAfter(&fa, k); + const got = x.insert(7, chat[0..150], d[0..9], true, 9) catch |err| { + try testing.expectEqual(error.OutOfMemory, err); + try consistent(&x); + try testing.expectEqual(before, shape(&x, 7)); + heal(&fa); + const again = try x.insert(7, chat[0..150], d[0..9], true, 9); + try testing.expect(again.new); + try consistent(&x); + var out: [150]i32 = undefined; + try x.idsOf(again.id, &out); + try testing.expectEqualSlices(i32, chat[0..150], &out); + var d10: [10]Digest = undefined; + prefix.chain(7, chat[0..160], page, &d10); + try testing.expectEqual(@as(?u32, again.id), x.find(7, chat[0..160], &d10, .ram)); + x.remove(again.id); + try consistent(&x); + try testing.expectEqual(before, shape(&x, 7)); + continue; + }; + // no allocation left to fail: every failure point before this one was injected and rolled back + try testing.expect(got.new and k > 0); + return; + } + return error.TestUnexpectedResult; +} + +test "a disk insert into an empty index failing at any allocation leaves it empty, and the same insert then succeeds" { + const chat = try seq(testing.allocator, 90, 0); + defer testing.allocator.free(chat); + var k: usize = 0; + while (k < 256) : (k += 1) { + var fa = testing.FailingAllocator.init(testing.allocator, .{}); + var x = T.init(fa.allocator(), page); + defer x.deinit(); + failAfter(&fa, k); + const got = x.insertFrom(7, prefix.none, chat, 5) catch |err| { + try testing.expectEqual(error.OutOfMemory, err); + try consistent(&x); + try testing.expectEqual(@as(u32, 0), x.by_digest.count()); // the tag's root and pages made before the failure + try testing.expectEqual(@as(u32, 0), x.by_key.count()); + heal(&fa); + const again = try x.insertFrom(7, prefix.none, chat, 5); + try testing.expect(again.new); + try testing.expectEqual(@as(u32, 1 + 5), x.by_digest.count()); + try consistent(&x); + continue; + }; + try testing.expect(got.new and k > 0); + return; + } + return error.TestUnexpectedResult; +} diff --git a/zig/src/sessions/sessions.zig b/zig/src/sessions/sessions.zig index 7bdc46b386..d13bef6029 100644 --- a/zig/src/sessions/sessions.zig +++ b/zig/src/sessions/sessions.zig @@ -28,6 +28,7 @@ test { _ = disk; _ = jobs; _ = store; + _ = @import("prefix_fail_test.zig"); _ = @import("disk_test.zig"); _ = @import("store_test.zig"); } From 3aeebb97b18f3799c83069d48fd6a3f36f150824 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:19:52 +0700 Subject: [PATCH 38/43] CUDA arena: a reuse whose address-map insertion fails keeps its range free allocLocked reserves room in `live` before it takes a range. A failed insertion used to pop a freed range and then lose it: the range was in neither the free list nor the live map. Both small-path inserts now cannot fail. A failure leaves the free list, the cursor and the chunks as they were, and the retry gets the same range. The new test spends the map's spare room so that the reuse's insertion must grow it, fails that allocation, and checks range, cursor and chunk conservation, the retry, and a clean deinit with every chunk released. Not in this change: tying memory.setHook / currentHook installation to the family that owns the arena. Jonathan Regan-Beasley's capture-ownership fix (5826f89b) is unchanged. Spark receipt: lane job doug-635-arena-01 on dgx-02 (aarch64, Zig 0.17.0), 2026-10-17T02:10:06Z, VERDICT PASS, candidate tree 663ff7bc2e2d0fa2eb404d5bc9d31eb55696bc68. zig fmt --check clean; lean 517/517 against merge-base 46a08c18, 0 new; zig build 12/12; test test-kernel-contract test-golden fuzz-cortheon 94/94 steps, 1430/1437 (7 skipped); test-cuda-host 39/39. With only the new test on abcad7d4, it fails (the freed range left its free list). Co-Authored-By: Claude Opus 5.5 --- zig/src/cuda/arena.zig | 35 +++++++++++++++++++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/zig/src/cuda/arena.zig b/zig/src/cuda/arena.zig index 3886c81403..de3da5718c 100644 --- a/zig/src/cuda/arena.zig +++ b/zig/src/cuda/arena.zig @@ -86,8 +86,9 @@ pub const Arena = struct { return p; } const n: u64 = roundUp(len, a.o.alignment); + try a.live.ensureUnusedCapacity(a.gpa, 1); // before a range is taken: a failure leaves free list and cursor alone if (a.free_small.getPtr(n)) |l| if (l.pop()) |p| { - try a.live.put(a.gpa, p, n); + a.live.putAssumeCapacityNoClobber(p, n); a.stats.reused += 1; return p; }; @@ -104,7 +105,7 @@ pub const Arena = struct { a.stats.chunk_bytes += cl; } const p = a.cursor; - try a.live.put(a.gpa, p, n); + a.live.putAssumeCapacityNoClobber(p, n); a.cursor += n; a.stats.small += 1; return p; @@ -372,3 +373,33 @@ test "arena: a long run of allocations and frees leaves both address maps free s } } } + +test "arena: a reuse whose address-map insertion fails keeps its range free, and the retry gets it" { + const t = std.testing; + var f: Fake = .{}; + var fa = t.FailingAllocator.init(t.allocator, .{}); + { + var a = Arena.init(fa.allocator(), f.backend(2 << 20), .{ .small_max = 1 << 20, .small_chunk = 4 << 20 }); + defer a.deinit(); // testing.allocator underneath reports any leaked host bytes + const p = a.alloc(4096).?; + try t.expect(a.free(p)); + // other sizes take back the room the free gave: inserting the reused range must grow the map + while (a.live.available > 0) _ = a.alloc(512).?; + const live = a.live.count(); + const cursor = a.cursor; + const chunks = f.chunks; + fa.fail_index = fa.alloc_index; // the next allocation fails, and every one after it until healed + try t.expectEqual(@as(?u64, null), a.alloc(4096)); + try t.expectEqual(live, a.live.count()); + try t.expectEqual(cursor, a.cursor); + try t.expectEqual(chunks, f.chunks); + try t.expectEqualSlices(u64, &.{p}, a.free_small.get(4096).?.items); + try t.expectEqual(@as(u64, 0), a.stats.reused); + fa.fail_index = std.math.maxInt(usize); + try t.expectEqual(p, a.alloc(4096).?); + try t.expectEqual(@as(u64, 1), a.stats.reused); + try t.expectEqual(@as(usize, 0), a.free_small.get(4096).?.items.len); + try t.expect(a.free(p)); + } + try t.expectEqual(f.chunks, f.released); // every chunk went back to the backend +} From 88418d8d0f4ecb875133b4da2a4de51faf3d4594 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Schr=C3=B6der?= Date: Sat, 17 Oct 2026 15:59:53 +0200 Subject: [PATCH 39/43] test(tp): reject malformed runner commands before configuration (cherry picked from commit 854b57cf74b90d306626f35e40de61c047574f72) Reported-by: Bizuayeu Co-Authored-By: Claude Opus 5.5 --- zig/build/tp.zig | 16 +++++++++ zig/src/fabric/tp/gpu_cli.zig | 66 +++++++++++++++++++++++++++++++++++ 2 files changed, 82 insertions(+) create mode 100644 zig/src/fabric/tp/gpu_cli.zig diff --git a/zig/build/tp.zig b/zig/build/tp.zig index 873127c6af..276f285ca1 100644 --- a/zig/build/tp.zig +++ b/zig/build/tp.zig @@ -41,6 +41,10 @@ fn kernelModule(b: *std.Build, image: ?std.Build.LazyPath) *std.Build.Module { pub fn add(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, test_step: *std.Build.Step) void { const host = b.graph.host; + const cli_tests = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/tp/gpu_cli.zig"), .target = host, .optimize = optimize }) })); + const cli_step = b.step("test-tp-cli", "TP runner command validation (no GPU)"); + cli_step.dependOn(&cli_tests.step); + test_step.dependOn(&cli_tests.step); if (host.result.os.tag == .linux) { const tests = b.addRunArtifact(b.addTest(.{ .root_module = module(b, host, .debug, runtime(b, host, .debug)) })); b.step("test-tp", "Tensor parallelism's host tests: bootstrap, fate channel, in-process ranks (no GPU)").dependOn(&tests.step); @@ -56,6 +60,18 @@ pub fn add(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti runner.addImport("tp", module(b, target, optimize, cuda)); runner.addImport("tp_kernel", kernelModule(b, image)); const exe = b.addExecutable(.{ .name = "tf-tp-test", .root_module = runner }); + if (host.result.os.tag == .linux and target.query.isNative()) { + const cli = @import("../src/fabric/tp/gpu_cli.zig"); + for (cli.invalid) |args| { + const check = b.addRunArtifact(exe); + check.addArgs(args); + check.setEnvironmentVariable("TF_TP_WORLD", "invalid"); + check.expectExitCode(2); + check.expectStdErrEqual(cli.usage); + cli_step.dependOn(&check.step); + test_step.dependOn(&check.step); + } + } b.installArtifact(exe); b.step("tf-tp-test", "TP=2 on two GPUs: bootstrap, every collective bit-exact, graph capture, fail-fast, latency").dependOn(&b.addInstallArtifact(exe, .{}).step); } diff --git a/zig/src/fabric/tp/gpu_cli.zig b/zig/src/fabric/tp/gpu_cli.zig new file mode 100644 index 0000000000..dc576fc03b --- /dev/null +++ b/zig/src/fabric/tp/gpu_cli.zig @@ -0,0 +1,66 @@ +//! Validate the runner's command before opening a driver, bootstrap or fate channel. +const std = @import("std"); + +pub const usage = "usage: tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | hca | kernel\n"; +pub const Mode = enum { fatal, kill, transport }; +pub const Failure = struct { victim: u32, mode: Mode, n: u32 = 50, hold: ?u32 = null }; +pub const Command = union(enum) { suite, failfast: Failure, @"roce-host", hca, kernel }; + +pub fn parse(args: []const []const u8) error{InvalidArguments}!Command { + if (args.len == 0) return error.InvalidArguments; + const tag = std.meta.stringToEnum(std.meta.Tag(Command), args[0]) orelse return error.InvalidArguments; + switch (tag) { + .suite, .@"roce-host", .hca, .kernel => { + if (args.len != 1) return error.InvalidArguments; + return switch (tag) { + .suite => .suite, + .@"roce-host" => .@"roce-host", + .hca => .hca, + .kernel => .kernel, + else => unreachable, + }; + }, + .failfast => { + if (args.len < 3 or args.len > 5) return error.InvalidArguments; + const mode = std.meta.stringToEnum(Mode, args[2]) orelse return error.InvalidArguments; + if (mode != .transport and args.len != 3) return error.InvalidArguments; + return .{ .failfast = .{ + .victim = try number(args[1]), + .mode = mode, + .n = if (args.len > 3) try number(args[3]) else 50, + .hold = if (args.len > 4) try number(args[4]) else null, + } }; + }, + } +} + +fn number(s: []const u8) error{InvalidArguments}!u32 { + return std.fmt.parseInt(u32, s, 10) catch error.InvalidArguments; +} + +// Also executed against the Linux entry point with an invalid environment, before any CUDA work. +pub const invalid = [_][]const []const u8{ + &.{}, &.{"UNKNOWN"}, &.{ "suite", "extra" }, &.{ "kernel", "extra" }, + &.{ "hca", "extra" }, &.{ "roce-host", "extra" }, &.{"failfast"}, &.{ "failfast", "0" }, + &.{ "failfast", "0", "typo" }, &.{ "failfast", "bad", "fatal" }, &.{ "failfast", "-1", "kill" }, &.{ "failfast", "4294967296", "fatal" }, + &.{ "failfast", "0", "fatal", "50" }, &.{ "failfast", "0", "kill", "50" }, &.{ "failfast", "0", "transport", "bad" }, &.{ "failfast", "0", "transport", "50", "bad" }, + &.{ "failfast", "0", "transport", "50", "0", "extra" }, +}; + +test "invalid commands are refused before resources open" { + for (invalid) |args| try std.testing.expectError(error.InvalidArguments, parse(args)); +} + +test "documented commands and failfast defaults" { + inline for (.{ "suite", "roce-host", "hca", "kernel" }) |name| { + try std.testing.expectEqualStrings(name, @tagName(try parse(&.{name}))); + } + inline for (.{ Mode.fatal, Mode.kill, Mode.transport }) |mode| { + const f = (try parse(&.{ "failfast", "2", @tagName(mode) })).failfast; + try std.testing.expectEqualDeep(Failure{ .victim = 2, .mode = mode }, f); + } + const explicit = (try parse(&.{ "failfast", "63", "transport", "0", "0" })).failfast; + try std.testing.expectEqualDeep(Failure{ .victim = 63, .mode = .transport, .n = 0, .hold = 0 }, explicit); + const n = (try parse(&.{ "failfast", "1", "transport", "123" })).failfast; + try std.testing.expectEqualDeep(Failure{ .victim = 1, .mode = .transport, .n = 123 }, n); +} From 4c69b0d4b434b782a785130ae649f4337f931f91 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Schr=C3=B6der?= Date: Sat, 17 Oct 2026 16:04:58 +0200 Subject: [PATCH 40/43] fix(tp): parse runner commands before opening sessions (cherry picked from commit 602a0fd4e1d621cf0715d01a2b466497dc3b4bb7) Reported-by: Bizuayeu Co-Authored-By: Claude Opus 5.5 --- zig/build/tp.zig | 13 ++++++++++++- zig/src/fabric/tp/gpu_main.zig | 32 ++++++++++++++++---------------- 2 files changed, 28 insertions(+), 17 deletions(-) diff --git a/zig/build/tp.zig b/zig/build/tp.zig index 276f285ca1..e07230abbc 100644 --- a/zig/build/tp.zig +++ b/zig/build/tp.zig @@ -47,7 +47,9 @@ pub fn add(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti test_step.dependOn(&cli_tests.step); if (host.result.os.tag == .linux) { const tests = b.addRunArtifact(b.addTest(.{ .root_module = module(b, host, .debug, runtime(b, host, .debug)) })); - b.step("test-tp", "Tensor parallelism's host tests: bootstrap, fate channel, in-process ranks (no GPU)").dependOn(&tests.step); + const tp_step = b.step("test-tp", "Tensor parallelism's host tests: bootstrap, fate channel, in-process ranks (no GPU)"); + tp_step.dependOn(&tests.step); + tp_step.dependOn(cli_step); test_step.dependOn(&tests.step); } if (target.result.os.tag != .linux) return; @@ -71,6 +73,15 @@ pub fn add(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti cli_step.dependOn(&check.step); test_step.dependOn(&check.step); } + if (image == null) { + const kernel = b.addRunArtifact(exe); + kernel.addArg("kernel"); + kernel.setEnvironmentVariable("TF_TP_WORLD", "invalid"); + kernel.expectExitCode(1); + kernel.expectStdErrEqual("mailbox kernel 0 B\n"); + cli_step.dependOn(&kernel.step); + test_step.dependOn(&kernel.step); + } } b.installArtifact(exe); b.step("tf-tp-test", "TP=2 on two GPUs: bootstrap, every collective bit-exact, graph capture, fail-fast, latency").dependOn(&b.addInstallArtifact(exe, .{}).step); diff --git a/zig/src/fabric/tp/gpu_main.zig b/zig/src/fabric/tp/gpu_main.zig index 82fc1edd5b..5eb3de5a88 100644 --- a/zig/src/fabric/tp/gpu_main.zig +++ b/zig/src/fabric/tp/gpu_main.zig @@ -5,24 +5,25 @@ const tp = @import("tp"); const tp_kernel = @import("tp_kernel"); const checks = @import("gpu_checks.zig"); const bench = @import("gpu_bench.zig"); +const cli = @import("gpu_cli.zig"); const max_bytes = (4 << 20) + 4096; pub fn main(init: std.process.Init) !u8 { const args = try init.minimal.args.toSlice(init.arena.allocator()); - if (args.len < 2) { - std.debug.print("usage: tf-tp-test suite | failfast RANK fatal|kill|transport [N [HOLD_S]] | roce-host | hca | kernel\n", .{}); + const command = cli.parse(args[1..]) catch { + std.debug.print("{s}", .{cli.usage}); return 2; - } + }; // `kernel`: whether the build embeds the mailbox kernel (-Dtp-nvcc or -Dtp-fatbin), for scripts to check first - if (std.mem.eql(u8, args[1], "kernel")) { + if (command == .kernel) { std.debug.print("mailbox kernel {d} B\n", .{tp_kernel.mailbox.len}); return if (tp_kernel.mailbox.len > 0) 0 else 1; } const cfg = try tp.Config.fromEnv(); - if (std.mem.eql(u8, args[1], "roce-host")) return @import("roce_host.zig").run(cfg, init.gpa); - if (std.mem.eql(u8, args[1], "hca")) return hca(cfg, init.io); - if (std.mem.eql(u8, args[1], "failfast") and (!cfg.failfast or args.len > 5) and !hostOnlyAllowed(cfg)) { + if (command == .@"roce-host") return @import("roce_host.zig").run(cfg, init.gpa); + if (command == .hca) return hca(cfg, init.io); + if (command == .failfast and (!cfg.failfast or command.failfast.hold != null) and !hostOnlyAllowed(cfg)) { std.debug.print("REFUSED: failfast with TF_TP_FAILFAST=0 or a hold is host-only: TF_TP_ALLOW_HOST_ONLY=1 and no RDMA wire\n", .{}); return 2; } @@ -40,15 +41,14 @@ pub fn main(init: std.process.Init) !u8 { try s.open(&driver, cfg, tp_kernel.mailbox); const up_ms = @as(f64, @floatFromInt(tp.sock.nowNs() - t0)) / 1e6; std.debug.print("INFO rank {d}/{d} on {s} (sm_{d}), backend {t}, mailbox kernel {d} B, bootstrap + NCCL init {d:.0} ms\n", .{ cfg.rank, cfg.world, try ctx.name(&name_buf), try ctx.capability(), s.comm().kind(), tp_kernel.mailbox.len, up_ms }); - if (std.mem.eql(u8, args[1], "suite")) return suite(&s, &driver, init.gpa, cfg); - if (std.mem.eql(u8, args[1], "failfast") and args.len >= 4 and std.mem.eql(u8, args[3], "transport")) { - const n = if (args.len > 4) try std.fmt.parseInt(u32, args[4], 10) else 50; - const hold = if (args.len > 5) try std.fmt.parseInt(u32, args[5], 10) else 0; - return @import("gpu_transport.zig").run(&s, &driver, init.gpa, cfg, try std.fmt.parseInt(u32, args[2], 10), n, hold); - } - if (std.mem.eql(u8, args[1], "failfast") and args.len == 4) return failfast(&s, &driver, init.gpa, cfg, try std.fmt.parseInt(u32, args[2], 10), args[3]); - std.debug.print("unknown command {s}\n", .{args[1]}); - return 2; + return switch (command) { + .suite => suite(&s, &driver, init.gpa, cfg), + .failfast => |f| if (f.mode == .transport) + @import("gpu_transport.zig").run(&s, &driver, init.gpa, cfg, f.victim, f.n, f.hold orelse 0) + else + failfast(&s, &driver, init.gpa, cfg, f.victim, @tagName(f.mode)), + else => unreachable, + }; } /// `hca` (host only, no GPU): both ranks' RoCE pick over the bootstrap; this rank's `NCCL_IB_HCA=...` line on stdout. From 4df28ced9871f0d32ecdb18a399efbc30c939a57 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Schr=C3=B6der?= Date: Sat, 17 Oct 2026 16:30:21 +0200 Subject: [PATCH 41/43] test(tp): retain verbs load failures in HCA diagnostics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit (cherry picked from commit 569577afd3c74c70525a4160e236b7654ac257e7) Reported-by: Christian Schröder Co-Authored-By: Claude Opus 5.5 --- zig/src/fabric/tp/roce_pick.zig | 59 ++++++++++++++++++++++++++++++--- 1 file changed, 54 insertions(+), 5 deletions(-) diff --git a/zig/src/fabric/tp/roce_pick.zig b/zig/src/fabric/tp/roce_pick.zig index b9355c401a..66589ba1ce 100644 --- a/zig/src/fabric/tp/roce_pick.zig +++ b/zig/src/fabric/tp/roce_pick.zig @@ -6,6 +6,7 @@ const bootstrap = @import("bootstrap.zig"); const sock = @import("sock.zig"); const Config = @import("config.zig").Config; const Verbs = @import("verbs").Verbs; +const OpenFn = *const fn (?[:0]const u8) @import("verbs").Error!Verbs; /// Ports a rank lists to its peers (a DGX Spark's ConnectX-7 shows four RDMA devices). pub const max_ports = 8; @@ -90,9 +91,13 @@ pub fn resolve(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSp /// `resolve` with how the devices were chosen, over this host's ports (a collective, as `resolve`). pub fn select(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSpec) !Pick { + return selectWithOpen(cfg, b, out, Verbs.open); +} + +fn selectWithOpen(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSpec, open: OpenFn) !Pick { var ports: [max_ports]rv.Port = undefined; var n_ports: usize = 0; - if (Verbs.open(null)) |opened| { + if (open(null)) |opened| { var v = opened; defer v.close(); n_ports = rv.scanPorts(&v, &ports); @@ -338,14 +343,14 @@ test "the port scan reads a host's devices through verbs: state, Ethernet, the I } /// One rank's `selectFrom` over a real bootstrap (rank 0 on `l`), its hca line or error into `out`. -fn hcaRank(cfg: Config, l: ?sock.Fd, ports: []const rv.Port, out: *[list_bytes + 16]u8, line: *HcaError![]const u8, err: *?anyerror) void { +fn hcaRank(cfg: Config, l: ?sock.Fd, ports: []const rv.Port, open: ?OpenFn, out: *[list_bytes + 16]u8, line: *HcaError![]const u8, err: *?anyerror) void { var b = (if (l) |fd| bootstrap.Bootstrap.acceptOn(fd, 2, 5 * std.time.ns_per_s) else bootstrap.Bootstrap.open(cfg)) catch |e| { err.* = e; return; }; defer b.close(); var specs: [rv.max_hcas]rv.HcaSpec = undefined; - const pick = selectFrom(cfg, &b, ports, &specs) catch |e| { + const pick = (if (open) |f| selectWithOpen(cfg, &b, &specs, f) else selectFrom(cfg, &b, ports, &specs)) catch |e| { err.* = e; return; }; @@ -353,18 +358,62 @@ fn hcaRank(cfg: Config, l: ?sock.Fd, ports: []const rv.Port, out: *[list_bytes + } fn hcaPair(cfgs: [2]Config, ports: [2][]const rv.Port, lines: *[2]HcaError![]const u8, bufs: *[2][list_bytes + 16]u8) !void { + return hcaPairWithOpen(cfgs, ports, .{ null, null }, lines, bufs); +} + +fn hcaPairWithOpen(cfgs: [2]Config, ports: [2][]const rv.Port, open: [2]?OpenFn, lines: *[2]HcaError![]const u8, bufs: *[2][list_bytes + 16]u8) !void { const l = try bootstrap.testListener(); defer sock.close(l.fd); var errs: [2]?anyerror = @splat(null); var c1 = cfgs[1]; c1.link_port = l.port; c1.connect_timeout_ns = 5 * std.time.ns_per_s; - const t = try std.Thread.spawn(.{}, hcaRank, .{ c1, null, ports[1], &bufs[1], &lines[1], &errs[1] }); - hcaRank(cfgs[0], l.fd, ports[0], &bufs[0], &lines[0], &errs[0]); + const t = try std.Thread.spawn(.{}, hcaRank, .{ c1, null, ports[1], open[1], &bufs[1], &lines[1], &errs[1] }); + hcaRank(cfgs[0], l.fd, ports[0], open[0], &bufs[0], &lines[0], &errs[0]); t.join(); for (errs) |e| if (e) |x| return x; } +fn missingLibrary(_: ?[:0]const u8) @import("verbs").Error!Verbs { + return error.LibraryUnavailable; +} + +fn missingSymbol(_: ?[:0]const u8) @import("verbs").Error!Verbs { + return error.MissingSymbol; +} + +test "hca keeps a verbs load failure while completing the peer's port exchange" { + var lines: [2]HcaError![]const u8 = undefined; + var bufs: [2][list_bytes + 16]u8 = undefined; + var scanned: [max_ports]rv.Port = undefined; + const peer = ScanFixture.scan(&hosts[1], &scanned); + const cfgs: [2]Config = .{ .{ .rank = 0 }, .{ .rank = 1 } }; + try hcaPairWithOpen(cfgs, .{ &.{}, peer }, .{ missingLibrary, null }, &lines, &bufs); + try std.testing.expectError(error.LibraryUnavailable, lines[0]); + try std.testing.expectError(error.NoSharedSubnet, lines[1]); +} + +test "hca distinguishes missing verbs symbols from an empty successful scan" { + const cfg: Config = .{ .world = 1 }; + var b = try bootstrap.Bootstrap.open(cfg); + defer b.close(); + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var buf: [list_bytes + 16]u8 = undefined; + const pick = try selectWithOpen(cfg, &b, &specs, missingSymbol); + try std.testing.expectError(error.MissingSymbol, hcaLine(pick, &specs, &buf)); + try std.testing.expectError(error.NoRoceDevice, hcaLine(try selectFrom(cfg, &b, &.{}, &specs), &specs, &buf)); +} + +test "hca keeps an explicit list when the verbs scan cannot open its library" { + const cfg: Config = .{ .world = 1, .roce_hca = "fx_rocep1s0f0:3" }; + var b = try bootstrap.Bootstrap.open(cfg); + defer b.close(); + var specs: [rv.max_hcas]rv.HcaSpec = undefined; + var buf: [list_bytes + 16]u8 = undefined; + const pick = try selectWithOpen(cfg, &b, &specs, missingLibrary); + try std.testing.expectEqualStrings("NCCL_IB_HCA=fx_rocep1s0f0\n", try hcaLine(pick, &specs, &buf)); +} + test "hca over a real bootstrap: both ranks print their peer-linked ports; an explicit rank keeps its list; no ports fail both" { var lines: [2]HcaError![]const u8 = undefined; var bufs: [2][list_bytes + 16]u8 = undefined; From 0a80da001f05f5225cceeaa16190bb25422c786a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Schr=C3=B6der?= Date: Sat, 17 Oct 2026 16:33:30 +0200 Subject: [PATCH 42/43] fix(tp): distinguish verbs load errors from absent HCA ports MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit (cherry picked from commit 8f75a6aa22c9c71d1342bbf0d914a67ae3a5318e) Reported-by: Christian Schröder Co-Authored-By: Claude Opus 5.5 --- zig/src/fabric/tp/gpu_main.zig | 9 ++++++--- zig/src/fabric/tp/roce_pick.zig | 23 ++++++++++++++++++----- 2 files changed, 24 insertions(+), 8 deletions(-) diff --git a/zig/src/fabric/tp/gpu_main.zig b/zig/src/fabric/tp/gpu_main.zig index 5eb3de5a88..1243ab794e 100644 --- a/zig/src/fabric/tp/gpu_main.zig +++ b/zig/src/fabric/tp/gpu_main.zig @@ -73,10 +73,13 @@ fn hca(cfg: tp.Config, io: std.Io) !u8 { var buf: [tp.roce_pick.list_bytes + 16]u8 = undefined; const line = tp.roce_pick.hcaLine(pick, &specs, &buf) catch |e| { const why = switch (e) { - error.NoRoceDevice => "no active Ethernet RDMA port with an IPv4 RoCE v2 GID on this host", - error.NoSharedSubnet => "no RoCE port here shares an IPv4 /24 with the peer's", + error.LibraryUnavailable => "could not load libibverbs.so.1 (LibraryUnavailable); check the RDMA userspace installation", + error.MissingSymbol => "libibverbs.so.1 lacks required symbols (MissingSymbol); check the RDMA userspace installation", + error.NoRoceDevice => "no active Ethernet RDMA port with an IPv4 RoCE v2 GID on this host: name the peer's ports in TF_TP_ROCE_HCA or NCCL_IB_HCA", + error.NoSharedSubnet => "no RoCE port here shares an IPv4 /24 with the peer's: name the peer's ports in TF_TP_ROCE_HCA or NCCL_IB_HCA", + else => @errorName(e), }; - std.debug.print("FAIL hca rank {d}: {s}: name the peer's ports in TF_TP_ROCE_HCA or NCCL_IB_HCA\n", .{ cfg.rank, why }); + std.debug.print("FAIL hca rank {d}: {s}\n", .{ cfg.rank, why }); return 1; }; var names: [tp.roce_pick.list_bytes]u8 = undefined; diff --git a/zig/src/fabric/tp/roce_pick.zig b/zig/src/fabric/tp/roce_pick.zig index 66589ba1ce..d1ed7a6529 100644 --- a/zig/src/fabric/tp/roce_pick.zig +++ b/zig/src/fabric/tp/roce_pick.zig @@ -13,7 +13,12 @@ pub const max_ports = 8; const blob_bytes = 1 + 4 * max_ports; /// How many devices, and how: on a /24 the peer shares (`shared`), named by TF_TP_ROCE_HCA (`explicit`), or neither. -pub const Pick = struct { n: usize, shared: bool, explicit: bool = false }; +pub const Pick = struct { + n: usize, + shared: bool, + explicit: bool = false, + verbs_open_error: ?@import("verbs").Error = null, +}; fn usable(p: rv.Port) bool { return p.active and p.ethernet and p.ipv4 != null; @@ -97,12 +102,17 @@ pub fn select(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSpe fn selectWithOpen(cfg: Config, b: *bootstrap.Bootstrap, out: *[rv.max_hcas]rv.HcaSpec, open: OpenFn) !Pick { var ports: [max_ports]rv.Port = undefined; var n_ports: usize = 0; + var open_error: ?@import("verbs").Error = null; if (open(null)) |opened| { var v = opened; defer v.close(); n_ports = rv.scanPorts(&v, &ports); - } else |_| {} - return selectFrom(cfg, b, ports[0..n_ports], out); + } else |e| { + open_error = e; + } + var pick = try selectFrom(cfg, b, ports[0..n_ports], out); + pick.verbs_open_error = open_error; + return pick; } /// `select` over given ports (the scan's, or a test's fixture ports). @@ -140,11 +150,14 @@ pub const HcaError = error{ NoRoceDevice, /// this host's ports share no IPv4 /24 with the peer's: NCCL would not reach the peer on them either NoSharedSubnet, -}; +} || @import("verbs").Error; /// `tf-tp-test hca`'s line for run_hosts.sh: `NCCL_IB_HCA=` the picked devices; never the first local ports alone. pub fn hcaLine(pick: Pick, specs: []const rv.HcaSpec, buf: *[list_bytes + 16]u8) HcaError![]const u8 { - if (pick.n == 0) return error.NoRoceDevice; + if (pick.n == 0) { + if (pick.verbs_open_error) |e| return e; + return error.NoRoceDevice; + } if (!pick.shared and !pick.explicit) return error.NoSharedSubnet; const prefix = "NCCL_IB_HCA="; @memcpy(buf[0..prefix.len], prefix); From adbdb8ab6ef873b75aea7cacddcf0f9e1f9d98f7 Mon Sep 17 00:00:00 2001 From: Jay Leaton <58495321+jayleaton@users.noreply.github.com> Date: Sun, 18 Oct 2026 18:00:11 +0700 Subject: [PATCH 43/43] CUDA launches: a not-ready launch returns NotReady unlogged, as Driver.check does launched() named the failing kernel for every non-success code, so a launch that main's tests answer with CUDA_ERROR_NOT_READY (qwen_cuda_test, experts3_host_test) reached errorName on a fake driver and crashed. It now returns error.NotReady first, like Driver.check. The ranks test reads the monotonic clock through std.Io, so it also compiles on Windows hosts (zig-windows). Co-Authored-By: Claude Opus 5.5 --- zig/src/cuda/graph_cache_ranks_test.zig | 4 +--- zig/src/cuda/launch.zig | 1 + 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/zig/src/cuda/graph_cache_ranks_test.zig b/zig/src/cuda/graph_cache_ranks_test.zig index 608ebe87a1..18b8a87e9d 100644 --- a/zig/src/cuda/graph_cache_ranks_test.zig +++ b/zig/src/cuda/graph_cache_ranks_test.zig @@ -53,9 +53,7 @@ const Board = struct { } fn now() u64 { - var ts: std.c.timespec = undefined; - _ = std.c.clock_gettime(.MONOTONIC, &ts); - return @as(u64, @intCast(ts.sec)) * std.time.ns_per_s + @as(u64, @intCast(ts.nsec)); + return @intCast(std.Io.Clock.awake.now(std.testing.io).toNanoseconds()); } }; diff --git a/zig/src/cuda/launch.zig b/zig/src/cuda/launch.zig index fe161fabf1..47cbf720b6 100644 --- a/zig/src/cuda/launch.zig +++ b/zig/src/cuda/launch.zig @@ -111,6 +111,7 @@ pub fn launch(f: Function, cfg: Config, stream: Stream, args: *Args) Error!void /// A launch's result; a failure names the kernel and its configuration with the code. fn launched(f: Function, cfg: Config, res: abi.Result, what: []const u8) Error!void { if (res == abi.success) return; + if (res == abi.error_not_ready) return error.NotReady; // as Driver.check: not a failure to log const d = f.d; var buf: [160]u8 = undefined; std.log.err("{s}: {s} ({d}) {s}; {s}{s}", .{ what, d.errorName(res), res, d.errorText(res), describe(&buf, d.functionName(f.handle), cfg), driver.shutdownNote(res) });