Support long code traces in Frontier replay
This commit is contained in:
114
runs/frontier-code-trace-v0/test_frontier_code_replay.py
Normal file
114
runs/frontier-code-trace-v0/test_frontier_code_replay.py
Normal file
@@ -0,0 +1,114 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import csv
|
||||||
|
import importlib.util
|
||||||
|
import json
|
||||||
|
import sys
|
||||||
|
import tempfile
|
||||||
|
import unittest
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
ROOT = Path(__file__).resolve().parent
|
||||||
|
S3_REAL = ROOT.parent / "frontier-s3-real-v0"
|
||||||
|
|
||||||
|
|
||||||
|
def load_replay_module():
|
||||||
|
sys.path.insert(0, str(S3_REAL))
|
||||||
|
spec = importlib.util.spec_from_file_location(
|
||||||
|
"frontier_prefix_replay", S3_REAL / "run_frontier_prefix_replay.py"
|
||||||
|
)
|
||||||
|
module = importlib.util.module_from_spec(spec)
|
||||||
|
spec.loader.exec_module(module)
|
||||||
|
return module
|
||||||
|
|
||||||
|
|
||||||
|
class FrontierCodeReplayTest(unittest.TestCase):
|
||||||
|
def test_summary_excludes_undefined_single_token_tpot(self) -> None:
|
||||||
|
module = load_replay_module()
|
||||||
|
with tempfile.TemporaryDirectory() as tmp:
|
||||||
|
root = Path(tmp)
|
||||||
|
trace = root / "frontier.csv"
|
||||||
|
with trace.open("w", newline="") as stream:
|
||||||
|
writer = csv.DictWriter(
|
||||||
|
stream,
|
||||||
|
fieldnames=(
|
||||||
|
"arrived_at",
|
||||||
|
"num_prefill_tokens",
|
||||||
|
"num_decode_tokens",
|
||||||
|
"session_id",
|
||||||
|
"block_hash_ids",
|
||||||
|
),
|
||||||
|
)
|
||||||
|
writer.writeheader()
|
||||||
|
writer.writerow(
|
||||||
|
{
|
||||||
|
"arrived_at": 0,
|
||||||
|
"num_prefill_tokens": 7,
|
||||||
|
"num_decode_tokens": 1,
|
||||||
|
"session_id": 1,
|
||||||
|
"block_hash_ids": "[]",
|
||||||
|
}
|
||||||
|
)
|
||||||
|
writer.writerow(
|
||||||
|
{
|
||||||
|
"arrived_at": 1,
|
||||||
|
"num_prefill_tokens": 16,
|
||||||
|
"num_decode_tokens": 2,
|
||||||
|
"session_id": 2,
|
||||||
|
"block_hash_ids": 123,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
metrics = root / "metrics" / "cell"
|
||||||
|
metrics.mkdir(parents=True)
|
||||||
|
with (metrics / "request_metrics.csv").open("w", newline="") as stream:
|
||||||
|
writer = csv.DictWriter(
|
||||||
|
stream,
|
||||||
|
fieldnames=(
|
||||||
|
"request_waiting_time_total",
|
||||||
|
"ttft",
|
||||||
|
"tpot",
|
||||||
|
"request_e2e_time",
|
||||||
|
"request_cached_prefill_tokens",
|
||||||
|
"request_prefix_cache_query_blocks",
|
||||||
|
"request_prefix_cache_hit_blocks",
|
||||||
|
),
|
||||||
|
)
|
||||||
|
writer.writeheader()
|
||||||
|
writer.writerow(
|
||||||
|
{
|
||||||
|
"request_waiting_time_total": 0,
|
||||||
|
"ttft": 10,
|
||||||
|
"tpot": "",
|
||||||
|
"request_e2e_time": 10,
|
||||||
|
"request_cached_prefill_tokens": 0,
|
||||||
|
"request_prefix_cache_query_blocks": 0,
|
||||||
|
"request_prefix_cache_hit_blocks": 0,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
writer.writerow(
|
||||||
|
{
|
||||||
|
"request_waiting_time_total": 0,
|
||||||
|
"ttft": 20,
|
||||||
|
"tpot": 5,
|
||||||
|
"request_e2e_time": 25,
|
||||||
|
"request_cached_prefill_tokens": 0,
|
||||||
|
"request_prefix_cache_query_blocks": 1,
|
||||||
|
"request_prefix_cache_hit_blocks": 0,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
(metrics / "system_metrics.json").write_text("{}")
|
||||||
|
(metrics / "frontier_stage_batch_ledger.jsonl").write_text(
|
||||||
|
json.dumps({"request_num_tokens": [1]}) + "\n"
|
||||||
|
)
|
||||||
|
summary = module.summarize(trace, root / "metrics", 2)
|
||||||
|
|
||||||
|
self.assertEqual(summary["requests"], 2)
|
||||||
|
self.assertEqual(summary["latency_ms"]["tpot"]["count"], 1)
|
||||||
|
self.assertEqual(summary["latency_ms"]["tpot"]["mean"], 5)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
192
runs/frontier-collective-joint-v0/run_frontier_with_curves.py
Normal file
192
runs/frontier-collective-joint-v0/run_frontier_with_curves.py
Normal file
@@ -0,0 +1,192 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Run Frontier with pure-decode serving-path MoE and/or collective curves."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import atexit
|
||||||
|
import csv
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
from collections import Counter
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
csv.field_size_limit(16 * 1024 * 1024)
|
||||||
|
|
||||||
|
|
||||||
|
COLLECTIVE_PAYLOAD = json.loads(
|
||||||
|
Path(os.environ["FRONTIER_COLLECTIVE_CURVE"]).read_text()
|
||||||
|
)
|
||||||
|
COLLECTIVE_VARIANT = os.environ.get("FRONTIER_COLLECTIVE_CURVE_VARIANT", "drop_mean")
|
||||||
|
if "curve_variants_ms_per_step" in COLLECTIVE_PAYLOAD:
|
||||||
|
try:
|
||||||
|
COLLECTIVE_CURVE = COLLECTIVE_PAYLOAD["curve_variants_ms_per_step"][
|
||||||
|
COLLECTIVE_VARIANT
|
||||||
|
]
|
||||||
|
except KeyError as error:
|
||||||
|
raise ValueError(
|
||||||
|
f"collective curve has no variant {COLLECTIVE_VARIANT!r}"
|
||||||
|
) from error
|
||||||
|
else:
|
||||||
|
if COLLECTIVE_VARIANT != "drop_mean":
|
||||||
|
raise ValueError("legacy collective curve only supports drop_mean")
|
||||||
|
COLLECTIVE_CURVE = COLLECTIVE_PAYLOAD.get(
|
||||||
|
"curve_ms_per_step", COLLECTIVE_PAYLOAD
|
||||||
|
)
|
||||||
|
MOE_CURVE_PATH = os.environ.get("FRONTIER_FUSED_MOE_CURVE")
|
||||||
|
MOE_CURVE = (
|
||||||
|
json.loads(Path(MOE_CURVE_PATH).read_text()) if MOE_CURVE_PATH else None
|
||||||
|
)
|
||||||
|
USAGE_PATH = Path(os.environ["FRONTIER_CURVE_USAGE"])
|
||||||
|
USAGE: Counter[str] = Counter()
|
||||||
|
|
||||||
|
|
||||||
|
def write_usage() -> None:
|
||||||
|
USAGE_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
USAGE_PATH.write_text(json.dumps(dict(sorted(USAGE.items())), indent=2) + "\n")
|
||||||
|
|
||||||
|
|
||||||
|
atexit.register(write_usage)
|
||||||
|
|
||||||
|
from frontier.execution_time_predictor.sklearn_moe_execution_time_predictor import ( # noqa: E402
|
||||||
|
SklearnMoEExecutionTimePredictor,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
_ORIGINAL_ATTN_COLLECTIVE = (
|
||||||
|
SklearnMoEExecutionTimePredictor._get_tensor_parallel_communication_time
|
||||||
|
)
|
||||||
|
_ORIGINAL_MOE_COLLECTIVE = (
|
||||||
|
SklearnMoEExecutionTimePredictor._get_moe_tensor_parallel_allreduce_time
|
||||||
|
)
|
||||||
|
_ORIGINAL_GROUPED_GEMM = SklearnMoEExecutionTimePredictor._get_grouped_gemm_time
|
||||||
|
_ORIGINAL_ATTN_NORM = (
|
||||||
|
SklearnMoEExecutionTimePredictor._get_attn_norm_layer_act_execution_time
|
||||||
|
)
|
||||||
|
_ORIGINAL_MLP_NORM = (
|
||||||
|
SklearnMoEExecutionTimePredictor._get_mlp_norm_layer_act_execution_time
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _pure_decode_point(self, batch) -> tuple[str, str] | None:
|
||||||
|
if batch is None or not bool(getattr(batch, "is_pure_decode_batch", False)):
|
||||||
|
return None
|
||||||
|
tp = str(int(self._replica_config.moe_tensor_parallel_size))
|
||||||
|
attn_tp = str(int(self._replica_config.attn_tensor_parallel_size))
|
||||||
|
if tp != attn_tp:
|
||||||
|
raise ValueError(
|
||||||
|
"Serving collective curve requires equal attention/MoE TP, got "
|
||||||
|
f"attn_tp={attn_tp}, moe_tp={tp}"
|
||||||
|
)
|
||||||
|
return tp, str(len(batch.requests))
|
||||||
|
|
||||||
|
|
||||||
|
def _collective_path_time(self, batch, path: str) -> float:
|
||||||
|
point = _pure_decode_point(self, batch)
|
||||||
|
if point is None:
|
||||||
|
original = (
|
||||||
|
_ORIGINAL_ATTN_COLLECTIVE
|
||||||
|
if path == "attention"
|
||||||
|
else _ORIGINAL_MOE_COLLECTIVE
|
||||||
|
)
|
||||||
|
return original(self, batch)
|
||||||
|
tp, decode_batch = point
|
||||||
|
if tp == "1":
|
||||||
|
original = (
|
||||||
|
_ORIGINAL_ATTN_COLLECTIVE
|
||||||
|
if path == "attention"
|
||||||
|
else _ORIGINAL_MOE_COLLECTIVE
|
||||||
|
)
|
||||||
|
value = original(self, batch)
|
||||||
|
if value != 0.0:
|
||||||
|
raise ValueError(f"TP1 {path} collective must be zero, got {value}")
|
||||||
|
USAGE[f"collective:{path}:tp1-b{decode_batch}:structural-zero"] += 1
|
||||||
|
return value
|
||||||
|
if tp not in COLLECTIVE_CURVE or decode_batch not in COLLECTIVE_CURVE[tp]:
|
||||||
|
raise ValueError(
|
||||||
|
"Collective curve has no exact pure-decode point for "
|
||||||
|
f"TP={tp}, batch={decode_batch}; refusing to extrapolate"
|
||||||
|
)
|
||||||
|
layers = int(self._num_layers_per_pipeline_stage)
|
||||||
|
if layers <= 0:
|
||||||
|
raise ValueError(f"invalid layers per pipeline stage: {layers}")
|
||||||
|
USAGE[f"collective:{path}:tp{tp}-b{decode_batch}"] += 1
|
||||||
|
# The serving curve is the sum of attention and MoE output reductions over
|
||||||
|
# the full model step. ExecutionTime later multiplies each per-layer path.
|
||||||
|
return float(COLLECTIVE_CURVE[tp][decode_batch]) / (2 * layers)
|
||||||
|
|
||||||
|
|
||||||
|
def _attention_collective_from_curve(self, batch) -> float:
|
||||||
|
return _collective_path_time(self, batch, "attention")
|
||||||
|
|
||||||
|
|
||||||
|
def _moe_collective_from_curve(self, batch) -> float:
|
||||||
|
return _collective_path_time(self, batch, "moe")
|
||||||
|
|
||||||
|
|
||||||
|
def _grouped_gemm_from_curve(self, num_tokens_or_allocation, batch=None) -> float:
|
||||||
|
if MOE_CURVE is None:
|
||||||
|
return _ORIGINAL_GROUPED_GEMM(
|
||||||
|
self, num_tokens_or_allocation, batch=batch
|
||||||
|
)
|
||||||
|
point = _pure_decode_point(self, batch)
|
||||||
|
if point is None:
|
||||||
|
return _ORIGINAL_GROUPED_GEMM(
|
||||||
|
self, num_tokens_or_allocation, batch=batch
|
||||||
|
)
|
||||||
|
tp, decode_batch = point
|
||||||
|
if tp not in MOE_CURVE or decode_batch not in MOE_CURVE[tp]:
|
||||||
|
raise ValueError(
|
||||||
|
"Fused MoE curve has no exact pure-decode point for "
|
||||||
|
f"TP={tp}, batch={decode_batch}; refusing to extrapolate"
|
||||||
|
)
|
||||||
|
layers = int(self._num_layers_per_pipeline_stage)
|
||||||
|
if layers <= 0:
|
||||||
|
raise ValueError(f"invalid layers per pipeline stage: {layers}")
|
||||||
|
USAGE[f"moe:tp{tp}-b{decode_batch}"] += 1
|
||||||
|
return float(MOE_CURVE[tp][decode_batch]) / layers
|
||||||
|
|
||||||
|
|
||||||
|
def _norm_without_fused_duplicate(self, batch, *, name: str) -> float:
|
||||||
|
original = _ORIGINAL_ATTN_NORM if name == "attn" else _ORIGINAL_MLP_NORM
|
||||||
|
value = original(self, batch)
|
||||||
|
point = _pure_decode_point(self, batch)
|
||||||
|
if point is None or point[0] == "1":
|
||||||
|
return value
|
||||||
|
if value < 0:
|
||||||
|
raise ValueError(f"negative {name} norm prediction cannot be deducted: {value}")
|
||||||
|
tp, decode_batch = point
|
||||||
|
USAGE[f"fused_norm_deduction:{name}:tp{tp}-b{decode_batch}"] += 1
|
||||||
|
# AllReduceFusionPattern 1 already contains residual add + RMSNorm. Returning
|
||||||
|
# zero removes the exact predictor value that would otherwise be emitted in
|
||||||
|
# the frozen Frontier attn_norm_time/mlp_norm_time ledger row.
|
||||||
|
return 0.0
|
||||||
|
|
||||||
|
|
||||||
|
def _attn_norm_without_fused_duplicate(self, batch) -> float:
|
||||||
|
return _norm_without_fused_duplicate(self, batch, name="attn")
|
||||||
|
|
||||||
|
|
||||||
|
def _mlp_norm_without_fused_duplicate(self, batch) -> float:
|
||||||
|
return _norm_without_fused_duplicate(self, batch, name="mlp")
|
||||||
|
|
||||||
|
|
||||||
|
SklearnMoEExecutionTimePredictor._get_tensor_parallel_communication_time = (
|
||||||
|
_attention_collective_from_curve
|
||||||
|
)
|
||||||
|
SklearnMoEExecutionTimePredictor._get_moe_tensor_parallel_allreduce_time = (
|
||||||
|
_moe_collective_from_curve
|
||||||
|
)
|
||||||
|
SklearnMoEExecutionTimePredictor._get_grouped_gemm_time = _grouped_gemm_from_curve
|
||||||
|
SklearnMoEExecutionTimePredictor._get_attn_norm_layer_act_execution_time = (
|
||||||
|
_attn_norm_without_fused_duplicate
|
||||||
|
)
|
||||||
|
SklearnMoEExecutionTimePredictor._get_mlp_norm_layer_act_execution_time = (
|
||||||
|
_mlp_norm_without_fused_duplicate
|
||||||
|
)
|
||||||
|
|
||||||
|
from frontier.main import main # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
105
runs/frontier-prefill-kvgrowth-fix-v0/run_replay_v5.py
Normal file
105
runs/frontier-prefill-kvgrowth-fix-v0/run_replay_v5.py
Normal file
@@ -0,0 +1,105 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Replay an S3-real cell with profile-v5-kvgrowth (CPU, no GPU needed).
|
||||||
|
|
||||||
|
Adapted from runs/frontier-s3-real-v0/run_frontier_prefix_replay.py with one
|
||||||
|
change: the attention profile input is overridden to profile-v5-kvgrowth so the
|
||||||
|
predictor retrains with the KV-context grid. Everything else (curves, trace,
|
||||||
|
config argv, prefix caching) is identical to the S3-real sim runs, so old-vs-new
|
||||||
|
differs by exactly one variable.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import csv
|
||||||
|
import importlib.util
|
||||||
|
import json
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
csv.field_size_limit(16 * 1024 * 1024)
|
||||||
|
|
||||||
|
|
||||||
|
ROOT = Path(__file__).resolve().parent
|
||||||
|
S3_REAL = ROOT.parent / "frontier-s3-real-v0"
|
||||||
|
|
||||||
|
|
||||||
|
def load_s3_module():
|
||||||
|
spec = importlib.util.spec_from_file_location(
|
||||||
|
"s3_prefix_replay", S3_REAL / "run_frontier_prefix_replay.py"
|
||||||
|
)
|
||||||
|
module = importlib.util.module_from_spec(spec)
|
||||||
|
sys.path.insert(0, str(S3_REAL)) # trace_utils import
|
||||||
|
spec.loader.exec_module(module)
|
||||||
|
return module
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
parser.add_argument("--trace", type=Path, required=True)
|
||||||
|
parser.add_argument("--output-root", type=Path, required=True)
|
||||||
|
parser.add_argument("--config", choices=("tp4_mns16", "tp2_mns16", "tp1_mns16"), required=True)
|
||||||
|
parser.add_argument("--label", required=True)
|
||||||
|
parser.add_argument("--max-tokens", type=int, required=True)
|
||||||
|
parser.add_argument("--duration-s", type=float)
|
||||||
|
parser.add_argument("--cache-root", type=Path)
|
||||||
|
parser.add_argument(
|
||||||
|
"--reuse-model-cache",
|
||||||
|
action="store_true",
|
||||||
|
help="Load predictor models from the content-addressed cache when available.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--attention-profile",
|
||||||
|
type=Path,
|
||||||
|
default=ROOT / "profiles/profile-v5-kvgrowth/attention.csv",
|
||||||
|
)
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
profile = args.attention_profile.resolve()
|
||||||
|
if not profile.is_file():
|
||||||
|
raise SystemExit(f"attention profile missing: {profile}")
|
||||||
|
|
||||||
|
module = load_s3_module()
|
||||||
|
original_replace = module.replace_flag
|
||||||
|
|
||||||
|
def replace_and_override(argv: list[str], flag: str, value: str) -> None:
|
||||||
|
original_replace(argv, flag, value)
|
||||||
|
# Piggyback on the first replace_flag call (trace file) to inject the
|
||||||
|
# profile override exactly once per run.
|
||||||
|
atten_flag = "--random_forrest_execution_time_predictor_config_atten_input_file"
|
||||||
|
if flag.endswith("trace_file") and atten_flag in argv:
|
||||||
|
original_replace(argv, atten_flag, str(profile))
|
||||||
|
no_cache_flag = "--random_forrest_execution_time_predictor_config_no_cache"
|
||||||
|
if args.reuse_model_cache and no_cache_flag in argv:
|
||||||
|
argv.remove(no_cache_flag)
|
||||||
|
|
||||||
|
module.replace_flag = replace_and_override
|
||||||
|
# The shared S3 parser predates TP1; reuse this wrapper's validated namespace.
|
||||||
|
module.parse_args = lambda: args
|
||||||
|
sys.argv = [
|
||||||
|
"run_frontier_prefix_replay.py",
|
||||||
|
"--trace", str(args.trace),
|
||||||
|
"--output-root", str(args.output_root),
|
||||||
|
"--config", args.config,
|
||||||
|
"--label", args.label,
|
||||||
|
"--max-tokens", str(args.max_tokens),
|
||||||
|
]
|
||||||
|
if args.duration_s:
|
||||||
|
sys.argv += ["--duration-s", str(args.duration_s)]
|
||||||
|
if args.cache_root:
|
||||||
|
sys.argv += ["--cache-root", str(args.cache_root)]
|
||||||
|
module.main()
|
||||||
|
|
||||||
|
manifest_path = args.output_root / "manifest.json"
|
||||||
|
manifest = json.loads(manifest_path.read_text())
|
||||||
|
manifest["attention_profile_override"] = str(profile)
|
||||||
|
manifest["attention_profile_sha256"] = module.sha256(profile)
|
||||||
|
manifest["reuse_model_cache"] = args.reuse_model_cache
|
||||||
|
manifest["schema"] = "frontier-prefill-kvgrowth-replay-v1"
|
||||||
|
manifest_path.write_text(json.dumps(manifest, indent=2))
|
||||||
|
print(f"replay done: {args.output_root}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
235
runs/frontier-s3-real-v0/run_frontier_prefix_replay.py
Normal file
235
runs/frontier-s3-real-v0/run_frontier_prefix_replay.py
Normal file
@@ -0,0 +1,235 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Replay a remapped trace with S1 joint injection and prefix caching enabled."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import csv
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import subprocess
|
||||||
|
from collections import Counter
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from trace_utils import distribution, sha256, write_json
|
||||||
|
|
||||||
|
|
||||||
|
ROOT = Path(__file__).resolve().parent
|
||||||
|
REPO = ROOT.parents[1]
|
||||||
|
REFERENCE = REPO / "runs/frontier-collective-joint-v0/counterfactual/joint-r2/manifest.json"
|
||||||
|
EXPECTED_FRONTIER_COMMIT = "deadc4a321f0baaa534c6ebd17f974123733cdc2"
|
||||||
|
MAX_CURVE_BATCH = 32
|
||||||
|
|
||||||
|
|
||||||
|
csv.field_size_limit(16 * 1024 * 1024)
|
||||||
|
|
||||||
|
|
||||||
|
def parse_args() -> argparse.Namespace:
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
parser.add_argument("--trace", type=Path, required=True)
|
||||||
|
parser.add_argument("--output-root", type=Path, required=True)
|
||||||
|
parser.add_argument("--config", choices=("tp4_mns16", "tp2_mns16"), default="tp4_mns16")
|
||||||
|
parser.add_argument("--label", required=True)
|
||||||
|
parser.add_argument("--max-tokens", type=int, required=True)
|
||||||
|
parser.add_argument("--duration-s", type=float)
|
||||||
|
parser.add_argument("--cache-root", type=Path)
|
||||||
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
def replace_flag(argv: list[str], flag: str, value: str) -> None:
|
||||||
|
try:
|
||||||
|
index = argv.index(flag)
|
||||||
|
except ValueError as error:
|
||||||
|
raise ValueError(f"template command is missing {flag}") from error
|
||||||
|
argv[index + 1] = value
|
||||||
|
|
||||||
|
|
||||||
|
def prepare_curves(reference: dict[str, Any], output_root: Path) -> dict[str, Any]:
|
||||||
|
inputs = output_root / "inputs"
|
||||||
|
inputs.mkdir(parents=True)
|
||||||
|
outputs: dict[str, Any] = {}
|
||||||
|
for kind, reference_key, filename in (
|
||||||
|
("collective", "collective_curve", "collective-curve-b4-extrapolated.json"),
|
||||||
|
("moe", "moe_curve", "fused-moe-curve-b4-extrapolated.json"),
|
||||||
|
):
|
||||||
|
payload = json.loads(Path(reference[reference_key]).read_text())
|
||||||
|
curves = (
|
||||||
|
payload["curve_variants_ms_per_step"].values()
|
||||||
|
if kind == "collective" and "curve_variants_ms_per_step" in payload
|
||||||
|
else (payload,)
|
||||||
|
)
|
||||||
|
for curve in curves:
|
||||||
|
for points in curve.values():
|
||||||
|
b4 = points["4"]
|
||||||
|
for batch in range(5, MAX_CURVE_BATCH + 1):
|
||||||
|
points[str(batch)] = b4
|
||||||
|
destination = inputs / filename
|
||||||
|
write_json(destination, payload)
|
||||||
|
outputs[kind] = str(destination)
|
||||||
|
outputs[f"{kind}_sha256"] = sha256(destination)
|
||||||
|
return outputs
|
||||||
|
|
||||||
|
|
||||||
|
def find_one(root: Path, name: str) -> Path:
|
||||||
|
matches = list(root.glob(f"**/{name}"))
|
||||||
|
if len(matches) != 1:
|
||||||
|
raise ValueError(f"expected one {name} below {root}, found {matches}")
|
||||||
|
return matches[0]
|
||||||
|
|
||||||
|
|
||||||
|
def read_csv(path: Path) -> list[dict[str, str]]:
|
||||||
|
with path.open(newline="") as stream:
|
||||||
|
return list(csv.DictReader(stream))
|
||||||
|
|
||||||
|
|
||||||
|
def summarize(trace: Path, metrics_root: Path, duration_s: float) -> dict[str, Any]:
|
||||||
|
trace_rows = read_csv(trace)
|
||||||
|
metrics_path = find_one(metrics_root, "request_metrics.csv")
|
||||||
|
metric_rows = read_csv(metrics_path)
|
||||||
|
if len(trace_rows) != len(metric_rows):
|
||||||
|
raise ValueError(f"request count mismatch: trace={len(trace_rows)}, metrics={len(metric_rows)}")
|
||||||
|
system_path = find_one(metrics_root, "system_metrics.json")
|
||||||
|
system = json.loads(system_path.read_text())
|
||||||
|
ledger_path = find_one(metrics_root, "frontier_stage_batch_ledger.jsonl")
|
||||||
|
waiting_ms = [float(row["request_waiting_time_total"]) for row in metric_rows]
|
||||||
|
ttft_ms = [float(row["ttft"]) for row in metric_rows]
|
||||||
|
# TPOT is undefined for one-token outputs because there is no inter-token
|
||||||
|
# interval. Frontier records those cells as an empty CSV value.
|
||||||
|
tpot_ms = [
|
||||||
|
float(row["tpot"])
|
||||||
|
for row in metric_rows
|
||||||
|
if row.get("tpot") is not None and row["tpot"].strip()
|
||||||
|
]
|
||||||
|
e2e_ms = [float(row["request_e2e_time"]) for row in metric_rows]
|
||||||
|
cached_tokens = [int(float(row.get("request_cached_prefill_tokens", 0))) for row in metric_rows]
|
||||||
|
query_blocks = [int(float(row.get("request_prefix_cache_query_blocks", 0))) for row in metric_rows]
|
||||||
|
hit_blocks = [int(float(row.get("request_prefix_cache_hit_blocks", 0))) for row in metric_rows]
|
||||||
|
arrivals = [float(row["arrived_at"]) for row in trace_rows]
|
||||||
|
prefill = [int(row["num_prefill_tokens"]) for row in trace_rows]
|
||||||
|
decode = [int(row["num_decode_tokens"]) for row in trace_rows]
|
||||||
|
completion_times = [arrival + e2e / 1000 for arrival, e2e in zip(arrivals, e2e_ms)]
|
||||||
|
batch_hist: Counter[int] = Counter()
|
||||||
|
with ledger_path.open() as stream:
|
||||||
|
for line in stream:
|
||||||
|
if not line.strip():
|
||||||
|
continue
|
||||||
|
row = json.loads(line)
|
||||||
|
tokens = row.get("request_num_tokens") or []
|
||||||
|
if tokens and all(int(token) == 1 for token in tokens):
|
||||||
|
batch_hist[len(tokens)] += 1
|
||||||
|
stages = sum(batch_hist.values())
|
||||||
|
prefix = system.get("prefix_cache_statistics") or {
|
||||||
|
"block_size_tokens": 16,
|
||||||
|
"requests": len(metric_rows),
|
||||||
|
"total_cached_prefill_tokens": sum(cached_tokens),
|
||||||
|
"total_query_blocks": sum(query_blocks),
|
||||||
|
"total_hit_blocks": sum(hit_blocks),
|
||||||
|
"hit_ratio": sum(hit_blocks) / sum(query_blocks) if sum(query_blocks) else 0.0,
|
||||||
|
}
|
||||||
|
return {
|
||||||
|
"schema": "frontier-s3-real-prefix-replay-summary-v1",
|
||||||
|
"requests": len(trace_rows),
|
||||||
|
"duration_s": duration_s,
|
||||||
|
"offered_load": {
|
||||||
|
"requests_per_s": len(trace_rows) / duration_s,
|
||||||
|
"prefill_tokens_per_s_raw": sum(prefill) / duration_s,
|
||||||
|
"prefill_tokens_per_s_after_prefix": (sum(prefill) - sum(cached_tokens)) / duration_s,
|
||||||
|
"decode_tokens_per_s": sum(decode) / duration_s,
|
||||||
|
},
|
||||||
|
"latency_ms": {
|
||||||
|
"waiting": distribution(waiting_ms),
|
||||||
|
"ttft": distribution(ttft_ms),
|
||||||
|
"tpot": distribution(tpot_ms),
|
||||||
|
"e2e": distribution(e2e_ms),
|
||||||
|
},
|
||||||
|
"drain": {
|
||||||
|
"last_arrival_s": max(arrivals),
|
||||||
|
"last_completion_s": max(completion_times),
|
||||||
|
"tail_after_last_arrival_s": max(completion_times) - max(arrivals),
|
||||||
|
},
|
||||||
|
"decode_batch": {
|
||||||
|
"stages": stages,
|
||||||
|
"histogram": dict(sorted(batch_hist.items())),
|
||||||
|
"share_gt_1": sum(count for batch, count in batch_hist.items() if batch > 1) / stages if stages else 0.0,
|
||||||
|
"share_gt_4": sum(count for batch, count in batch_hist.items() if batch > 4) / stages if stages else 0.0,
|
||||||
|
"max": max(batch_hist, default=0),
|
||||||
|
},
|
||||||
|
"prefix_cache": prefix,
|
||||||
|
"artifacts": {
|
||||||
|
"request_metrics": str(metrics_path),
|
||||||
|
"system_metrics": str(system_path),
|
||||||
|
"stage_ledger": str(ledger_path),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
args = parse_args()
|
||||||
|
args.trace = args.trace.resolve()
|
||||||
|
args.output_root = args.output_root.resolve()
|
||||||
|
args.cache_root = (args.cache_root or args.output_root.parent / "cache").resolve()
|
||||||
|
if args.output_root.exists():
|
||||||
|
raise ValueError(f"refusing to overwrite {args.output_root}")
|
||||||
|
reference = json.loads(REFERENCE.read_text())
|
||||||
|
frontier = Path(reference["frontier_checkout"])
|
||||||
|
commit = subprocess.check_output(["git", "-C", str(frontier), "rev-parse", "HEAD"], text=True).strip()
|
||||||
|
if commit != EXPECTED_FRONTIER_COMMIT or commit != reference["frontier_commit"]:
|
||||||
|
raise ValueError(f"Frontier commit drift: {commit}")
|
||||||
|
if subprocess.check_output(["git", "-C", str(frontier), "status", "--porcelain"], text=True).strip():
|
||||||
|
raise ValueError("Frontier checkout must be clean")
|
||||||
|
args.output_root.mkdir(parents=True)
|
||||||
|
curves = prepare_curves(reference, args.output_root)
|
||||||
|
cell = reference["cells"][args.config]
|
||||||
|
argv = list(cell["argv"])
|
||||||
|
replace_flag(argv, "--trace_request_generator_config_trace_file", str(args.trace))
|
||||||
|
replace_flag(argv, "--trace_request_generator_config_max_tokens", str(args.max_tokens))
|
||||||
|
replace_flag(argv, "--random_forrest_execution_time_predictor_config_prediction_max_tokens_per_request", str(args.max_tokens))
|
||||||
|
replace_flag(argv, "--metrics_config_output_dir", str(args.output_root / "metrics"))
|
||||||
|
replace_flag(argv, "--metrics_config_run_id", args.label)
|
||||||
|
replace_flag(argv, "--metrics_config_cache_dir", str(args.cache_root / "model"))
|
||||||
|
replace_flag(argv, "--vidur_cc_backend_config_cache_dir", str(args.cache_root / "cc"))
|
||||||
|
argv.extend(("--vllm_v1_scheduler_config_enable_prefix_caching", "--log_level", "warning", "--cluster_event_log_level", "WARNING"))
|
||||||
|
usage = args.output_root / "usage.json"
|
||||||
|
env = os.environ.copy()
|
||||||
|
env.update(
|
||||||
|
{
|
||||||
|
"CUDA_VISIBLE_DEVICES": "",
|
||||||
|
"PYTHONDONTWRITEBYTECODE": "1",
|
||||||
|
"PYTHONPATH": os.pathsep.join([str(frontier), *reference["python_dependency_roots"]]),
|
||||||
|
"FRONTIER_COLLECTIVE_CURVE": curves["collective"],
|
||||||
|
"FRONTIER_COLLECTIVE_CURVE_VARIANT": reference["collective_curve_variant"],
|
||||||
|
"FRONTIER_FUSED_MOE_CURVE": curves["moe"],
|
||||||
|
"FRONTIER_CURVE_USAGE": str(usage),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
manifest = {
|
||||||
|
"schema": "frontier-s3-real-prefix-replay-v1",
|
||||||
|
"frontier_commit": commit,
|
||||||
|
"reference_manifest": str(REFERENCE.resolve()),
|
||||||
|
"reference_manifest_sha256": sha256(REFERENCE),
|
||||||
|
"trace": str(args.trace),
|
||||||
|
"trace_sha256": sha256(args.trace),
|
||||||
|
"config": args.config,
|
||||||
|
"prefix_caching": True,
|
||||||
|
"block_size": 16,
|
||||||
|
"max_tokens": args.max_tokens,
|
||||||
|
"argv": argv,
|
||||||
|
"curves": curves,
|
||||||
|
}
|
||||||
|
write_json(args.output_root / "manifest.json", manifest)
|
||||||
|
log_path = args.output_root / "sim.log"
|
||||||
|
with log_path.open("w") as log:
|
||||||
|
completed = subprocess.run(argv, cwd=frontier, env=env, stdout=log, stderr=subprocess.STDOUT, check=False)
|
||||||
|
if completed.returncode:
|
||||||
|
raise SystemExit(f"sim failed with exit code {completed.returncode}; see {log_path}")
|
||||||
|
with args.trace.open(newline="") as stream:
|
||||||
|
trace_rows = list(csv.DictReader(stream))
|
||||||
|
duration_s = args.duration_s or max(float(row["arrived_at"]) for row in trace_rows) or 1.0
|
||||||
|
summary = summarize(args.trace, args.output_root / "metrics", duration_s)
|
||||||
|
write_json(args.output_root / "summary.json", summary)
|
||||||
|
print(json.dumps(summary, sort_keys=True))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
Reference in New Issue
Block a user