From 5282572134fc777d38860cea31623239f7706b5c 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 1/2] test(tp): reject malformed runner commands before configuration (cherry picked from commit 854b57cf74b90d306626f35e40de61c047574f72) Reported-by: Bizuayeu (cherry picked from commit 88418d8d0f4ecb875133b4da2a4de51faf3d4594) --- 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 6419f55b70..46465c5cf0 100644 --- a/zig/build/tp.zig +++ b/zig/build/tp.zig @@ -42,6 +42,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); 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)) @@ -68,6 +72,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 69b1e5142c24e8a1a9a3c0847b30c81a93287138 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 2/2] fix(tp): parse runner commands before opening sessions (cherry picked from commit 602a0fd4e1d621cf0715d01a2b466497dc3b4bb7) Reported-by: Bizuayeu (cherry picked from commit 4c69b0d4b434b782a785130ae649f4337f931f91) --- 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 46465c5cf0..1c3207f64e 100644 --- a/zig/build/tp.zig +++ b/zig/build/tp.zig @@ -51,7 +51,9 @@ pub fn add(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti 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); @@ -83,6 +85,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.