tp: parse the runner's command before taking any resource (carried from #636)

zig/build/tp.zig

@@ -42,12 +42,18 @@ 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);
     const portable = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/tp/portable_test.zig"), .target = host, .optimize = .debug, .link_libc = true, .imports = &.{.{ .name = "cuda", .module = runtime(b, host, .debug) }} });
     const tests = b.addRunArtifact(b.addTest(.{ .root_module = if (host.result.os.tag == .linux)
         module(b, host, .debug, runtime(b, host, .debug))
     else
         portable }));
-    b.step("test-tp", "TP bootstrap, failure, settings and in-process transport tests").dependOn(&tests.step);
+    const tp_step = b.step("test-tp", "TP bootstrap, failure, settings and in-process transport tests");
+    tp_step.dependOn(&tests.step);
+    tp_step.dependOn(cli_step);
     test_step.dependOn(&tests.step);
     const linux = b.resolveTargetQuery(.{ .cpu_arch = host.result.cpu.arch, .os_tag = .linux, .abi = .gnu });
     const check_cuda = runtime(b, linux, .debug);
@@ -68,6 +74,27 @@ 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);
+        }
+        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);
 }

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);
+}

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.