dsv41: a rank whose peer dies ends with exit 75 (peer link watch)

zig/src/families/dsv41/dsv41.zig

@@ -16,6 +16,7 @@ pub const weights = @import("weights.zig");
 pub const engram_io = @import("engram_io.zig");
 pub const engram_aio = @import("engram_aio.zig");
 pub const link = @import("link.zig");
+pub const peer_watch = @import("peer_watch.zig");
 pub const comm = @import("comm.zig");
 pub const tri = @import("tri.zig");
 pub const tri_basic = @import("tri_basic.zig");
@@ -64,6 +65,7 @@ test {
     refAll(engram_io);
     refAll(engram_aio);
     refAll(link);
+    refAll(peer_watch);
     refAll(comm);
     refAll(tri);
     refAll(tri_basic);

zig/src/families/dsv41/lanes.zig

@@ -20,6 +20,8 @@ const round = @import("round.zig");
 const draft = @import("draft.zig");
 const sampling = @import("sampling.zig");
 const link = @import("link.zig");
+const peer_watch = @import("peer_watch.zig");
+const rdma = @import("rdma.zig");
 const kept_order = @import("kept_order.zig");
 const cuda = @import("cuda");
 
@@ -295,6 +297,8 @@ pub const Lanes = struct {
     gpa: std.mem.Allocator,
     m: *Model,
     peer: ?link.Link, // rank 1 (null: one rank); every follower gets each frame (Model.followers)
+    watch: ?*peer_watch.Watch = null, // the links to the followers, watched for a rank that dies (peer_watch.zig)
+    stop_requested: bool = false, // set by the server's stop (native.zig deinitFn): deinit then ends the watch orderly
     vision: ?*@import("vision.zig").Vision = null, // rank 0's tower (a vision checkpoint served with its kit's vision/)
     step: u64 = 0,
     wire: Wire = .{},
@@ -315,23 +319,33 @@ pub const Lanes = struct {
     snaps: std.ArrayList(*Snap) = .empty, // kept prompt states, oldest first
     next_snap: u64 = 1,
 
-    pub fn init(gpa: std.mem.Allocator, m: *Model) Lanes {
+    pub fn init(gpa: std.mem.Allocator, m: *Model) !Lanes {
         var ln: Lanes = .{ .gpa = gpa, .m = m, .peer = if (m.world > 1) m.peer() else null };
+        if (m.world > 1) ln.watch = try watchLinks(gpa, m, m.followers());
         // the round costs by rows: ROUND_MS[min(rows, len) - 1] (the served table, or TF_DS_ROUND_MS)
         const t = m.round_ms;
         for (&ln.wc, 0..) |*c, i| c.* = .{ .width = @intCast(i + 1), .ms = t[@min(i, t.len - 1)] };
         for (&ln.sc, 0..) |*c, i| c.* = .{ .width = @intCast(i + 1), .ms = t[@min(i, t.len - 1)] };
         return ln;
     }
 
+    /// Ends the peer link watch (Watch.end) before the model frees its links and rings; may exit 75 on `failed`.
+    pub fn endWatch(self: *Lanes, how: peer_watch.How) void {
+        if (self.watch) |w| _ = w.end(self.gpa, how);
+        self.watch = null;
+    }
+
     /// Tells rank 1 to stop, then frees the host state (the model is the caller's).
     pub fn deinit(self: *Lanes) void {
         for (self.snaps.items) |sn| {
             sn.buf.free();
             self.gpa.destroy(sn);
         }
         self.snaps.deinit(self.gpa);
-        if (self.peer != null) self.send(.done) catch {};
+        // a SIGTERM stop is orderly; else a follower gone before the done frame, or a done frame that fails, is a fault
+        const lost = if (self.watch) |w| w.faulted() else false;
+        const sent = if (self.peer != null) (if (self.send(.done)) |_| true else |_| false) else true;
+        self.endWatch(if (self.stop_requested or (!lost and sent)) .done else .failed);
         self.streams.deinit(self.gpa);
         for (&self.seqs) |*q| q.deinit(self.gpa);
         self.wire.deinit(self.gpa);
@@ -961,6 +975,25 @@ fn refused(err: anyerror) void {
 /// A follower's verify times with Model.prof (a profile on ranks 1-3: rounds, issue and wall ns summed).
 pub var follow_prof: struct { rounds: u64 = 0, enqueue: u64 = 0, forward: u64 = 0 } = .{};
 
+/// A watch on `links`; if it cannot start, the error with TF_DS_EXIT_ON_FAULT=1 (exit 75 needs it), else null.
+fn watchLinks(gpa: std.mem.Allocator, m: *Model, links: []const ?link.Link) !?*peer_watch.Watch {
+    var fds: [3]std.posix.fd_t = undefined;
+    var n: usize = 0;
+    for (links) |l| if (l) |x| if (n < fds.len) {
+        fds[n] = x.fd;
+        n += 1;
+    };
+    if (n == 0) return null;
+    // every ring a GPU wait of this rank reads (the same list that Model.synchronize checks)
+    var rings: [6]*rdma.Ring = undefined;
+    const grace = peer_watch.faultGrace();
+    return peer_watch.Watch.start(gpa, fds[0..n], m.rdmaRings(&rings), m.rank, grace) catch |err| {
+        std.debug.print("[tensorfold] the peer link watch did not start ({s})\n", .{@errorName(err)});
+        if (grace != null) return err;
+        return null;
+    };
+}
+
 /// Rank 1 (each of ranks 1-3 on the four-node split): the primitives rank 0 sends, run in order until it says done.
 pub fn follow(gpa: std.mem.Allocator, m: *Model) !void {
     const peer = m.peer();
@@ -981,6 +1014,11 @@ pub fn follow(gpa: std.mem.Allocator, m: *Model) !void {
         while (ki.next()) |b| b.free();
         kept.deinit(gpa);
     }
+    // a dead rank 0 ends this rank too (peer_watch.zig); started after the kept states, so it ends before they free
+    const watch = try watchLinks(gpa, m, &.{peer});
+    errdefer {
+        if (watch) |w| _ = w.end(gpa, .failed);
+    }
     while (true) {
         // the next primitive, polled for 20 ms before the read blocks (rank 0 sends it a fraction of a millisecond
         // after this rank's GPU goes idle in a decode; a sleeping thread would add its core's wake-up to every frame)
@@ -1118,7 +1156,10 @@ pub fn follow(gpa: std.mem.Allocator, m: *Model) !void {
                     b.free();
                 }
             },
-            .done => return,
+            .done => {
+                if (watch) |w| _ = w.end(gpa, .done);
+                return;
+            },
         }
     }
 }

zig/src/families/dsv41/lanes_main.zig

@@ -185,7 +185,7 @@ pub fn main(init: std.process.Init) !u8 {
     const text = try std.Io.Dir.cwd().readFileAlloc(io, requests_file orelse return error.NoRequests, a, .limited(1 << 30));
     const Requests = struct { requests: []Request };
     const reqs = (try std.json.parseFromSliceLeaky(Requests, a, text, .{ .ignore_unknown_fields = true })).requests;
-    var ln = dsv41.lanes.Lanes.init(gpa, m);
+    var ln = try dsv41.lanes.Lanes.init(gpa, m);
     defer ln.deinit();
     ln.fix_k = fix_k;
     if (profile) {

zig/src/families/dsv41/link.zig

@@ -3,6 +3,11 @@
 const std = @import("std");
 const c = std.c;
 
+/// TCP_USER_TIMEOUT (Linux): the longest time sent data stays not acknowledged before the kernel ends the link.
+const tcp_user_timeout: u32 = 18;
+/// The same bound as the keepalive (10 s + 3 x 3 s, rounded up); also the end of an idle link's unanswered probes.
+const user_timeout_ms: u32 = 20_000;
+
 pub const Link = struct {
     fd: c.fd_t,
 
@@ -21,6 +26,12 @@ pub const Link = struct {
         const busy_us: c_int = 50;
         _ = c.setsockopt(fd, c.IPPROTO.TCP, std.posix.TCP.NODELAY, std.mem.asBytes(&one), @sizeOf(c_int));
         _ = c.setsockopt(fd, c.SOL.SOCKET, c.SO.BUSY_POLL, std.mem.asBytes(&busy_us), @sizeOf(c_int));
+        // a vanished node (no FIN) shows in about 20 s to a read and to peer_watch.zig: keepalive, TCP_USER_TIMEOUT
+        _ = c.setsockopt(fd, c.SOL.SOCKET, c.SO.KEEPALIVE, std.mem.asBytes(&one), @sizeOf(c_int));
+        for ([_][2]u32{ .{ std.posix.TCP.KEEPIDLE, 10 }, .{ std.posix.TCP.KEEPINTVL, 3 }, .{ std.posix.TCP.KEEPCNT, 3 }, .{ tcp_user_timeout, user_timeout_ms } }) |o| {
+            const v: c_int = @intCast(o[1]);
+            _ = c.setsockopt(fd, c.IPPROTO.TCP, o[0], std.mem.asBytes(&v), @sizeOf(c_int));
+        }
     }
 
     fn addr(ip: [4]u8, port: u16) c.sockaddr.in {
@@ -64,7 +75,8 @@ pub const Link = struct {
     fn writeAll(l: Link, bytes: []const u8) !void {
         var done: usize = 0;
         while (done < bytes.len) {
-            const n = c.send(l.fd, bytes.ptr + done, bytes.len - done, 0);
+            // NOSIGNAL: a send to a gone peer fails with PeerDown (a SIGPIPE would end the rank with no exit 75)
+            const n = c.send(l.fd, bytes.ptr + done, bytes.len - done, c.MSG.NOSIGNAL);
             if (n <= 0) {
                 if (n < 0 and c.errno(n) == .INTR) continue;
                 return error.PeerDown;
@@ -155,6 +167,38 @@ test "addresses parse as four bytes" {
     try std.testing.expectError(error.BadAddress, parseIp("10.200.30"));
 }
 
+fn tcpOption(fd: c.fd_t, level: i32, name: u32) !c_int {
+    var v: c_int = 0;
+    var len: c.socklen_t = @sizeOf(c_int);
+    if (c.getsockopt(fd, level, name, &v, &len) != 0) return error.GetSockOptFailed;
+    return v;
+}
+
+test "both ends of a link find a gone peer within 20 s: keepalive and TCP_USER_TIMEOUT" {
+    const io = std.testing.io;
+    const port: u16 = 28400 + @as(u16, @intCast(@mod(std.Io.Timestamp.now(io, .awake).nanoseconds, 500)));
+    const Peer = struct {
+        fn run(p: u16, out: *[2]c_int) void {
+            var l = Link.connect(std.testing.io, .{ 127, 0, 0, 1 }, p, 10) catch return;
+            defer l.close();
+            out[0] = tcpOption(l.fd, c.IPPROTO.TCP, tcp_user_timeout) catch -1;
+            out[1] = tcpOption(l.fd, c.SOL.SOCKET, c.SO.KEEPALIVE) catch -1;
+            var buf: [8]u8 = undefined;
+            _ = l.recv(&buf) catch {}; // until rank 0 closes its end
+        }
+    };
+    var follower: [2]c_int = .{ 0, 0 };
+    const t = try std.Thread.spawn(.{}, Peer.run, .{ port, &follower });
+    var l = try Link.listen(.{ 127, 0, 0, 1 }, port);
+    // rank 0's end (the accepted socket) and rank 1's end (connect)
+    try std.testing.expectEqual(@as(c_int, user_timeout_ms), try tcpOption(l.fd, c.IPPROTO.TCP, tcp_user_timeout));
+    try std.testing.expectEqual(@as(c_int, 1), try tcpOption(l.fd, c.SOL.SOCKET, c.SO.KEEPALIVE));
+    try std.testing.expectEqual(@as(c_int, 10), try tcpOption(l.fd, c.IPPROTO.TCP, std.posix.TCP.KEEPIDLE));
+    l.close();
+    t.join();
+    try std.testing.expectEqual([2]c_int{ user_timeout_ms, 1 }, follower);
+}
+
 test "frames cross a loopback link in order" {
     const io = std.testing.io;
     const port: u16 = 29000 + @as(u16, @intCast(@mod(std.Io.Timestamp.now(io, .awake).nanoseconds, 500)));

zig/src/families/dsv41/model.zig

@@ -448,6 +448,31 @@ pub const Model = struct {
     fn synchronize(m: *Model) !void {
         try m.stream.synchronize();
         if (m.neighbor_gather) |*ng| try ng.check();
+        try m.ringsOk();
+    }
+
+    /// After a stream sync: a ring the peer watch aborted during the work may leave stale peer slots, so it fails.
+    fn ringsOk(m: *Model) !void {
+        var buf: [6]*rdma.Ring = undefined;
+        if (rdma.anyFailed(m.rdmaRings(&buf))) return error.RdmaTransportFailed;
+    }
+
+    /// Every RDMA ring a GPU wait of this rank reads (TP2, ring2d.zig, neighbor.zig): the peer watch aborts these.
+    pub fn rdmaRings(m: *Model, buf: *[6]*rdma.Ring) []*rdma.Ring {
+        var k: usize = 0;
+        if (m.ring) |r| {
+            buf[k] = r;
+            k += 1;
+        }
+        if (m.rdma_rings) |rr| for ([_]*rdma.Ring{ rr.all4, rr.col, rr.pair }) |r| {
+            buf[k] = r;
+            k += 1;
+        };
+        if (m.neighbor_gather) |ng| for ([_]*rdma.Ring{ ng.first, ng.second }) |r| {
+            buf[k] = r;
+            k += 1;
+        };
+        return buf[0..k];
     }
 
     /// TP2's RDMA ring for the decode gathers: created on both ranks, the queue-pair infos exchanged over NCCL,
@@ -831,5 +856,7 @@ pub const Model = struct {
         const pool = &m.dpool.?;
         for (q0, slots) |q, s| pool.absorbed[@intCast(s)] = @intCast(q);
         try dr.pass(&m.eng, &m.ch, pool, tokens, q0, slots, steps, null);
+        // the drafter syncs its stream itself (draft.zig), so its drafts get the same ring check as synchronize
+        try m.ringsOk();
     }
 };

zig/src/families/dsv41/native.zig

@@ -159,7 +159,8 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [
         .streams = if (getenv("TF_DS_STREAMS")) |v| std.fmt.parseInt(usize, v, 10) catch return error.BadStreams else model.max_streams_default,
     });
     errdefer st.m.close();
-    st.lanes = lanes_mod.Lanes.init(gpa, st.m);
+    st.lanes = try lanes_mod.Lanes.init(gpa, st.m);
+    errdefer st.lanes.endWatch(.failed); // a failed start closes the model next: the watch ends first
     if (getenv("TF_DS_FIX_K")) |t| st.lanes.fix_k = std.fmt.parseInt(usize, t, 10) catch return error.BadFixK;
     if (getenv("TF_DS_SERVED_K")) |t| st.lanes.served_k = if (std.mem.eql(u8, t, "1") or std.mem.eql(u8, t, "all")) .all else if (std.mem.eql(u8, t, "solo")) .solo else .off;
     // a vision checkpoint's tower on this rank, from the kit's vision/ (torch's attention cubin, the routing bias)
@@ -175,6 +176,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [
 
 fn deinitFn(ptr: *anyopaque) void {
     const st: *State = @ptrCast(@alignCast(ptr));
+    st.lanes.stop_requested = true; // the server stops (its SIGTERM, or its start failed after open): an orderly end
     st.lanes.deinit(); // rank 1 hears done
     if (st.vision) |v| v.close();
     st.m.close();