164 lines
4.7 KiB
Python
164 lines
4.7 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import math
|
|
import sys
|
|
import tempfile
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
|
|
HERE = Path(__file__).resolve().parent
|
|
sys.path.insert(0, str(HERE))
|
|
|
|
import pilot_controller as controller # noqa: E402
|
|
import prepare_pilot as prepare # noqa: E402
|
|
from analyze_pilot import campaign_gpu_accounting # noqa: E402
|
|
|
|
|
|
@dataclass
|
|
class Request:
|
|
row_id: str
|
|
sampling_u: float
|
|
arrival_s: float = 0.0
|
|
prompt_tokens_hint: int = 1
|
|
|
|
|
|
def main() -> None:
|
|
requests = [
|
|
Request("a", 0.1),
|
|
Request("b", 0.2),
|
|
Request("c", 0.2),
|
|
Request("d", 0.9),
|
|
]
|
|
anchor, selected = prepare.attainable_anchor(requests, target_count=2)
|
|
assert anchor == 0.2
|
|
assert [request.row_id for request in selected] == ["a", "b", "c"]
|
|
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
root = Path(temporary)
|
|
source = root / "source.jsonl"
|
|
rows = []
|
|
for index, role in enumerate(prepare.ROLES):
|
|
rows.append(
|
|
{
|
|
"request_id": role,
|
|
"timestamp": float(index),
|
|
"sampling_u": (index + 0.5) / len(prepare.ROLES),
|
|
"input_length": 16 + index,
|
|
"messages": [{"role": "user", "content": role}],
|
|
}
|
|
)
|
|
source.write_text(
|
|
"".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8"
|
|
)
|
|
windows, stats = prepare.materialize_bands(
|
|
source,
|
|
{
|
|
"window_id": "source",
|
|
"trace_type": "chat",
|
|
"window_start": 0.0,
|
|
"window_end": 600.0,
|
|
},
|
|
root / "private",
|
|
)
|
|
assert windows.is_file()
|
|
assert all(stats[role]["rows"] == 1 for role in prepare.ROLES)
|
|
for role in prepare.ROLES:
|
|
row = json.loads((root / "private" / "traces" / f"{role}.jsonl").read_text())
|
|
assert row["fidelity_pilot_band"] == role
|
|
assert abs(float(row["sampling_u"]) - 0.5) < 1e-12
|
|
|
|
prior = root / "prior-state.json"
|
|
primary = root / "primary-state.json"
|
|
prior.write_text(
|
|
json.dumps(
|
|
{
|
|
"status": "failed",
|
|
"gpu_hours_total": 0.02,
|
|
"hard_cap_h20_hours": 3.5,
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
primary.write_text(
|
|
json.dumps(
|
|
{
|
|
"status": "complete",
|
|
"gpu_hours_total": 1.5,
|
|
"hard_cap_h20_hours": 3.5,
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
accounting = campaign_gpu_accounting(primary, (prior,))
|
|
assert math.isclose(accounting["aggregate_h20_hours"], 1.52)
|
|
assert all(accounting["invariants"].values())
|
|
|
|
assert len(controller.ORDER) == 6
|
|
assert set(controller.ORDER) == set(prepare.CELLS)
|
|
assert math.isclose(
|
|
sum(
|
|
controller.CELL_ESTIMATE_H20_HOURS[int(config["tp"])]
|
|
for config in prepare.CELLS.values()
|
|
) + controller.SAFETY_H20_HOURS,
|
|
3.0,
|
|
)
|
|
selection = {
|
|
"selected_count": 122,
|
|
"request_id_order_sha256": "request-hash",
|
|
"arrival_order_sha256": "arrival-hash",
|
|
"input_length_order_sha256": "length-hash",
|
|
}
|
|
warmup = {
|
|
"kind": "warmup",
|
|
"selection": {"count": 16},
|
|
"invariants": {
|
|
"warmup_16": True,
|
|
"warmup_exact_16": True,
|
|
"warmup_long": True,
|
|
},
|
|
}
|
|
controller.validate_result_selection(
|
|
result=warmup,
|
|
selection=selection,
|
|
cell="tp1_mns8",
|
|
role="burnin",
|
|
warmup=True,
|
|
)
|
|
measured = {
|
|
"kind": "anchor",
|
|
"selection": {
|
|
"count": 122,
|
|
"request_id_order_sha256": "request-hash",
|
|
"arrival_order_sha256": "arrival-hash",
|
|
"raw_length_order_sha256": "length-hash",
|
|
},
|
|
"invariants": {},
|
|
}
|
|
controller.validate_result_selection(
|
|
result=measured,
|
|
selection=selection,
|
|
cell="tp1_mns8",
|
|
role="low1",
|
|
warmup=False,
|
|
)
|
|
try:
|
|
controller.validate_result_selection(
|
|
result=warmup,
|
|
selection=selection,
|
|
cell="tp1_mns8",
|
|
role="low1",
|
|
warmup=False,
|
|
)
|
|
except RuntimeError as error:
|
|
assert "selection count mismatch" in str(error)
|
|
else:
|
|
raise AssertionError("measured selection accepted a warmup subset")
|
|
print("fidelity pilot tools: PASS")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|