111 lines
3.2 KiB
Python
111 lines
3.2 KiB
Python
import importlib.util
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
SCRIPT = (
|
|
Path(__file__).resolve().parents[1]
|
|
/ "runs/frontier-workload-regime-taxonomy-v0/materialize_workload_regimes.py"
|
|
)
|
|
SPEC = importlib.util.spec_from_file_location("workload_regimes", SCRIPT)
|
|
MODULE = importlib.util.module_from_spec(SPEC)
|
|
assert SPEC.loader is not None
|
|
sys.modules[SPEC.name] = MODULE
|
|
SPEC.loader.exec_module(MODULE)
|
|
|
|
|
|
def source_rows():
|
|
return [
|
|
{
|
|
"source_index": index,
|
|
"source_arrived_at": arrival,
|
|
"input_length": input_length,
|
|
"output_length": output_length,
|
|
"session_id": 100 + index,
|
|
"runtime_block_ids": list(range(input_length // 16)),
|
|
"body": {"prompt": f"prompt-{index}", "model": "fixture"},
|
|
}
|
|
for index, (arrival, input_length, output_length) in enumerate(
|
|
[(10.0, 16, 4), (10.5, 32, 8), (14.0, 48, 12)]
|
|
)
|
|
]
|
|
|
|
|
|
def test_arrival_contracts_preserve_rate_and_burst_pattern():
|
|
rows = source_rows()
|
|
uniform = MODULE.arrivals_for(rows, "uniform", 2.0)
|
|
traced = MODULE.arrivals_for(rows, "trace", 2.0)
|
|
|
|
assert uniform == [0.0, 0.5, 1.0]
|
|
assert (len(traced) - 1) / (traced[-1] - traced[0]) == 2.0
|
|
assert traced[1] - traced[0] < traced[2] - traced[1]
|
|
|
|
|
|
def test_family_axes_change_only_the_declared_contract():
|
|
rows = source_rows()
|
|
token_ids = [10, 11, 12, 13]
|
|
families = {family.name: family for family in MODULE.FAMILIES}
|
|
|
|
w1 = MODULE.build_family_rows(
|
|
rows,
|
|
families["w1-mean-fixed-uniform-none"],
|
|
target_rate=1.0,
|
|
fixed_token_ids=token_ids,
|
|
)
|
|
w3 = MODULE.build_family_rows(
|
|
rows,
|
|
families["w3-heterogeneous-uniform-none"],
|
|
target_rate=1.0,
|
|
fixed_token_ids=token_ids,
|
|
)
|
|
w5 = MODULE.build_family_rows(
|
|
rows,
|
|
families["w5-heterogeneous-uniform-prefix"],
|
|
target_rate=1.0,
|
|
fixed_token_ids=token_ids,
|
|
)
|
|
|
|
assert {(row["input_length"], row["output_length"]) for row in w1} == {(32, 8)}
|
|
assert [(row["input_length"], row["output_length"]) for row in w3] == [
|
|
(16, 4),
|
|
(32, 8),
|
|
(48, 12),
|
|
]
|
|
assert all(row["runtime_block_ids"] == [] for row in w1 + w3)
|
|
assert [row["runtime_block_ids"] for row in w5] == [
|
|
[0],
|
|
[0, 1],
|
|
[0, 1, 2],
|
|
]
|
|
assert [row["body"]["prompt"] for row in w3] == [
|
|
"prompt-0",
|
|
"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
|
|
|
|
|
|
def test_rho_labels_do_not_collide_below_one_percent():
|
|
assert MODULE.rho_label(0.0025) == "rho0p0025"
|
|
assert MODULE.rho_label(0.005) == "rho0p005"
|
|
assert MODULE.rho_label(0.01) == "rho0p01"
|