Add fixed-cohort Frontier rank protocol

runs/frontier-multicase-sufficiency-v0/best_effort/test_fixed_cohort_rank.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+import argparse
+import csv
+import importlib.util
+import json
+import sys
+from pathlib import Path
+
+
+SCRIPT = Path(__file__).with_name("fixed_cohort_rank.py")
+SPEC = importlib.util.spec_from_file_location("fixed_cohort_rank", 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 _write_trace(path: Path) -> None:
+    lengths = [100, 1500, 3000, 6000, 12000, 24000]
+    rows = []
+    for index in range(120):
+        rows.append(
+            {
+                "timestamp": index * 0.5,
+                "sampling_u": (index + 1) / 121,
+                "input_length": lengths[index % len(lengths)],
+                "output_length": 10,
+                "prompt": f"raw prompt {index}",
+            }
+        )
+    path.write_text("".join(json.dumps(row) + "\n" for row in rows))
+
+
+def test_prepare_protocol_keeps_identical_cohort_and_scales_only_arrivals(
+    tmp_path: Path,
+) -> None:
+    trace = tmp_path / "trace.jsonl"
+    _write_trace(trace)
+    manifest_path = MODULE.prepare_protocol(
+        argparse.Namespace(
+            trace=trace,
+            expected_trace_sha256="",
+            output_root=tmp_path / "protocol",
+            cohort_size=24,
+            seed=7,
+            rate=[0.2, 0.4],
+        )
+    )
+    manifest = json.loads(manifest_path.read_text())
+    assert manifest["selection"]["cohort_bin_quotas"] == [4, 4, 4, 4, 4, 4]
+    assert manifest["load_contract"]["binary_search"] is False
+    assert manifest["load_contract"]["monotonicity_assumed"] is False
+
+    materialized = []
+    for record in manifest["rates"].values():
+        with Path(record["path"]).open(newline="") as source:
+            rows = list(csv.DictReader(source))
+        materialized.append(rows)
+        duration = float(rows[-1]["arrived_at"]) - float(rows[0]["arrived_at"])
+        assert len(rows) / duration == record["offered_request_rate"]
+    assert [row["source_row_index"] for row in materialized[0]] == [
+        row["source_row_index"] for row in materialized[1]
+    ]
+    assert [row["num_prefill_tokens"] for row in materialized[0]] == [
+        row["num_prefill_tokens"] for row in materialized[1]
+    ]
+    assert MODULE.audit_protocol(manifest_path)["status"] == "passed"
+
+
+def test_score_requests_applies_all_slos_to_same_outcomes() -> None:
+    requests = [
+        {"input_tokens": 100, "success": True, "ttft_ms": 900.0},
+        {"input_tokens": 16000, "success": True, "ttft_ms": 2500.0},
+        {"input_tokens": 24000, "success": False, "ttft_ms": None},
+    ]
+    score = MODULE.score_requests(requests)
+    assert score["scores"]["linear_8k_primary"]["passed_request_count"] == 2
+    assert score["scores"]["legacy_step_1s_2s"]["passed_request_count"] == 1
+    assert score["ttft_ms"]["count"] == 2
+
+
+def test_rank_surface_records_nonmonotone_feasibility_without_binary_search() -> None:
+    config_results = [
+        {
+            "config": {"name": "a", "tp": 4},
+            "loads": [
+                {
+                    "offered_request_rate": 0.2,
+                    "scores": {"linear_8k_primary": {"feasible": False, "slo_pass_rate": 0.9}},
+                },
+                {
+                    "offered_request_rate": 0.4,
+                    "scores": {"linear_8k_primary": {"feasible": True, "slo_pass_rate": 0.95}},
+                },
+            ],
+        },
+        {
+            "config": {"name": "b", "tp": 8},
+            "loads": [
+                {
+                    "offered_request_rate": 0.2,
+                    "scores": {"linear_8k_primary": {"feasible": True, "slo_pass_rate": 1.0}},
+                },
+                {
+                    "offered_request_rate": 0.4,
+                    "scores": {"linear_8k_primary": {"feasible": False, "slo_pass_rate": 0.9}},
+                },
+            ],
+        },
+    ]
+    ranking = MODULE.rank_surface(config_results, slo_name="linear_8k_primary")
+    a = next(item for item in ranking if item["config"]["name"] == "a")
+    assert a["maximum_tested_feasible_request_rate"] == 0.4
+    assert a["monotonicity_violations"] == [[0.2, 0.4]]
+
+
+def test_average_ranks_use_midrank_for_capacity_ties() -> None:
+    assert MODULE._average_ranks({"a": 0.2, "b": 0.2, "c": 0.1}) == {
+        "a": 1.5,
+        "b": 1.5,
+        "c": 3.0,
+    }
+
+
+def test_community_study_uses_affine_primary_slo(tmp_path: Path) -> None:
+    payload = MODULE._community_study_payload(
+        tp=8,
+        repo=tmp_path,
+        python=tmp_path / "python",
+        vllm=tmp_path / "vllm",
+        model=tmp_path / "model",
+        trace=tmp_path / "trace.jsonl",
+        windows=tmp_path / "windows.json",
+        port=18918,
+    )
+    assert payload["slo"]["ttft_rule"] == {
+        "kind": "linear_ms",
+        "intercept_ms": 1000.0,
+        "per_token_ms": 0.125,
+    }
+    assert payload["trace"]["request_mode"] == "raw_completion"
+    assert payload["trace"]["completion_tokens_override"] == 1