Files
aituner/runs/intervention-response-v2/test_analysis.py

201 lines
6.3 KiB
Python

#!/usr/bin/env python3
from __future__ import annotations
import importlib.util
import math
from pathlib import Path
from types import SimpleNamespace
HERE = Path(__file__).resolve().parent
def load_module():
spec = importlib.util.spec_from_file_location(
"intervention_response_phase_aware_v2", HERE / "analyze_existing.py"
)
module = importlib.util.module_from_spec(spec)
assert spec.loader is not None
spec.loader.exec_module(module)
return module
def load_prepare_module():
spec = importlib.util.spec_from_file_location(
"intervention_response_phase_aware_prepare", HERE / "prepare_pilot.py"
)
module = importlib.util.module_from_spec(spec)
assert spec.loader is not None
spec.loader.exec_module(module)
return module
def load_controller_module():
spec = importlib.util.spec_from_file_location(
"intervention_response_phase_aware_controller", HERE / "pilot_controller.py"
)
module = importlib.util.module_from_spec(spec)
assert spec.loader is not None
spec.loader.exec_module(module)
return module
def load_pilot_analysis_module():
spec = importlib.util.spec_from_file_location(
"intervention_response_phase_aware_pilot_analysis", HERE / "analyze_pilot.py"
)
module = importlib.util.module_from_spec(spec)
assert spec.loader is not None
spec.loader.exec_module(module)
return module
def main() -> None:
module = load_module()
assert module.common_decile_fractions(
trace_duration_s=60.0, minimum_elapsed_s=19.448
) == (0.1, 0.2, 0.3)
assert module.common_decile_fractions(
trace_duration_s=60.0, minimum_elapsed_s=60.0
)[-1] == 1.0
stats = module.numeric([0.0, 1.0, 2.0])
assert stats == {
"n": 3,
"min": 0.0,
"max": 2.0,
"distinct_n": 3,
"median": 1.0,
}
assert math.isclose(module._pearson([1.0, 2.0], [2.0, 4.0]), 1.0)
assert module._pearson([1.0, 1.0], [2.0, 3.0]) is None
prepare = load_prepare_module()
requests = [
SimpleNamespace(
sampling_u=index / 10.0,
row_id=f"r{index}",
arrival_s=float(index),
prompt_tokens_hint=100 + index,
)
for index in range(1, 6)
]
_anchor, selected = prepare.attainable_anchor(requests, 3)
assert len(selected) == 3
record = prepare.selection_record(selected, duration_s=3.0)
assert record["selected_count"] == 3
assert record["offered_req_s_per_gpu"] == 0.25
assert len(prepare.SESSION_ORDER) == 6
assert {mns for _replicate, mns in prepare.SESSION_ORDER} == {16, 64}
first_id = prepare.private_request_id(
source_sha256="a" * 64, line_number=1, original_id="1"
)
assert first_id == prepare.private_request_id(
source_sha256="a" * 64, line_number=1, original_id="1"
)
assert first_id != prepare.private_request_id(
source_sha256="b" * 64, line_number=1, original_id="1"
)
controller = load_controller_module()
assert math.isclose(controller.remaining_projection(6, 0), 7.7)
assert math.isclose(controller.remaining_projection(6, 5), 1.45)
parsed = controller.parser().parse_args(
[
"--manifest",
"/tmp/manifest.json",
"--run-root",
"/tmp/run",
"--aituner-root",
"/tmp/aituner",
"--vllm-source",
"/tmp/vllm",
"--venv",
"/tmp/venv",
"--model",
"/tmp/model",
"--client",
"/tmp/client.py",
"--dry-run",
]
)
assert parsed.dry_run is True
pilot_analysis = load_pilot_analysis_module()
stable = pilot_analysis.stable_adjacent_features(
[
{"end_fraction": 0.1, "qualifying_response_features": ["queue"]},
{
"end_fraction": 0.25,
"qualifying_response_features": ["kv", "queue"],
},
{"end_fraction": 0.5, "qualifying_response_features": ["queue"]},
]
)
assert stable == {"0.10->0.25": ["queue"], "0.25->0.50": ["queue"]}
load_consistency = {
"0.10->0.25:queue": {"passes_two_regimes": True},
"0.25->0.50:queue": {"passes_two_regimes": True},
}
mechanism = pilot_analysis.mechanism_gate(stable, load_consistency)
assert mechanism["passes"] is False
stable["0.25->0.50"].append("kv")
load_consistency["0.25->0.50:kv"] = {"passes_two_regimes": True}
mechanism = pilot_analysis.mechanism_gate(stable, load_consistency)
assert mechanism["passes"] is True
assert mechanism["passing_transitions"] == ["0.25->0.50"]
efficacy = pilot_analysis.stable_adjacent_efficacy_features(
[
{
"end_fraction": 0.1,
"efficacy": {"telemetry_qualifying_features": ["early"]},
},
{
"end_fraction": 0.25,
"efficacy": {"telemetry_qualifying_features": ["queue"]},
},
{
"end_fraction": 0.5,
"efficacy": {"telemetry_qualifying_features": ["kv", "queue"]},
},
]
)
assert efficacy == {"0.25->0.50": ["queue"]}
coverage = pilot_analysis.telemetry_coverage(
[
{"step_index": 1, "submit_mono_ns": 100_000_000},
{"step_index": 2, "submit_mono_ns": 200_000_000},
],
start_ns=0,
end_ns=300_000_000,
)
assert coverage == {
"start_gap_s": 0.1,
"end_gap_s": 0.1,
"max_internal_gap_s": 0.1,
}
coverage_gate = pilot_analysis.cumulative_coverage_gate(
[
{
"trial_sanity": [
{
"trial_id": "a",
"admitted_fraction": 0.25,
"completed_fraction": 0.2,
}
]
},
{
"trial_sanity": [
{
"trial_id": "a",
"admitted_fraction": 0.5,
"completed_fraction": 0.4,
}
]
},
]
)
assert coverage_gate["red_flags"] == []
print("phase-aware intervention response v2 analysis: PASS")
if __name__ == "__main__":
main()