86 lines
2.4 KiB
Python
86 lines
2.4 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
|
|
|
|
|
|
@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
|
|
|
|
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,
|
|
)
|
|
print("fidelity pilot tools: PASS")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|