glm cuda: rank 0's NCCL id goes over the ranks' lines, so a /32 ring's followers bootstrap on NCCL_SOCKET_IFNAME

zig/src/cuda/tp_watch.zig

@@ -1,14 +1,19 @@
-//! Ranks' liveness lines: a TCP line from each rank to rank 0, open before NCCL, carrying ready, error, wake records.
+//! Ranks' liveness lines: a TCP line from each rank to rank 0, open before NCCL: rank 0's NCCL id, ready, error, wake.
 const std = @import("std");
 const platform = @import("sys");
 const posix = @import("sys").posix;
 
-/// One record on a line: ready, an error and its name, or rank 0's wake.
+/// One record on a line: ready, an error and its name, rank 0's wake, or a part of rank 0's NCCL id.
 const Record = extern struct { kind: u32, rank: u32, text: [56]u8 = @splat(0) };
 const ready_kind: u32 = 1;
 const error_kind: u32 = 2;
 const wake_kind: u32 = 3;
+const id_kind: u32 = 4;
 const max_lines = 16;
+const id_len = 128; // ncclUniqueId's bytes
+const id_part = 48; // an id record's bytes, from text[8]; its part's number is text[0]
+const id_parts = (id_len + id_part - 1) / id_part;
+const id_all: u32 = (1 << id_parts) - 1;
 
 /// Set by a rank that ends on purpose, before it tells the others: a line closing after that is expected.
 var ending = std.atomic.Value(bool).init(false);
@@ -21,6 +26,8 @@ var peer_rank: [max_lines]u32 = undefined;
 var wakes = std.atomic.Value(u32).init(0); // wake records this rank's line carried (another rank's)
 var slept: u32 = 0; // the wakes this rank's sleeps took
 var watch_io: std.Io = undefined; // the watchers' futex wakes
+var id_bytes: [id_len]u8 = undefined; // rank 0's NCCL id as this rank's line carried it
+var id_seen = std.atomic.Value(u32).init(0); // a bit for each part of it taken
 
 /// This rank ends on purpose (rank 0 sends its end header, or a rank leaves on one).
 pub fn ended() void {
@@ -152,6 +159,40 @@ pub fn sleep(io: std.Io) void {
     }
 }
 
+/// Rank 0: its NCCL id to every other rank, before NCCL starts there (as the fabric's bootstrap broadcasts its id).
+pub fn sendId(id: *const [id_len]u8) void {
+    for (0..id_parts) |part| send(idRecord(id, part));
+}
+
+/// A rank other than 0: rank 0's NCCL id from this rank's line, within 15 minutes.
+pub fn takeId(io: std.Io) ![id_len]u8 {
+    var waited: usize = 0;
+    while (id_seen.load(.acquire) != id_all) : (waited += 1) {
+        if (waited >= 90_000) return error.RankNeverSentId;
+        std.Io.sleep(io, .fromMilliseconds(10), .awake) catch {};
+    }
+    return id_bytes;
+}
+
+fn idRecord(id: *const [id_len]u8, part: usize) Record {
+    var r: Record = .{ .kind = id_kind, .rank = @intCast(my_rank) };
+    const at = part * id_part;
+    const n = @min(id_part, id_len - at);
+    r.text[0] = @intCast(part);
+    @memcpy(r.text[8..][0..n], id[at..][0..n]);
+    return r;
+}
+
+/// A part of rank 0's id into id_bytes; a part past the id is not taken.
+fn takePart(r: *const Record) void {
+    const part = r.text[0];
+    if (part >= id_parts) return;
+    const at = @as(usize, part) * id_part;
+    const n = @min(id_part, id_len - at);
+    @memcpy(id_bytes[at..][0..n], r.text[8..][0..n]);
+    _ = id_seen.fetchOr(@as(u32, 1) << @intCast(part), .acq_rel);
+}
+
 /// A rank's transport failed: with lines open every rank hears `name` and this one ends now; with none, nothing.
 pub fn fatal(name: []const u8) void {
     if (line_count == 0) return;
@@ -192,6 +233,10 @@ fn watch(fd: posix.socket_t, line: usize) void {
             watch_io.futexWake(u32, &wakes.raw, 1);
             continue;
         }
+        if (r.kind == id_kind) {
+            takePart(&r);
+            continue;
+        }
         if (r.kind == error_kind and !ending.load(.acquire)) {
             const text = std.mem.sliceTo(&r.text, 0);
             std.log.err("rank {d}: rank {d} stopped ({s}); this rank stops too", .{ my_rank, r.rank, text });
@@ -314,6 +359,70 @@ test "a lost wake and an early failure both stop and join the test watcher" {
     try std.testing.expectEqual(count, line_count);
 }
 
+const IdTest = struct {
+    /// The watcher took every part within 2 s (takeId alone would wait 15 minutes).
+    fn taken() !void {
+        for (0..2000) |_| {
+            if (id_seen.load(.acquire) == id_all) return;
+            std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {};
+        }
+        return error.IdTimeout;
+    }
+
+    fn bytes(seed: u8) [id_len]u8 {
+        var id: [id_len]u8 = undefined;
+        for (&id, 0..) |*b, i| b.* = @as(u8, @truncate(i *% 37)) +% seed;
+        return id;
+    }
+};
+
+test "a rank takes rank 0's NCCL id from its line in any order, not a part past the id, and as sendId sends it" {
+    if (@import("builtin").os.tag == .windows) return error.SkipZigTest;
+    var pair: [2]c_int = undefined;
+    if (ProcessTest.socketpair(posix.AF.UNIX, posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketFailed;
+    defer _ = posix.system.close(pair[0]);
+    var peer_open = true;
+    defer if (peer_open) {
+        _ = posix.system.close(pair[1]);
+    };
+    const was_ending = ending.load(.acquire);
+    const was_seen = id_seen.load(.acquire);
+    const was_count = line_count;
+    defer {
+        ending.store(was_ending, .release);
+        id_seen.store(was_seen, .release);
+        line_count = was_count;
+    }
+    id_seen.store(0, .release);
+    ending.store(false, .release);
+    const watcher = try std.Thread.spawn(.{}, watch, .{ pair[0], 0 });
+    defer {
+        ending.store(true, .release);
+        _ = posix.system.close(pair[1]);
+        peer_open = false;
+        watcher.join();
+    }
+    const id = IdTest.bytes(11);
+    var past = idRecord(&id, 1);
+    past.text[0] = id_parts;
+    @memset(past.text[8..], 0xff);
+    for ([_]Record{ idRecord(&id, 2), past, idRecord(&id, 0), idRecord(&id, 1) }) |r| {
+        const b = std.mem.asBytes(&r);
+        try std.testing.expectEqual(@as(isize, b.len), posix.system.write(pair[1], b.ptr, b.len));
+    }
+    try IdTest.taken();
+    const got = try takeId(std.testing.io);
+    try std.testing.expectEqualSlices(u8, &id, &got);
+    id_seen.store(0, .release);
+    lines[0] = pair[1];
+    line_count = 1;
+    const next = IdTest.bytes(5);
+    sendId(&next);
+    try IdTest.taken();
+    const sent = try takeId(std.testing.io);
+    try std.testing.expectEqualSlices(u8, &next, &sent);
+}
+
 const ProcessTest = struct {
     var last_pid: c_int = 0;
     var last_fd: c_int = -1;

zig/src/families/glm/cuda_native.zig

@@ -104,6 +104,7 @@ const Owned = struct {
 };
 
 extern "c" fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: c_int) c_int;
+extern "c" fn unsetenv(name: [*:0]const u8) c_int;
 
 /// Every peer's send/recv connected now, before the weights fill memory (connected on first use, ibv_reg_mr can fail).
 fn warmPeers(d: *const cuda.Driver, g: *const cuda.tp_group.Group) !void {
@@ -162,10 +163,15 @@ fn openRank(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [
     // NCCL's cuMem host buffers wore out GB10's driver pool under load (NV_ERR_NO_MEMORY; next ranks failed to start)
     if (std.c.getenv("NCCL_CUMEM_HOST_ENABLE") == null) _ = setenv("NCCL_CUMEM_HOST_ENABLE", "0", 0);
     if (std.c.getenv("NCCL_CUMEM_ENABLE") == null) _ = setenv("NCCL_CUMEM_ENABLE", "0", 0);
+    // NCCL_COMM_ID holds NCCL's bootstrap to its subnet past NCCL_SOCKET_IFNAME: a follower naming one lets it pick
+    if (p.rank > 0 and std.c.getenv("NCCL_SOCKET_IFNAME") != null) _ = unsetenv("NCCL_COMM_ID");
     var lib = try cuda.nccl.Library.open();
     errdefer lib.close();
     var id: cuda.nccl.UniqueId = undefined;
-    try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); // NCCL_COMM_ID: the same id on every rank
+    if (p.rank == 0) {
+        try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); // the root at NCCL_COMM_ID
+        cuda.tp_watch.sendId(&id.internal);
+    } else id.internal = try cuda.tp_watch.takeId(io); // rank 0's, over this rank's line
     var comm: cuda.nccl.Comm = null;
     try lib.check(lib.api.ncclCommInitRank(&comm, @intCast(p.world), id, @intCast(p.rank)), "ncclCommInitRank");
     own.group = .{ .lib = lib, .comm = comm, .rank = p.rank, .world = p.world };