flashnext: a shard in a folder inside the model folder (#1220)

zig/build/cuda.zig

@@ -418,7 +418,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st
     const h3_checkpoint = b.createModule(.{ .root_source_file = b.path("zig/src/families/h3/checkpoint.zig"), .target = host, .optimize = .debug, .link_libc = true });
     // the checkpoint maps its file with mmap: not Windows
     if (host.result.os.tag != .windows) step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = h3_checkpoint })).step);
-
+    @import("pull_tests.zig").add(b, host, .debug, native_modules.engines, mods.core, step);
     b.step("test-save-reopen-host", "Unresolved save through actual checkpoint admission").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = mods.core, .filters = &.{"Save recovery:"} })).step);
     b.step("test-trainer-close-host", "Actual native close and real Flash Next trainer ownership sentinels").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = native, .filters = &.{"Save recovery:"} })).step);
     // the HTTP server's unit tests (zig/src/server/root.zig), which the macOS build runs, on Linux hosts too
@@ -479,7 +479,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st
     }
     // the server goldens and Cortheon's fuzz step here when the target is not macOS, whose Metal build defines them
     if (server_steps) {
-        golden_build.step(b, host, server_tests, native_modules.api, mods.tokenizer, template);
+        golden_build.step(b, host, server_tests, native_modules.api, mods.tokenizer, template, native_modules.engines);
         @import("cortheon_fuzz.zig").add(b, host, native_modules.api, mods.tokenizer, template);
     }
     // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them

zig/build/golden.zig

@@ -2,8 +2,8 @@
 
 const std = @import("std");
 
-/// `zig build test-golden` over the server tests' module and the imports it was built with.
-pub fn step(b: *std.Build, target: std.Build.ResolvedTarget, server: *std.Build.Module, api: *std.Build.Module, tokenizer: *std.Build.Module, template: *std.Build.Module) void {
+/// `zig build test-golden` over the server tests' module; `engines` is the platform registry `cli/hub.zig` reads.
+pub fn step(b: *std.Build, target: std.Build.ResolvedTarget, server: *std.Build.Module, api: *std.Build.Module, tokenizer: *std.Build.Module, template: *std.Build.Module, engines: *std.Build.Module) void {
     const fake_serve = b.addExecutable(.{ .name = "fake_serve", .root_module = b.createModule(.{
         .root_source_file = b.path("zig/tests/server/fake_serve.zig"),
         .target = target,
@@ -45,6 +45,7 @@ pub fn step(b: *std.Build, target: std.Build.ResolvedTarget, server: *std.Build.
             .target = target,
             .optimize = .Debug,
             .link_libc = true,
+            .imports = &.{.{ .name = "native_engines", .module = engines }},
         }), .filters = &.{"security:"} });
         const run = b.addRunArtifact(checks);
         run.setCwd(b.path("."));

zig/build/metal_targets.zig

@@ -113,7 +113,7 @@ pub fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s
     // the engine seam's own tests (lane_host.zig), as zig/tests/server/zig_test.sh runs them
     test_step.?.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = api })).step);
     // the frozen server parity goldens, checked by Zig driving fake_serve
-    golden_build.step(b, target, server_tests, api, tokenizer, template);
+    golden_build.step(b, target, server_tests, api, tokenizer, template, engines);
     b.step("test-native", "The Metal engines module's host-side tests").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = engines })).step);
     const reuse = b.addExecutable(.{ .name = "tf-flashnext-reuse", .root_module = b.createModule(.{
         .root_source_file = b.path("zig/tests/flashnext_reuse.zig"),

zig/build/pull_tests.zig

@@ -0,0 +1,16 @@
+//! `pull`'s own tests: `cli/hub.zig` reads the engine registry, so they need `native_engines` and EXLR8's import.
+const std = @import("std");
+
+pub fn add(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, engines: *std.Build.Module, core: *std.Build.Module, test_step: *std.Build.Step) void {
+    if (target.result.os.tag == .windows) return; // the fake hub listens on a POSIX socket
+    const mod = b.createModule(.{
+        .root_source_file = b.path("zig/src/cli/pull_test.zig"),
+        .target = target,
+        .optimize = optimize,
+        .link_libc = true,
+        .imports = &.{ .{ .name = "native_engines", .module = engines }, @import("exlr8.zig").cliImport(b, target, optimize, core) },
+    });
+    const run = b.addRunArtifact(b.addTest(.{ .root_module = mod }));
+    b.step("test-pull-host", "Checkpoint download paths: a nested shard folder, retries, resume and the whole-file fallback").dependOn(&run.step);
+    test_step.dependOn(&run.step);
+}

zig/src/cli/hub.zig

@@ -197,7 +197,8 @@ test "security: negative quantization metadata is refused without a cast trap" {
 test "family detection matches the registry and refuses unknown types" {
     try std.testing.expectEqualStrings("Nemotron 3.5 Lightning", family("nemotron_h").?.title);
     try std.testing.expectEqualStrings("Qwen3.8 Flash Next", family("qwen4_exp").?.title);
-    try std.testing.expectEqualStrings("Gemma 4 12B", family("gemma4_unified").?.title);
+    // Every family the linked registry serves resolves; the titles are display names, so only their presence is fixed.
+    for (engines.families) |f| try std.testing.expect(family(f.model_type) != null);
     try std.testing.expect(family("gemma3n") == null);
     try std.testing.expect(family("unknown") == null);
 }
@@ -236,3 +237,17 @@ test "cache layout helpers read a synthetic cache" {
     const missing = try cachedSnapshot(a, io, hub, "Org/Other");
     try std.testing.expect(missing == null);
 }
+
+test "a checkpoint's size counts shards in a folder inside it" {
+    const io = std.testing.io;
+    const gpa = std.testing.allocator;
+    var tmp = std.testing.tmpDir(.{});
+    defer tmp.cleanup();
+    // The publisher's export keeps every shard in `tensors/`, so a walk that stopped at the top level undercounted it.
+    try tmp.dir.createDirPath(io, "tensors");
+    try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = "{}" });
+    try tmp.dir.writeFile(io, .{ .sub_path = "tensors/model-00001-of-00001.safetensors", .data = "xxxxxxx" });
+    var buf: [64]u8 = undefined;
+    const dir = try std.fmt.bufPrint(&buf, ".zig-cache/tmp/{s}", .{tmp.sub_path});
+    try std.testing.expectEqual(@as(u64, 9), try sizeOf(gpa, io, dir));
+}

zig/src/cli/pull_test.zig

@@ -16,16 +16,22 @@ const FakeHub = struct {
     range_starts: std.atomic.Value(u64) = .init(0),
     weights: []const u8,
     config: []const u8,
+    nested: []const u8,
     weights_sha_hex: []const u8,
+    nested_sha_hex: []const u8,
     config_git_hex: [40]u8,
     flaky_git_hex: [40]u8,
     port: u16,
 
     const weights_path = "/Org/Flash/resolve/aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1/weights.safetensors";
     const config_path = "/Org/Flash/resolve/aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1/config.json";
+    const nested_path = "/Org/Flash/resolve/aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1/tensors/model-00001-of-00001.safetensors";
     const flaky_body = "{\"chat_template\": \"{{ messages }}\"}";
     const drafter_config = "{\"model_type\": \"qwen3\", \"architectures\": [\"DFlash2DraftModel\"], \"dflash_config\": {\"block_size\": 8}}";
 
+    var nested_body_storage: [1234]u8 = @splat(0x5A);
+    const nested_body: []const u8 = &nested_body_storage;
+
     fn open() !FakeHub {
         const posix = std.posix;
         const rc = posix.system.socket(posix.AF.INET, posix.SOCK.STREAM, 0);
@@ -45,11 +51,17 @@ const FakeHub = struct {
         hash.update(weights_body);
         var digest: [32]u8 = undefined;
         hash.final(&digest);
+        var nested_hash = std.crypto.hash.sha2.Sha256.init(.{});
+        nested_hash.update(nested_body);
+        var nested_digest: [32]u8 = undefined;
+        nested_hash.final(&nested_digest);
         return .{
             .server_fd = fd,
             .weights = weights_body,
             .config = config_body,
+            .nested = nested_body,
             .weights_sha_hex = try std.fmt.allocPrint(std.testing.allocator, "{s}", .{std.fmt.bytesToHex(digest, .lower)}),
+            .nested_sha_hex = try std.fmt.allocPrint(std.testing.allocator, "{s}", .{std.fmt.bytesToHex(nested_digest, .lower)}),
             .config_git_hex = gitHex(config_body),
             .flaky_git_hex = gitHex(flaky_body),
             .port = port,
@@ -162,7 +174,7 @@ const FakeHub = struct {
             respond(fd, drafter_config);
         } else if (std.mem.eql(u8, path, "/api/models/Org/Flash/tree/aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1") or std.mem.eql(u8, path, "/api/models/Org/Whole/tree/bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb2")) {
             // The hub's tree sends LFS oids as bare hex.
-            const tree = std.fmt.allocPrint(std.testing.allocator, "[{{\"type\": \"file\", \"oid\": \"{s}\", \"path\": \"config.json\", \"size\": {d}}}, {{\"type\": \"file\", \"oid\": \"1111111111111111111111111111111111111111\", \"path\": \"weights.safetensors\", \"size\": {d}, \"lfs\": {{\"oid\": \"{s}\", \"size\": {d}}}}}, {{\"type\": \"file\", \"oid\": \"{s}\", \"path\": \"flaky.json\", \"size\": {d}}}]", .{ &fake.config_git_hex, fake.config.len, fake.weights.len, fake.weights_sha_hex, fake.weights.len, &fake.flaky_git_hex, flaky_body.len }) catch return false;
+            const tree = std.fmt.allocPrint(std.testing.allocator, "[{{\"type\": \"file\", \"oid\": \"{s}\", \"path\": \"config.json\", \"size\": {d}}}, {{\"type\": \"file\", \"oid\": \"1111111111111111111111111111111111111111\", \"path\": \"weights.safetensors\", \"size\": {d}, \"lfs\": {{\"oid\": \"{s}\", \"size\": {d}}}}}, {{\"type\": \"file\", \"oid\": \"{s}\", \"path\": \"flaky.json\", \"size\": {d}}}, {{\"type\": \"file\", \"oid\": \"3333333333333333333333333333333333333333\", \"path\": \"tensors/model-00001-of-00001.safetensors\", \"size\": {d}, \"lfs\": {{\"oid\": \"{s}\", \"size\": {d}}}}}]", .{ &fake.config_git_hex, fake.config.len, fake.weights.len, fake.weights_sha_hex, fake.weights.len, &fake.flaky_git_hex, flaky_body.len, fake.nested.len, fake.nested_sha_hex, fake.nested.len }) catch return false;
             defer std.testing.allocator.free(tree);
             respond(fd, tree);
         } else if (std.mem.eql(u8, path, config_path) or std.mem.eql(u8, path, "/Org/Whole/resolve/bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb2/config.json")) {
@@ -179,6 +191,8 @@ const FakeHub = struct {
         } else if (std.mem.eql(u8, path, "/Org/Whole/resolve/bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb2/weights.safetensors")) {
             _ = fake.weight_requests.fetchAdd(1, .monotonic);
             whole(fd, fake.weights);
+        } else if (std.mem.eql(u8, path, nested_path) or std.mem.eql(u8, path, "/Org/Whole/resolve/bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb2/tensors/model-00001-of-00001.safetensors")) {
+            whole(fd, fake.nested);
         } else if (std.mem.eql(u8, path, weights_path)) {
             _ = fake.weight_requests.fetchAdd(1, .monotonic);
             const from = range_start orelse {
@@ -238,6 +252,7 @@ const FakeHub = struct {
         if (fake.thread) |t| t.join();
         _ = std.posix.system.close(fake.server_fd);
         std.testing.allocator.free(fake.weights_sha_hex);
+        std.testing.allocator.free(fake.nested_sha_hex);
     }
 };
 
@@ -271,6 +286,11 @@ test "pull downloads, verifies, retries a cut body, resumes and refuses a family
     try std.testing.expectEqualStrings(FakeHub.config_body, config_read);
     const weights_read = try std.Io.Dir.cwd().readFileAlloc(io, weights_link, a, .limited(1 << 20));
     try std.testing.expectEqualSlices(u8, FakeHub.weights_body, weights_read);
+    // A shard in a folder inside the model folder lands as a real subfolder in the snapshot, still a blob link.
+    const nested_link = try std.fs.path.join(a, &.{ snapshot, "tensors/model-00001-of-00001.safetensors" });
+    const nested_read = try std.Io.Dir.cwd().readFileAlloc(io, nested_link, a, .limited(1 << 20));
+    try std.testing.expectEqualSlices(u8, FakeHub.nested_body, nested_read);
+    try std.Io.Dir.cwd().access(io, try std.fs.path.join(a, &.{ root, "models--Org--Flash/blobs", fake_hub.nested_sha_hex }), .{});
     const ref = try std.Io.Dir.cwd().readFileAlloc(io, try std.fs.path.join(a, &.{ root, "models--Org--Flash/refs/main" }), a, .limited(64));
     try std.testing.expectEqualStrings("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1", std.mem.trim(u8, ref, " \r\n"));
     try std.testing.expect(std.mem.indexOf(u8, out.written(), "weights.safetensors") != null);
@@ -288,6 +308,7 @@ test "pull downloads, verifies, retries a cut body, resumes and refuses a family
     try std.testing.expectEqual(@as(u8, 0), try pull.run(a, io, &out2.writer, &err_out.writer, &env, root, "Org/Flash"));
     try std.testing.expectEqual(@as(u32, 1), fake_hub.weight_requests.load(.monotonic));
     try std.testing.expect(std.mem.indexOf(u8, out2.written(), "weights.safetensors: cached") != null);
+    try std.testing.expect(std.mem.indexOf(u8, out2.written(), "tensors/model-00001-of-00001.safetensors: cached") != null);
     try std.testing.expect(std.mem.indexOf(u8, out2.written(), "flaky.json: cached") != null);
     try std.testing.expectEqual(@as(u32, 2), fake_hub.flaky_requests.load(.monotonic));
 
@@ -448,3 +469,25 @@ test "security: a hub response cannot write beyond the manifest file size" {
     };
     try std.testing.expect(written <= 8);
 }
+
+test "security: a shard in a folder inside the model folder lands inside the snapshot" {
+    var arena_state = std.heap.ArenaAllocator.init(std.testing.allocator);
+    defer arena_state.deinit();
+    const a = arena_state.allocator();
+    const io = std.testing.io;
+    var fake_hub = FakeHub.open() catch return error.SkipZigTest;
+    fake_hub.start();
+    defer fake_hub.stopServer();
+    const root = try std.fmt.allocPrint(a, ".tf-pull-nested-{d}", .{std.Io.Clock.awake.now(io).toNanoseconds()});
+    defer std.Io.Dir.cwd().deleteTree(io, root) catch {};
+    const env = try testEnv(a, try fake_hub.endpoint(a));
+    var out: std.Io.Writer.Allocating = .init(a);
+    var err_out: std.Io.Writer.Allocating = .init(a);
+    try std.testing.expectEqual(@as(u8, 0), try pull.run(a, io, &out.writer, &err_out.writer, &env, root, "Org/Flash"));
+    const snapshot = try std.fs.path.join(a, &.{ root, "models--Org--Flash/snapshots/aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1" });
+    // NVIDIA's own export puts every shard in `tensors/`; the snapshot keeps that folder and still links the blob.
+    const link = try std.fs.path.join(a, &.{ snapshot, "tensors/model-00001-of-00001.safetensors" });
+    const body = try std.Io.Dir.cwd().readFileAlloc(io, link, a, .limited(1 << 20));
+    try std.testing.expectEqualSlices(u8, FakeHub.nested_body, body);
+    try std.Io.Dir.cwd().access(io, try std.fs.path.join(a, &.{ root, "models--Org--Flash/blobs", fake_hub.nested_sha_hex }), .{});
+}

zig/src/core/checkpoint.zig

@@ -345,3 +345,26 @@ test "Save recovery: an unresolved partial save refuses actual checkpoint reopen
         return error.UncertainCheckpointReopened;
     } else |err| try std.testing.expectEqual(error.CheckpointSaveUncertain, err);
 }
+
+test "a checkpoint whose index names shards in a folder opens and reads them" {
+    const io = std.testing.io;
+    const gpa = std.testing.allocator;
+    var tmp = std.testing.tmpDir(.{});
+    defer tmp.cleanup();
+    // NVIDIA's own export writes `tensors/model-00001-of-00001.safetensors`.
+    try tmp.dir.createDirPath(io, "tensors");
+    const image = try testImage(gpa, "a", 4);
+    defer gpa.free(image);
+    try tmp.dir.writeFile(io, .{ .sub_path = "tensors/model-00001-of-00001.safetensors", .data = image });
+    try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "{\"weight_map\":{\"a\":\"tensors/model-00001-of-00001.safetensors\"}}" });
+    var buf: [64]u8 = undefined;
+    const dir = try std.fmt.bufPrint(&buf, ".zig-cache/tmp/{s}", .{tmp.sub_path});
+    const listed = try shardFiles(gpa, io, dir);
+    defer freeShardFiles(gpa, listed);
+    try std.testing.expectEqual(@as(usize, 1), listed.len);
+    try std.testing.expectEqualStrings("tensors/model-00001-of-00001.safetensors", listed[0][dir.len + 1 ..]);
+    var ck = try Checkpoint.openModel(gpa, io, dir);
+    defer ck.close();
+    _ = try ck.get("a");
+    try std.testing.expectEqual(@as(usize, 0), ck.unused());
+}