diff --git a/runs/frontier-workload-regime-taxonomy-v0/materialize_workload_regimes.py b/runs/frontier-workload-regime-taxonomy-v0/materialize_workload_regimes.py index 13fcb11..3265e65 100644 --- a/runs/frontier-workload-regime-taxonomy-v0/materialize_workload_regimes.py +++ b/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()), diff --git a/tests/test_materialize_workload_regimes.py b/tests/test_materialize_workload_regimes.py index 3545db3..e9772ac 100644 --- a/tests/test_materialize_workload_regimes.py +++ b/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