Add action-conditioned intervention feasibility model
This commit is contained in:
196
runs/active-intervention-v0/test_pipeline.py
Normal file
196
runs/active-intervention-v0/test_pipeline.py
Normal file
@@ -0,0 +1,196 @@
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
def load(name: str, path: Path):
|
||||
spec = importlib.util.spec_from_file_location(name, path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec.loader is not None
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def write_json(path: Path, payload) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(payload) + "\n", encoding="utf-8")
|
||||
|
||||
|
||||
def write_jsonl(path: Path, rows) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(
|
||||
"".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8"
|
||||
)
|
||||
|
||||
|
||||
def engine_record(index: int, timestamp_ns: int) -> dict:
|
||||
alternate = index % 2
|
||||
return {
|
||||
"step_index": index,
|
||||
"submit_mono_ns": timestamp_ns,
|
||||
"model_executed": True,
|
||||
"scheduled_requests": 1 + alternate,
|
||||
"decode_batch_size": alternate,
|
||||
"prefill_tokens": 8 + alternate,
|
||||
"decode_tokens": alternate,
|
||||
"preemptions": 0,
|
||||
"queues": {"waiting": alternate, "running": 1 + alternate},
|
||||
"kv": {"usage": 0.1 + 0.01 * alternate},
|
||||
"cudagraph": {
|
||||
"runtime_mode": "FULL" if alternate else "NONE",
|
||||
"bucket_tokens": 16,
|
||||
"padding_tokens": alternate,
|
||||
},
|
||||
"dropped_records_before": 0,
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
extractor = load("active_intervention_extract_test", HERE / "extract_training.py")
|
||||
trainer = load("active_intervention_train_test", HERE / "train_policy.py")
|
||||
with tempfile.TemporaryDirectory() as temporary:
|
||||
root = Path(temporary)
|
||||
run_root = root / "runs"
|
||||
configs = [
|
||||
{"id": "a_base", "mns": 16, "mbbt": 8192},
|
||||
{"id": "a_mns", "mns": 64, "mbbt": 8192},
|
||||
{"id": "a_mbbt", "mns": 16, "mbbt": 16384},
|
||||
{"id": "b_base", "mns": 64, "mbbt": 2048},
|
||||
{"id": "b_mns", "mns": 128, "mbbt": 2048},
|
||||
{"id": "b_mbbt", "mns": 64, "mbbt": 8192},
|
||||
]
|
||||
manifest = {
|
||||
"engine": {"duration_s": 300.0, "tp": 4},
|
||||
"configs": configs,
|
||||
"repetitions": {
|
||||
str(rep): {"selection": {"offered_req_s_per_gpu": 0.01}}
|
||||
for rep in (1, 2, 3)
|
||||
},
|
||||
"regimes": {
|
||||
"A": {
|
||||
"source": "a_base",
|
||||
"actions": {"mns": "a_mns", "mbbt": "a_mbbt"},
|
||||
},
|
||||
"B": {
|
||||
"source": "b_base",
|
||||
"actions": {"mns": "b_mns", "mbbt": "b_mbbt"},
|
||||
},
|
||||
},
|
||||
}
|
||||
manifest_path = root / "manifest.json"
|
||||
write_json(manifest_path, manifest)
|
||||
|
||||
streams = []
|
||||
source_starts: dict[tuple[str, int], int] = {}
|
||||
for source_index, source_id in enumerate(("a_base", "b_base")):
|
||||
rows = []
|
||||
index = 0
|
||||
for repetition in (1, 2, 3):
|
||||
start_ns = int((source_index * 2000 + repetition * 400) * 1e9)
|
||||
source_starts[(source_id, repetition)] = start_ns
|
||||
for second in (1, 30, 76, 105, 151, 180, 226, 255):
|
||||
rows.append(engine_record(index, start_ns + int(second * 1e9)))
|
||||
index += 1
|
||||
stream_path = root / f"{source_id}-stream.jsonl"
|
||||
write_jsonl(stream_path, rows)
|
||||
streams.append(
|
||||
{
|
||||
"config_id": source_id,
|
||||
"stream": str(stream_path),
|
||||
"stream_sha256": extractor.sha256_file(stream_path),
|
||||
}
|
||||
)
|
||||
|
||||
request_rows = [
|
||||
{
|
||||
"request_id": f"r{index}",
|
||||
"arrival_s": arrival,
|
||||
"completed_elapsed_s": arrival + 10,
|
||||
"slo_pass": index != 3,
|
||||
"ttft_ms": 1000 + index * 100,
|
||||
"tpot_ms": 20 + index,
|
||||
"raw_input_tokens": 1000 + index * 100,
|
||||
}
|
||||
for index, arrival in enumerate((5.0, 80.0, 155.0, 230.0), 1)
|
||||
]
|
||||
for source_id in ("a_base", "b_base"):
|
||||
for repetition in (1, 2, 3):
|
||||
write_jsonl(
|
||||
run_root
|
||||
/ "sessions"
|
||||
/ source_id
|
||||
/ f"rep{repetition}"
|
||||
/ "requests.jsonl",
|
||||
request_rows,
|
||||
)
|
||||
|
||||
goodput = {
|
||||
"a_base": 0.020,
|
||||
"a_mns": 0.036,
|
||||
"a_mbbt": 0.028,
|
||||
"b_base": 0.032,
|
||||
"b_mns": 0.030,
|
||||
"b_mbbt": 0.038,
|
||||
}
|
||||
runs = []
|
||||
for config in configs:
|
||||
for repetition in (1, 2, 3):
|
||||
item = {
|
||||
"config_id": config["id"],
|
||||
"repetition": repetition,
|
||||
"outcome": {
|
||||
"slo_goodput_req_s": goodput[config["id"]]
|
||||
+ repetition * 0.0001
|
||||
},
|
||||
}
|
||||
if config["id"] in ("a_base", "b_base"):
|
||||
start_ns = source_starts[(config["id"], repetition)]
|
||||
item["state"] = {
|
||||
"interval": {
|
||||
"start_ns": start_ns,
|
||||
"end_ns": start_ns + int(300 * 1e9),
|
||||
}
|
||||
}
|
||||
runs.append(item)
|
||||
audit = {
|
||||
"schema": "action-aware-constraint-pilot-audit-v0",
|
||||
"sanity": {"red_flags": []},
|
||||
"streams": streams,
|
||||
"runs": runs,
|
||||
}
|
||||
audit_path = root / "audit.json"
|
||||
write_json(audit_path, audit)
|
||||
|
||||
dataset = extractor.build_dataset(
|
||||
audit_path=audit_path, manifest_path=manifest_path, run_root=run_root
|
||||
)
|
||||
assert dataset["status"] == "VALID"
|
||||
assert len(dataset["examples"]) == 72
|
||||
assert not dataset["sanity"]["red_flags"]
|
||||
assert all(
|
||||
"exclusive" not in feature
|
||||
for example in dataset["examples"]
|
||||
for feature in example["source"]["telemetry"]
|
||||
)
|
||||
dataset_path = root / "dataset.json"
|
||||
write_json(dataset_path, dataset)
|
||||
policy = trainer.build_policy(dataset_path)
|
||||
assert policy["status"] in {
|
||||
"RETROSPECTIVE_INCREMENTAL_SIGNAL",
|
||||
"NO_RETROSPECTIVE_INCREMENTAL_SIGNAL",
|
||||
}
|
||||
assert not policy["sanity"]["red_flags"]
|
||||
print("active intervention pipeline: PASS")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user