diff --git a/runs/intervention-response-v2/prepare_pilot.py b/runs/intervention-response-v2/prepare_pilot.py index 2b735d5..18f1b78 100644 --- a/runs/intervention-response-v2/prepare_pilot.py +++ b/runs/intervention-response-v2/prepare_pilot.py @@ -27,7 +27,7 @@ REPLAY_TIME_SCALE = 0.5 EXPECTED_DURATION_S = 300.0 SAFETY_DEADLINE_S = 360.0 LOADS_REQ_S_GPU = {"low": 1.5, "mid": 2.125, "high": 3.125} -REPLICATE_ROLES = ("high1", "high2", "high3") +REPLICATE_ROLE_PAIRS = (("low1", "high1"), ("low2", "high2"), ("low3", "high3")) LOAD_ORDERS = { 1: ("low", "mid", "high"), 2: ("high", "low", "mid"), @@ -103,12 +103,50 @@ def selection_record(selected: list[Any], *, duration_s: float) -> dict[str, Any } -def materialize_study(source: Path, target: Path, *, replicate: int) -> Path: +def resolve_role_trace(base: dict[str, Any], role: str) -> Path: + windows_path = Path(base["private"]["windows"]) + windows = json.loads(windows_path.read_text(encoding="utf-8"))["windows"] + window_id = f"fidelity_pilot_{role}" + record = next(item for item in windows if item["window_id"] == window_id) + trace = Path(record["trace_file"]) + if not trace.is_absolute(): + trace = (windows_path.parent / trace).resolve() + return trace + + +def merge_role_traces(sources: tuple[Path, Path], target: Path) -> dict[str, Any]: + target.parent.mkdir(parents=True, exist_ok=True) + temporary = target.with_suffix(target.suffix + ".tmp") + rows = 0 + with temporary.open("w", encoding="utf-8") as output: + for source in sources: + with source.open(encoding="utf-8") as input_file: + for line in input_file: + if not line.strip(): + continue + json.loads(line) + output.write(line if line.endswith("\n") else line + "\n") + rows += 1 + os.replace(temporary, target) + return { + "path": str(target), + "sha256": sha256_file(target), + "bytes": target.stat().st_size, + "rows": rows, + "sources": [str(source) for source in sources], + "source_sha256": [sha256_file(source) for source in sources], + } + + +def materialize_study( + source: Path, target: Path, *, replicate: int, trace_override: Path +) -> Path: payload = json.loads(source.read_text(encoding="utf-8")) payload["study_id"] = f"phase-aware-telemetry-v2-rep{replicate}" payload["hardware"]["host_candidates"] = ["dash0"] payload["engine"]["engine_version"] = "0.24.1.dev3+opprof" trace = payload["trace"] + trace["trace_file_override"] = str(trace_override) trace["replay_time_scale"] = REPLAY_TIME_SCALE trace["early_stop_max_lag_s"] = None trace["early_stop_max_elapsed_s"] = SAFETY_DEADLINE_S @@ -128,10 +166,21 @@ def build_manifest( repetitions = {} selection_hashes = [] selected_ids_by_replicate: dict[int, set[str]] = {} - for replicate, role in enumerate(REPLICATE_ROLES, start=1): - source_study = Path(base["private"]["studies"][role]["tp4"]) + merged_traces = {} + for replicate, roles in enumerate(REPLICATE_ROLE_PAIRS, start=1): + source_study = Path(base["private"]["studies"][roles[1]]["tp4"]) + source_traces = tuple(resolve_role_trace(base, role) for role in roles) + merged_trace = private_root / "traces" / f"rep{replicate}.jsonl" + merged_traces[str(replicate)] = merge_role_traces( + source_traces, merged_trace + ) target_study = private_root / "studies" / f"rep{replicate}-tp4.json" - materialize_study(source_study, target_study, replicate=replicate) + materialize_study( + source_study, + target_study, + replicate=replicate, + trace_override=merged_trace, + ) study = load_study_spec(target_study) window, requests = load_trace_requests(study, study_spec_path=target_study) duration_s = float(window.window_end - window.window_start) @@ -159,7 +208,8 @@ def build_manifest( selection_hashes.append(record["request_id_order_sha256"]) selected_ids_by_replicate[replicate] = selected_ids_by_level["high"] repetitions[str(replicate)] = { - "source_role": role, + "source_roles": list(roles), + "merged_trace": merged_traces[str(replicate)], "study": str(target_study), "study_sha256": sha256_file(target_study), "duration_s": duration_s, @@ -219,6 +269,7 @@ def build_manifest( "disable_slo_early_stop": True, }, "burnin": burnin, + "private": {"merged_traces": merged_traces}, "repetitions": repetitions, "sessions": sessions, "checkpoints": {