fix: thread data materialization schema through gui

Author: apstenku123

cppmega_v4/jsonrpc/schema.py

@@ -15,7 +15,7 @@
 from pydantic import BaseModel, ConfigDict, Field
 
 
-SCHEMA_VERSION: str = "1.1.0"
+SCHEMA_VERSION: str = "1.1.1"
 
 JsonRpcVersion = Literal["2.0"]
 
@@ -369,6 +369,9 @@ class VerifyParams(BaseModel):
     side_channels: SideChannelSpecPayload = Field(
         default_factory=SideChannelSpecPayload
     )
+    data_materialization: DataMaterializationSpecPayload = Field(
+        default_factory=DataMaterializationSpecPayload
+    )
     available_side_channels: list[str] = Field(default_factory=lambda: ["doc_ids", "token_ids"])
 
 

tests/v4/test_jsonrpc_schema.py

@@ -148,12 +148,23 @@ def test_verify_params_round_trip():
     assert parsed.graph.nodes[0].id == "a"
     assert parsed.training is True
     assert parsed.side_channels.mode == "auto"
+    assert parsed.data_materialization.packing_policy == "best_fit"
+    assert parsed.data_materialization.max_seq_len == 4096
     assert parsed.side_channels.families["platform"].columns == [
         "platform_ids", "source_platform_ids",
     ]
     serial = parsed.model_dump(mode="json")
     assert "graph" in serial
     assert serial["side_channels"]["families"]["structure"]["mode"] == "if_available"
+    assert serial["data_materialization"]["required_token_fields"] == [
+        "input_ids",
+        "target_ids",
+        "loss_mask",
+        "doc_ids",
+        "pack_id",
+        "valid_token_count",
+        "num_docs",
+    ]
 
 
 def test_verify_params_rejects_unknown_field():

vbgui/src/App.tsx

@@ -763,6 +763,7 @@ function buildVerifyParams(
     },
     training: true,
     side_channels: spec.side_channels,
+    data_materialization: spec.data_materialization,
     available_side_channels: availableSideChannels,
   };
 }

vbgui/src/lib/types.ts

@@ -1,5 +1,5 @@
 // Wire-format types matching cppmega_v4.jsonrpc.schema (Pydantic side).
-// Keep in sync with SCHEMA_VERSION 1.1.0. When the backend bumps the
+// Keep in sync with SCHEMA_VERSION 1.1.1. When the backend bumps the
 // version, regen this file (codegen lands in F-D).
 
 export type JsonRpcVersion = "2.0";
@@ -24,7 +24,7 @@ export interface JsonRpcResponse<R = unknown> {
   error?: JsonRpcError;
 }
 
-export const SCHEMA_VERSION = "1.1.0";
+export const SCHEMA_VERSION = "1.1.1";
 
 // ---------------------------------------------------------------------------
 // Domain payloads
@@ -81,6 +81,16 @@ export interface SideChannelSpecPayload {
   inference: InferenceEnrichmentSpecPayload;
 }
 
+export type PackingPolicy = "sequential" | "best_fit";
+
+export interface DataMaterializationSpecPayload {
+  packing_policy: PackingPolicy;
+  max_seq_len: number;
+  pad_to_max: boolean;
+  include_provenance: boolean;
+  required_token_fields: string[];
+}
+
 export interface EdgeResolution {
   src: string;
   dst: string;

vbgui/src/state/migrations.ts

@@ -12,7 +12,7 @@
 // by a newer build) throw — the caller is expected to surface this as
 // a UI error rather than silently corrupting the spec.
 
-export const CURRENT_SCHEMA_VERSION = 1;
+export const CURRENT_SCHEMA_VERSION = 2;
 
 export type SchemaVersion = number;
 
@@ -37,9 +37,46 @@ function migrate_v0_to_v1(spec: VersionedSpec): VersionedSpec {
   return { ...spec, schema_version: 1 };
 }
 
+function defaultDataMaterialization(): Record<string, unknown> {
+  return {
+    packing_policy: "best_fit",
+    max_seq_len: 4096,
+    pad_to_max: true,
+    include_provenance: true,
+    required_token_fields: [
+      "input_ids",
+      "target_ids",
+      "loss_mask",
+      "doc_ids",
+      "pack_id",
+      "valid_token_count",
+      "num_docs",
+    ],
+  };
+}
+
+/** v1 → v2: add packed-row materialization defaults to saved GUI specs. */
+function migrate_v1_to_v2(spec: VersionedSpec): VersionedSpec {
+  const nested = spec.spec;
+  if (nested && typeof nested === "object" && !Array.isArray(nested)) {
+    const specPayload = nested as Record<string, unknown>;
+    return {
+      ...spec,
+      schema_version: 2,
+      spec: {
+        ...specPayload,
+        data_materialization:
+          specPayload.data_materialization ?? defaultDataMaterialization(),
+      },
+    };
+  }
+  return { ...spec, schema_version: 2 };
+}
+
 const MIGRATIONS: Record<number,
   (s: VersionedSpec) => VersionedSpec> = {
   0: migrate_v0_to_v1,
+  1: migrate_v1_to_v2,
 };
 
 export function migrate(spec: VersionedSpec): VersionedSpec {

vbgui/src/state/spec.ts

@@ -84,6 +84,16 @@ export interface SideChannelState {
   inference: InferenceEnrichmentState;
 }
 
+export type PackingPolicy = "sequential" | "best_fit";
+
+export interface DataMaterializationState {
+  packing_policy: PackingPolicy;
+  max_seq_len: number;
+  pad_to_max: boolean;
+  include_provenance: boolean;
+  required_token_fields: string[];
+}
+
 export type TopologyFactory =
   | "h100_8x" | "h200_8x" | "a100_8x" | "b100_8x"
   | "gb10_quarter" | "tpu_v6e_8" | "tpu_v5p_4" | "m3_ultra_solo";
@@ -115,6 +125,7 @@ export interface SpecState {
   optim: OptimState;
   rewriters: RewriterState[];
   side_channels: SideChannelState;
+  data_materialization: DataMaterializationState;
   sharding: ShardingState;
   gotchas: GotchaState[];
   worst_rank_bytes: number;
@@ -217,6 +228,21 @@ export const INITIAL_SPEC: SpecState = {
       cache_enabled: true,
     },
   },
+  data_materialization: {
+    packing_policy: "best_fit",
+    max_seq_len: 4096,
+    pad_to_max: true,
+    include_provenance: true,
+    required_token_fields: [
+      "input_ids",
+      "target_ids",
+      "loss_mask",
+      "doc_ids",
+      "pack_id",
+      "valid_token_count",
+      "num_docs",
+    ],
+  },
   sharding: {
     topology: "h100_8x",
     axis_assignments: [{ axis_name: "dp", kind: "fsdp2", degree: 8 }],

vbgui/tests/App.integration.test.tsx

@@ -110,10 +110,21 @@ describe("App integration — run pipeline", () => {
     const pipelineCall = calls.find((c) => c.method === "pipeline.run")!;
     const pipelineParams = pipelineCall.params as {
       pipeline: { stages: string[] };
+      spec: {
+        data_materialization: {
+          packing_policy: string;
+          max_seq_len: number;
+          required_token_fields: string[];
+        };
+      };
     };
     expect(pipelineParams.pipeline.stages).toContain("parse");
     expect(pipelineParams.pipeline.stages).toContain("dry_forward");
     expect(pipelineParams.pipeline.stages).not.toContain("train");
+    expect(pipelineParams.spec.data_materialization.packing_policy).toBe("best_fit");
+    expect(pipelineParams.spec.data_materialization.max_seq_len).toBe(4096);
+    expect(pipelineParams.spec.data_materialization.required_token_fields)
+      .toContain("loss_mask");
 
     await waitFor(() => {
       expect(screen.getByTestId("run-result-modal")).toBeTruthy();

vbgui/tests/migrations.test.ts

@@ -7,7 +7,7 @@ describe("V7-H04 spec migrations", () => {
   it("stamps schema_version on a v0 (pre-versioning) spec", () => {
     const v0 = { projectName: "demo", spec: {}, nodes: [], edges: [] };
     const out = migrate(v0);
-    expect(out.schema_version).toBe(1);
+    expect(out.schema_version).toBe(CURRENT_SCHEMA_VERSION);
     expect(out.projectName).toBe("demo");
   });
 
@@ -30,9 +30,40 @@ describe("V7-H04 spec migrations", () => {
       spec: { loss: { kind: "cross_entropy" }, optim: { kind: "adamw" } },
       nodes: [{ id: "a" }],
     };
-    const out = migrate(v0) as typeof v0 & { schema_version: number };
-    expect(out.schema_version).toBe(1);
-    expect(out.spec).toEqual(v0.spec);
+    const out = migrate(v0) as typeof v0 & {
+      schema_version: number;
+      spec: Record<string, unknown>;
+    };
+    expect(out.schema_version).toBe(CURRENT_SCHEMA_VERSION);
+    expect(out.spec.loss).toEqual(v0.spec.loss);
+    expect(out.spec.optim).toEqual(v0.spec.optim);
+    expect(out.spec.data_materialization).toBeDefined();
     expect(out.nodes).toEqual(v0.nodes);
   });
+
+  it("adds data materialization defaults when loading v1 specs", () => {
+    const v1 = {
+      schema_version: 1,
+      spec: { loss: { kind: "cross_entropy" } },
+    };
+    const out = migrate(v1) as typeof v1 & {
+      spec: { data_materialization?: unknown };
+    };
+    expect(out.schema_version).toBe(CURRENT_SCHEMA_VERSION);
+    expect(out.spec.data_materialization).toEqual({
+      packing_policy: "best_fit",
+      max_seq_len: 4096,
+      pad_to_max: true,
+      include_provenance: true,
+      required_token_fields: [
+        "input_ids",
+        "target_ids",
+        "loss_mask",
+        "doc_ids",
+        "pack_id",
+        "valid_token_count",
+        "num_docs",
+      ],
+    });
+  });
 });

vbgui/tests/spec.test.ts

@@ -106,6 +106,24 @@ describe("side-channel defaults", () => {
     ]);
     expect(INITIAL_SPEC.side_channels.inference.fail_policy).toBe("drop_family");
   });
+
+  it("INITIAL_SPEC exposes packed-row materialization defaults", () => {
+    expect(INITIAL_SPEC.data_materialization).toEqual({
+      packing_policy: "best_fit",
+      max_seq_len: 4096,
+      pad_to_max: true,
+      include_provenance: true,
+      required_token_fields: [
+        "input_ids",
+        "target_ids",
+        "loss_mask",
+        "doc_ids",
+        "pack_id",
+        "valid_token_count",
+        "num_docs",
+      ],
+    });
+  });
 });
 
 describe("memory helpers", () => {