fix: derive complete trace blocks from private artifact

runs/frontier-workload-regime-taxonomy-v0/materialize_workload_regimes.py

@@ -44,7 +44,11 @@ FAMILIES = (
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
-    parser.add_argument("--source-public", type=Path, required=True)
+    parser.add_argument(
+        "--source-public",
+        type=Path,
+        help="Optional Frontier CSV; otherwise project complete blocks from private rows.",
+    )
     parser.add_argument("--source-private", type=Path, required=True)
     parser.add_argument("--model", type=Path, required=True)
     parser.add_argument("--output-root", type=Path, required=True)
@@ -87,12 +91,32 @@ def vector_sha256(rows: Iterable[dict[str, Any]]) -> str:
     return digest.hexdigest()
 
 
-def load_source(public_path: Path, private_path: Path) -> list[dict[str, Any]]:
-    with public_path.open(newline="") as source:
-        public_rows = list(csv.DictReader(source))
+def load_source(
+    public_path: Path | None, private_path: Path
+) -> list[dict[str, Any]]:
     private_rows = [json.loads(line) for line in private_path.open() if line.strip()]
-    if not public_rows or len(public_rows) != len(private_rows):
-        raise ValueError("source public/private request count mismatch")
+    if not private_rows:
+        raise ValueError("empty private source")
+    if public_path is None:
+        public_rows = [
+            {
+                "arrived_at": row["arrived_at"],
+                "num_prefill_tokens": row["input_length"],
+                "num_decode_tokens": row["output_length"],
+                "block_hash_ids": "|".join(
+                    str(value)
+                    for value in row["runtime_block_ids"][
+                        : int(row["input_length"]) // 16
+                    ]
+                ),
+            }
+            for row in private_rows
+        ]
+    else:
+        with public_path.open(newline="") as source:
+            public_rows = list(csv.DictReader(source))
+        if not public_rows or len(public_rows) != len(private_rows):
+            raise ValueError("source public/private request count mismatch")
 
     rows = []
     for index, (public, private) in enumerate(
@@ -233,7 +257,7 @@ def write_case(
     rho: float,
     target_rate: float,
     reference_capacity: float,
-    source_public: Path,
+    source_public: Path | None,
     source_private: Path,
 ) -> dict[str, Any]:
     public_root = root / "public"
@@ -292,8 +316,17 @@ def write_case(
         },
         "first_arrival_s": arrivals[0],
         "last_arrival_s": arrivals[-1],
-        "source_public": str(source_public.resolve()),
-        "source_public_sha256": sha256(source_public),
+        "source_public": (
+            str(source_public.resolve()) if source_public is not None else None
+        ),
+        "source_public_sha256": (
+            sha256(source_public) if source_public is not None else None
+        ),
+        "frontier_block_projection": (
+            "source_public_complete_blocks"
+            if source_public is not None
+            else "private_runtime_ids_first_floor_isl_div_16_blocks"
+        ),
         "source_private": str(source_private.resolve()),
         "source_private_sha256": sha256(source_private),
         "public_csv": str(public_path.resolve()),

tests/test_materialize_workload_regimes.py

@@ -1,4 +1,5 @@
 import importlib.util
+import json
 import sys
 from pathlib import Path
 
@@ -82,3 +83,22 @@ def test_family_axes_change_only_the_declared_contract():
         "prompt-1",
         "prompt-2",
     ]
+
+
+def test_private_source_projects_only_complete_prefix_blocks(tmp_path):
+    private_path = tmp_path / "source.jsonl"
+    source = {
+        "source_index": 7,
+        "arrived_at": 1.5,
+        "input_length": 17,
+        "output_length": 4,
+        "session_id": 11,
+        "runtime_block_ids": [101, 102],
+        "body": {"prompt": "fixture"},
+    }
+    private_path.write_text(json.dumps(source) + "\n")
+
+    rows = MODULE.load_source(None, private_path)
+
+    assert rows[0]["runtime_block_ids"] == [101]
+    assert rows[0]["source_arrived_at"] == 1.5