Compare commits
3 Commits
8b4116fad0
...
51a9e4a007
| Author | SHA1 | Date | |
|---|---|---|---|
| 51a9e4a007 | |||
| 0f15bbc3f1 | |||
| 6f8e3c95c1 |
@@ -92,17 +92,39 @@ def parse_args() -> argparse.Namespace:
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def stable_uniform(*, seed: int, window_id: str, index: int, row: dict[str, Any]) -> float:
|
||||
def resolve_session_root(row: dict[str, Any], root_of: dict[Any, Any]) -> Any:
|
||||
"""Resolve the session root chat_id for a trace row.
|
||||
|
||||
Sessions are multi-turn chains linked via parent_chat_id (turn>1 points to the
|
||||
parent turn's chat_id, the root turn has parent_chat_id=-1). Because parent
|
||||
turns precede their children in time, a single streaming pass that records
|
||||
chat_id -> root resolves the full chain. Rows whose parent is not yet known
|
||||
(e.g. it fell outside the materialized span) fall back to the parent id so
|
||||
siblings still group together.
|
||||
"""
|
||||
chat_id = row.get("chat_id")
|
||||
parent = row.get("parent_chat_id")
|
||||
parent_is_root = (
|
||||
parent is None
|
||||
or (isinstance(parent, (int, float)) and not isinstance(parent, bool) and int(parent) < 0)
|
||||
)
|
||||
root = chat_id if parent_is_root else root_of.get(parent, parent)
|
||||
if chat_id is not None:
|
||||
root_of[chat_id] = root
|
||||
return root
|
||||
|
||||
|
||||
def session_uniform(*, seed: int, window_id: str, session_root: Any) -> float:
|
||||
"""Deterministic per-session uniform score in [0, 1).
|
||||
|
||||
All turns of a session share one score, so thresholding sampling_u keeps or
|
||||
drops whole sessions and preserves intra-session prefix (KV-cache) reuse.
|
||||
"""
|
||||
payload = json.dumps(
|
||||
{
|
||||
"seed": seed,
|
||||
"window_id": window_id,
|
||||
"index": index,
|
||||
"timestamp": row.get("timestamp"),
|
||||
"input_length": row.get("input_length"),
|
||||
"output_length": row.get("output_length"),
|
||||
"chat_id": row.get("chat_id"),
|
||||
"turn": row.get("turn"),
|
||||
"session_root": session_root,
|
||||
},
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
@@ -241,12 +263,16 @@ def materialize_windows(
|
||||
bucket = grouped[(trace_path, prompt_path)]
|
||||
bucket.sort(key=lambda item: (float(item["window_start"]), str(item["window_id"])))
|
||||
matched_rows = 0
|
||||
root_of: dict[Any, Any] = {}
|
||||
with trace_path.open() as trace_handle, prompt_path.open() as prompt_handle:
|
||||
for trace_raw, prompt_raw in zip(trace_handle, prompt_handle):
|
||||
trace_raw = trace_raw.strip()
|
||||
if not trace_raw:
|
||||
continue
|
||||
trace_row = json.loads(trace_raw)
|
||||
# Resolve session linkage for every row (even unmatched ones)
|
||||
# so multi-turn chains crossing the window edge still group.
|
||||
session_root = resolve_session_root(trace_row, root_of)
|
||||
timestamp = float(trace_row.get("timestamp") or 0.0)
|
||||
matched_window: dict[str, Any] | None = None
|
||||
for window in bucket:
|
||||
@@ -267,11 +293,11 @@ def materialize_windows(
|
||||
start = float(matched_window["window_start"])
|
||||
out["source_timestamp"] = timestamp
|
||||
out["timestamp"] = timestamp - start
|
||||
out["sampling_u"] = stable_uniform(
|
||||
out["session_root"] = session_root
|
||||
out["sampling_u"] = session_uniform(
|
||||
seed=sample_seed,
|
||||
window_id=window_id,
|
||||
index=stats_by_window[window_id].num_requests,
|
||||
row=merged,
|
||||
session_root=session_root,
|
||||
)
|
||||
handles[window_id].write(json.dumps(out, ensure_ascii=False) + "\n")
|
||||
stats_by_window[window_id].record(out)
|
||||
@@ -311,7 +337,7 @@ def build_output_window(
|
||||
output["num_excluded_too_long"] = 0
|
||||
output["sampling_u_field"] = "sampling_u"
|
||||
output["sampling_seed"] = int(sample_seed)
|
||||
output["sampling_strategy"] = "fixed_uniform_score"
|
||||
output["sampling_strategy"] = "session_coherent_uniform_score"
|
||||
output["first_request_ts"] = stats.first_request_ts
|
||||
output["last_request_ts"] = stats.last_request_ts
|
||||
output["first_request_index"] = stats.first_request_index
|
||||
|
||||
@@ -10,6 +10,7 @@ from aituner.llm import (
|
||||
load_capability_profile,
|
||||
parse_proposal_text,
|
||||
)
|
||||
from aituner.lca import build_study_workload_profile
|
||||
from aituner.spec import load_study_spec
|
||||
from aituner.store import StudyStore
|
||||
from aituner.trace import load_trace_requests, summarize_window
|
||||
@@ -89,6 +90,7 @@ def main() -> int:
|
||||
window_summary=summarize_window(requests, window),
|
||||
state=state,
|
||||
capability_profile=capability_profile,
|
||||
workload_profile=build_study_workload_profile(study, requests, window),
|
||||
)
|
||||
prompt_name = f"prompt-{state.next_trial_index:04d}"
|
||||
store.write_prompt(study.study_id, prompt_name, prompt)
|
||||
|
||||
@@ -14,6 +14,7 @@ from .harness import (
|
||||
)
|
||||
from .job import append_job, build_trial_job
|
||||
from .lca import (
|
||||
build_study_workload_profile,
|
||||
build_workload_profile,
|
||||
resolve_length_mode,
|
||||
similarity_report,
|
||||
@@ -140,6 +141,7 @@ def cmd_study_prompt(args: argparse.Namespace) -> int:
|
||||
window_summary=summarize_window(requests, window),
|
||||
state=state,
|
||||
capability_profile=capability_profile,
|
||||
workload_profile=build_study_workload_profile(study, requests, window),
|
||||
)
|
||||
prompt_name = args.prompt_name or f"prompt-{state.next_trial_index:04d}"
|
||||
path = store.write_prompt(study.study_id, prompt_name, prompt)
|
||||
@@ -160,6 +162,7 @@ def cmd_study_llm_propose(args: argparse.Namespace) -> int:
|
||||
window_summary=summarize_window(requests, window),
|
||||
state=state,
|
||||
capability_profile=capability_profile,
|
||||
workload_profile=build_study_workload_profile(study, requests, window),
|
||||
)
|
||||
proposal_text = call_llm_for_proposal(
|
||||
policy=study.llm,
|
||||
@@ -242,11 +245,13 @@ def cmd_study_tune(args: argparse.Namespace) -> int:
|
||||
break
|
||||
window, requests = load_trace_requests(study, study_spec_path=spec_path)
|
||||
window_summary = summarize_window(requests, window)
|
||||
workload_profile = build_study_workload_profile(study, requests, window)
|
||||
harness_context = (
|
||||
build_harness_context(
|
||||
study=study,
|
||||
window_summary=window_summary,
|
||||
state=state,
|
||||
workload_profile=workload_profile,
|
||||
)
|
||||
if study.llm.use_harness
|
||||
else None
|
||||
@@ -256,6 +261,7 @@ def cmd_study_tune(args: argparse.Namespace) -> int:
|
||||
window_summary=window_summary,
|
||||
state=state,
|
||||
capability_profile=capability_profile,
|
||||
workload_profile=workload_profile,
|
||||
)
|
||||
prompt_name = f"prompt-{state.next_trial_index:04d}"
|
||||
store.write_prompt(study.study_id, prompt_name, prompt)
|
||||
|
||||
@@ -4,6 +4,7 @@ import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from .lca import EPSILON, WorkloadProfile
|
||||
from .spec import ConfigPatch, Proposal, StudySpec, StudyState, TrialSummary
|
||||
|
||||
|
||||
@@ -30,6 +31,7 @@ def build_harness_context(
|
||||
study: StudySpec,
|
||||
window_summary: dict[str, Any],
|
||||
state: StudyState,
|
||||
workload_profile: WorkloadProfile | None = None,
|
||||
) -> dict[str, Any]:
|
||||
recent_diagnostics = _recent_trial_diagnostics(state)
|
||||
trial_profiles = _trial_profiles(study, recent_diagnostics)
|
||||
@@ -52,7 +54,7 @@ def build_harness_context(
|
||||
"feature_model": "L-C-A: request lengths, inter-request KV-cache reuse, and arrival dynamics.",
|
||||
"trial_policy": "Profile measured trials, rank bottleneck hypotheses, score generic candidate actions, and stop only when no useful measured hypothesis remains.",
|
||||
},
|
||||
"workload_lca_profile": _workload_lca_profile(window_summary),
|
||||
"workload_lca_profile": _workload_lca_profile(window_summary, workload_profile),
|
||||
"recent_trial_diagnostics": recent_diagnostics,
|
||||
"trial_profiles": trial_profiles,
|
||||
"bottleneck_hypotheses": bottleneck_hypotheses,
|
||||
@@ -141,7 +143,12 @@ def render_harness_context(context: dict[str, Any]) -> str:
|
||||
return json.dumps(context, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
def _workload_lca_profile(window_summary: dict[str, Any]) -> dict[str, Any]:
|
||||
def _workload_lca_profile(
|
||||
window_summary: dict[str, Any],
|
||||
workload_profile: WorkloadProfile | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if workload_profile is not None:
|
||||
return _canonical_lca_profile(workload_profile)
|
||||
prefix_cache = window_summary.get("prefix_cache")
|
||||
if not isinstance(prefix_cache, dict):
|
||||
prefix_cache = {}
|
||||
@@ -178,6 +185,54 @@ def _workload_lca_profile(window_summary: dict[str, Any]) -> dict[str, Any]:
|
||||
}
|
||||
|
||||
|
||||
def _canonical_lca_profile(profile: WorkloadProfile) -> dict[str, Any]:
|
||||
"""Authoritative L-C-A block: the paper's 10-dim RobustScaler vector.
|
||||
|
||||
Sourced from lca.WorkloadProfile so the prompt's L-C-A is the same metric
|
||||
used for the workload-similarity computations, not an ad-hoc re-derivation.
|
||||
The regime labels reuse the heuristic classifiers but are fed from the
|
||||
canonical stats.
|
||||
"""
|
||||
stats = profile.stats if isinstance(profile.stats, dict) else {}
|
||||
length = stats.get("length") if isinstance(stats.get("length"), dict) else {}
|
||||
cache = stats.get("cache") if isinstance(stats.get("cache"), dict) else {}
|
||||
arrival = stats.get("arrival") if isinstance(stats.get("arrival"), dict) else {}
|
||||
length_p95 = _as_float(length.get("p95"))
|
||||
length_p50 = _as_float(length.get("p50"))
|
||||
tail_ratio = float(length_p95 / max(length_p50, EPSILON)) if length_p95 else 0.0
|
||||
repeated_token_ratio = _as_float(cache.get("input_hit_rate"))
|
||||
fano_1s = _as_float(arrival.get("fano_1s"))
|
||||
interarrival_cv = _as_float(arrival.get("interarrival_cv"))
|
||||
return {
|
||||
"metric": "paper L-C-A (10-dim, RobustScaler-normalized) from lca.WorkloadProfile",
|
||||
"length_mode": profile.length_mode,
|
||||
"feature_names": list(profile.feature_names),
|
||||
"vector": list(profile.vector),
|
||||
"L_request_lengths": {
|
||||
"mean": _as_float(length.get("mean")),
|
||||
"p50": length_p50,
|
||||
"p95": length_p95,
|
||||
"cv": _as_float(length.get("cv")),
|
||||
"tail_ratio_p95_p50": tail_ratio,
|
||||
"regime": _length_regime(length_p95, tail_ratio),
|
||||
},
|
||||
"C_prefix_cache": {
|
||||
"hit_rate": _as_float(cache.get("hit_rate")),
|
||||
"input_hit_rate": repeated_token_ratio,
|
||||
"repeated_block_ratio": _as_float(cache.get("repeated_block_ratio")),
|
||||
"rows_with_hash_ids": int(cache.get("rows_with_hash_ids") or 0),
|
||||
"regime": _cache_regime(repeated_token_ratio),
|
||||
},
|
||||
"A_arrivals": {
|
||||
"request_rate": _as_float(arrival.get("request_rate")),
|
||||
"request_rate_per_gpu": _as_float(arrival.get("request_rate_per_gpu")),
|
||||
"interarrival_cv": interarrival_cv,
|
||||
"fano_1s": fano_1s,
|
||||
"regime": _arrival_regime(fano_1s, interarrival_cv),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _knob_harnesses(
|
||||
study: StudySpec,
|
||||
window_summary: dict[str, Any],
|
||||
|
||||
@@ -4,10 +4,13 @@ import json
|
||||
import math
|
||||
import statistics
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Sequence
|
||||
|
||||
from .trace import TraceRequest, WindowRecord
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .spec import StudySpec
|
||||
|
||||
|
||||
EPSILON = 1e-9
|
||||
|
||||
@@ -178,6 +181,28 @@ def build_workload_profile(
|
||||
)
|
||||
|
||||
|
||||
def build_study_workload_profile(
|
||||
study: "StudySpec",
|
||||
requests: list[TraceRequest],
|
||||
window: WindowRecord,
|
||||
) -> WorkloadProfile:
|
||||
"""Canonical L-C-A profile for a study's loaded window.
|
||||
|
||||
This is the single source of truth for the paper's 10-dimensional L-C-A
|
||||
feature vector used by the harness prompt and (later) by Stop-A.
|
||||
"""
|
||||
mode = resolve_length_mode(
|
||||
request_mode=study.trace.request_mode,
|
||||
length_mode="auto",
|
||||
)
|
||||
return build_workload_profile(
|
||||
requests,
|
||||
window,
|
||||
gpu_count=study.hardware.gpu_count,
|
||||
length_mode=mode,
|
||||
)
|
||||
|
||||
|
||||
def fit_robust_scale(profiles: Sequence[WorkloadProfile]) -> RobustScale:
|
||||
if not profiles:
|
||||
raise ValueError("At least one profile is required to fit a robust scale.")
|
||||
@@ -234,6 +259,151 @@ def similarity_report(profiles: Sequence[WorkloadProfile]) -> dict[str, Any]:
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ConvergencePoint:
|
||||
converged: bool
|
||||
stop_index: int
|
||||
stop_time_s: float
|
||||
fraction: float
|
||||
family_similarity: dict[str, float]
|
||||
checks: list[dict[str, Any]]
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"converged": self.converged,
|
||||
"stop_index": self.stop_index,
|
||||
"stop_time_s": self.stop_time_s,
|
||||
"fraction": self.fraction,
|
||||
"family_similarity": self.family_similarity,
|
||||
"checks": self.checks,
|
||||
}
|
||||
|
||||
|
||||
def find_convergence_prefix(
|
||||
requests: list[TraceRequest],
|
||||
window: WindowRecord,
|
||||
*,
|
||||
gpu_count: int,
|
||||
length_mode: str = "total",
|
||||
tau: float = 0.9,
|
||||
tau_c: float = 0.92,
|
||||
stable_checks: int = 3,
|
||||
max_checks: int = 20,
|
||||
min_fraction: float = 0.1,
|
||||
) -> ConvergencePoint:
|
||||
"""Earliest arrival-ordered prefix whose offered L-C-A converges to the full set.
|
||||
|
||||
The L-C-A vector is a deterministic function of the trace metadata, so the
|
||||
convergence of prefix-vs-full is itself deterministic (the paper's Fig. 9
|
||||
curve). Stop-A replays only up to this prefix. A prefix counts as converged
|
||||
when the L and A family similarities reach ``tau`` and the (slowest) C family
|
||||
similarity reaches the stricter ``tau_c`` for ``stable_checks`` consecutive
|
||||
checkpoints. If that never happens within the window the point reports the
|
||||
full set (converged=False), which keeps the C-gate honest: an unconverged C
|
||||
means the probe must replay the whole window rather than stop early.
|
||||
"""
|
||||
total = len(requests)
|
||||
if total == 0:
|
||||
return ConvergencePoint(
|
||||
converged=False,
|
||||
stop_index=0,
|
||||
stop_time_s=0.0,
|
||||
fraction=1.0,
|
||||
family_similarity={"L": 1.0, "C": 1.0, "A": 1.0},
|
||||
checks=[],
|
||||
)
|
||||
# Compare each arrival-ordered prefix to the whole set, both measured over
|
||||
# their own elapsed span so the A (rate) dimension is comparable rather than
|
||||
# diluted by the fixed window length.
|
||||
target = _prefix_profile(
|
||||
requests, total, window, gpu_count=gpu_count, length_mode=length_mode
|
||||
)
|
||||
indices = _checkpoint_indices(
|
||||
total, max_checks=max_checks, min_fraction=min_fraction
|
||||
)
|
||||
checks: list[dict[str, Any]] = []
|
||||
consecutive = 0
|
||||
converged_index: int | None = None
|
||||
converged_sims: dict[str, float] | None = None
|
||||
for index in indices:
|
||||
prefix = _prefix_profile(
|
||||
requests, index, window, gpu_count=gpu_count, length_mode=length_mode
|
||||
)
|
||||
sims = _family_similarity(target.vector, prefix.vector)
|
||||
checks.append(
|
||||
{
|
||||
"index": index,
|
||||
"fraction": float(index / total),
|
||||
"time_s": float(requests[index - 1].arrival_s),
|
||||
"family_similarity": sims,
|
||||
}
|
||||
)
|
||||
passed = sims["L"] >= tau and sims["A"] >= tau and sims["C"] >= tau_c
|
||||
consecutive = consecutive + 1 if passed else 0
|
||||
if consecutive >= stable_checks and converged_index is None:
|
||||
converged_index = index
|
||||
converged_sims = sims
|
||||
break
|
||||
if converged_index is None:
|
||||
last_sims = checks[-1]["family_similarity"] if checks else {"L": 1.0, "C": 1.0, "A": 1.0}
|
||||
return ConvergencePoint(
|
||||
converged=False,
|
||||
stop_index=total,
|
||||
stop_time_s=float(requests[-1].arrival_s),
|
||||
fraction=1.0,
|
||||
family_similarity=last_sims,
|
||||
checks=checks,
|
||||
)
|
||||
return ConvergencePoint(
|
||||
converged=True,
|
||||
stop_index=converged_index,
|
||||
stop_time_s=float(requests[converged_index - 1].arrival_s),
|
||||
fraction=float(converged_index / total),
|
||||
family_similarity=converged_sims or {},
|
||||
checks=checks,
|
||||
)
|
||||
|
||||
|
||||
def _prefix_profile(
|
||||
requests: list[TraceRequest],
|
||||
index: int,
|
||||
window: WindowRecord,
|
||||
*,
|
||||
gpu_count: int,
|
||||
length_mode: str,
|
||||
) -> WorkloadProfile:
|
||||
prefix = requests[:index]
|
||||
end = float(prefix[-1].arrival_s) if prefix else float(window.window_start)
|
||||
prefix_window = WindowRecord(
|
||||
window_id=window.window_id,
|
||||
trace_path=window.trace_path,
|
||||
trace_type=window.trace_type,
|
||||
window_start=window.window_start,
|
||||
window_end=end,
|
||||
source_payload=window.source_payload,
|
||||
)
|
||||
return build_workload_profile(
|
||||
prefix, prefix_window, gpu_count=gpu_count, length_mode=length_mode
|
||||
)
|
||||
|
||||
|
||||
def _checkpoint_indices(total: int, *, max_checks: int, min_fraction: float) -> list[int]:
|
||||
start = max(1, int(math.ceil(min_fraction * total)))
|
||||
if total <= max_checks:
|
||||
candidates = range(start, total + 1)
|
||||
else:
|
||||
step = max(1, total // max_checks)
|
||||
candidates = list(range(start, total + 1, step))
|
||||
if candidates and candidates[-1] != total:
|
||||
candidates.append(total)
|
||||
seen: list[int] = []
|
||||
for value in candidates:
|
||||
clamped = min(total, max(1, int(value)))
|
||||
if not seen or seen[-1] != clamped:
|
||||
seen.append(clamped)
|
||||
return seen
|
||||
|
||||
|
||||
def dumps_profile(profile: WorkloadProfile) -> str:
|
||||
return json.dumps(profile.to_dict(), ensure_ascii=False, indent=2) + "\n"
|
||||
|
||||
|
||||
@@ -3,12 +3,15 @@ from __future__ import annotations
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from .harness import build_harness_context, render_harness_context
|
||||
from .http_client import chat_completion, stream_text_completion
|
||||
from .spec import LLMPolicySpec, Proposal, SpecError, StudySpec, StudyState
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .lca import WorkloadProfile
|
||||
|
||||
|
||||
def _parse_bool_like(value: Any, *, context: str) -> bool:
|
||||
if isinstance(value, bool):
|
||||
@@ -178,6 +181,7 @@ def build_prompt(
|
||||
window_summary: dict[str, Any],
|
||||
state: StudyState,
|
||||
capability_profile: dict[str, Any] | None,
|
||||
workload_profile: "WorkloadProfile | None" = None,
|
||||
) -> str:
|
||||
objective_notes: list[str] = []
|
||||
if study.trace.request_mode == "decode_only":
|
||||
@@ -409,6 +413,7 @@ def build_prompt(
|
||||
study=study,
|
||||
window_summary=window_summary,
|
||||
state=state,
|
||||
workload_profile=workload_profile,
|
||||
)
|
||||
),
|
||||
"",
|
||||
|
||||
@@ -321,6 +321,59 @@ class InputLengthFilterSpec:
|
||||
return spec
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AdaptiveStopSpec:
|
||||
"""Stop-A: truncate per-probe replay once the offered L-C-A converges.
|
||||
|
||||
Disabled by default; the thresholds are calibrated per workload (Phase 3)
|
||||
before being switched on, so existing studies are unaffected.
|
||||
"""
|
||||
|
||||
enabled: bool = False
|
||||
tau: float = 0.9
|
||||
tau_c: float = 0.92
|
||||
stable_checks: int = 3
|
||||
max_checks: int = 20
|
||||
min_fraction: float = 0.1
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Any) -> "AdaptiveStopSpec":
|
||||
if data is None:
|
||||
return cls()
|
||||
m = _require_mapping(data, context="trace.adaptive_stop")
|
||||
enabled = (
|
||||
_require_bool(m.get("enabled"), context="trace.adaptive_stop.enabled")
|
||||
if m.get("enabled") is not None
|
||||
else False
|
||||
)
|
||||
tau = _require_float(m.get("tau", 0.9), context="trace.adaptive_stop.tau")
|
||||
tau_c = _require_float(m.get("tau_c", 0.92), context="trace.adaptive_stop.tau_c")
|
||||
stable_checks = _require_int(
|
||||
m.get("stable_checks", 3), context="trace.adaptive_stop.stable_checks"
|
||||
)
|
||||
max_checks = _require_int(
|
||||
m.get("max_checks", 20), context="trace.adaptive_stop.max_checks"
|
||||
)
|
||||
min_fraction = _require_float(
|
||||
m.get("min_fraction", 0.1), context="trace.adaptive_stop.min_fraction"
|
||||
)
|
||||
for name, value in (("tau", tau), ("tau_c", tau_c), ("min_fraction", min_fraction)):
|
||||
if not 0.0 < value <= 1.0:
|
||||
raise SpecError(f"trace.adaptive_stop.{name} must be in (0, 1].")
|
||||
if stable_checks <= 0 or max_checks <= 0:
|
||||
raise SpecError(
|
||||
"trace.adaptive_stop.stable_checks and max_checks must be > 0."
|
||||
)
|
||||
return cls(
|
||||
enabled=enabled,
|
||||
tau=tau,
|
||||
tau_c=tau_c,
|
||||
stable_checks=stable_checks,
|
||||
max_checks=max_checks,
|
||||
min_fraction=min_fraction,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TraceSpec:
|
||||
windows_path: str
|
||||
@@ -338,6 +391,7 @@ class TraceSpec:
|
||||
early_stop_max_lag_s: float | None = None
|
||||
early_stop_max_elapsed_s: float | None = None
|
||||
restart_engine_after_early_stop: bool = False
|
||||
adaptive_stop: AdaptiveStopSpec = AdaptiveStopSpec()
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Mapping[str, Any]) -> "TraceSpec":
|
||||
@@ -429,6 +483,7 @@ class TraceSpec:
|
||||
if data.get("restart_engine_after_early_stop") is not None
|
||||
else request_mode == "decode_only"
|
||||
),
|
||||
adaptive_stop=AdaptiveStopSpec.from_dict(data.get("adaptive_stop")),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -16,6 +16,7 @@ from typing import Any, Callable
|
||||
|
||||
from .engine import build_launch_recipe
|
||||
from .http_client import HttpClientError, stream_chat_completion, wait_for_server
|
||||
from .lca import find_convergence_prefix, resolve_length_mode
|
||||
from .search import ThresholdProbe, binary_search_max_feasible
|
||||
from .slo import RequestOutcome, evaluate_request, summarize_evaluations
|
||||
from .spec import ConfigPatch, SamplingSearchSpec, TrialSpec, load_study_spec, to_jsonable
|
||||
@@ -209,6 +210,45 @@ def _probe_outcome_details(
|
||||
}
|
||||
|
||||
|
||||
def _adaptive_replay_set(
|
||||
selected: list[TraceRequest],
|
||||
*,
|
||||
study: Any,
|
||||
window: Any,
|
||||
) -> tuple[list[TraceRequest], dict[str, Any] | None]:
|
||||
"""Stop-A: truncate the replay to the offered-L-C-A convergence prefix.
|
||||
|
||||
Returns the (possibly shortened) request list to replay and a certificate of
|
||||
the convergence decision. When Stop-A is disabled, or C never converges, the
|
||||
full selected set is replayed (the C-gate: no early stop on a cold cache).
|
||||
"""
|
||||
spec = study.trace.adaptive_stop
|
||||
if not getattr(spec, "enabled", False) or not selected:
|
||||
return selected, None
|
||||
point = find_convergence_prefix(
|
||||
selected,
|
||||
window,
|
||||
gpu_count=study.hardware.gpu_count,
|
||||
length_mode=resolve_length_mode(request_mode=study.trace.request_mode),
|
||||
tau=spec.tau,
|
||||
tau_c=spec.tau_c,
|
||||
stable_checks=spec.stable_checks,
|
||||
max_checks=spec.max_checks,
|
||||
min_fraction=spec.min_fraction,
|
||||
)
|
||||
replay = selected[: point.stop_index] if point.stop_index > 0 else selected
|
||||
certificate = {
|
||||
"enabled": True,
|
||||
"converged": point.converged,
|
||||
"stop_index": point.stop_index,
|
||||
"total_selected": len(selected),
|
||||
"fraction": point.fraction,
|
||||
"stop_time_s": point.stop_time_s,
|
||||
"family_similarity": point.family_similarity,
|
||||
}
|
||||
return replay, certificate
|
||||
|
||||
|
||||
def _best_feasible_probe_record(probe_history: list[dict[str, Any]]) -> dict[str, Any] | None:
|
||||
feasible = [
|
||||
item
|
||||
@@ -519,9 +559,12 @@ def run_trial(trial_spec_path: Path) -> dict[str, Any]:
|
||||
def evaluator(threshold: float) -> ThresholdProbe[ProbePayload]:
|
||||
nonlocal process
|
||||
selected = select_requests_for_threshold(requests, threshold=threshold)
|
||||
replay_set, adaptive_stop_certificate = _adaptive_replay_set(
|
||||
selected, study=study, window=window
|
||||
)
|
||||
restart_after_early_stop = study.trace.restart_engine_after_early_stop
|
||||
outcomes, early_stopped, early_stop_reason = _replay_requests(
|
||||
selected,
|
||||
replay_set,
|
||||
base_url=recipe.base_url,
|
||||
timeout_s=recipe.request_timeout_s,
|
||||
max_concurrency=study.trace.max_concurrency,
|
||||
@@ -534,12 +577,13 @@ def run_trial(trial_spec_path: Path) -> dict[str, Any]:
|
||||
evaluations, summary = summarize_evaluations(outcomes, study.slo)
|
||||
probe_details = _probe_outcome_details(
|
||||
threshold=threshold,
|
||||
selected=selected,
|
||||
selected=replay_set,
|
||||
outcomes=outcomes,
|
||||
evaluations=evaluations,
|
||||
early_stopped=early_stopped,
|
||||
early_stop_reason=early_stop_reason,
|
||||
)
|
||||
probe_details["adaptive_stop"] = adaptive_stop_certificate
|
||||
with probe_details_path.open("a", encoding="utf-8") as details_handle:
|
||||
details_handle.write(
|
||||
json.dumps(probe_details, ensure_ascii=False) + "\n"
|
||||
@@ -580,12 +624,14 @@ def run_trial(trial_spec_path: Path) -> dict[str, Any]:
|
||||
probe_record = {
|
||||
"threshold": threshold,
|
||||
"request_count": payload.request_count,
|
||||
"replayed_request_count": len(replay_set),
|
||||
"pass_rate": payload.pass_rate,
|
||||
"request_rate": payload.request_rate,
|
||||
"feasible": payload.feasible,
|
||||
"early_stopped": payload.early_stopped,
|
||||
"early_stop_reason": payload.early_stop_reason,
|
||||
"latency_summary": payload.latency_summary,
|
||||
"adaptive_stop": adaptive_stop_certificate,
|
||||
}
|
||||
probe_history.append(probe_record)
|
||||
StudyStore.write_json(Path(trial.probe_log_path), probe_history)
|
||||
|
||||
@@ -28,7 +28,9 @@ from aituner.harness import (
|
||||
build_harness_stop_proposal,
|
||||
)
|
||||
from aituner.lca import (
|
||||
build_study_workload_profile,
|
||||
build_workload_profile,
|
||||
find_convergence_prefix,
|
||||
profile_similarity,
|
||||
resolve_length_mode,
|
||||
similarity_report,
|
||||
@@ -37,6 +39,7 @@ from aituner.llm import _extract_response_text, build_prompt, parse_proposal_tex
|
||||
from aituner.search import ThresholdProbe, binary_search_max_feasible
|
||||
from aituner.slo import RequestOutcome, evaluate_request, summarize_evaluations
|
||||
from aituner.spec import (
|
||||
AdaptiveStopSpec,
|
||||
ConfigPatch,
|
||||
LLMEndpointSpec,
|
||||
Proposal,
|
||||
@@ -48,6 +51,7 @@ from aituner.spec import (
|
||||
from aituner.store import StudyStore
|
||||
from aituner.trace import load_trace_requests, summarize_window
|
||||
from aituner.worker import (
|
||||
_adaptive_replay_set,
|
||||
_best_feasible_probe_record,
|
||||
_latency_summary,
|
||||
_run_one_request,
|
||||
@@ -298,6 +302,162 @@ class CoreFlowTests(unittest.TestCase):
|
||||
self.assertAlmostEqual(profile.stats["arrival"]["fano_1s"], 0.5)
|
||||
self.assertEqual(resolve_length_mode(request_mode="decode_only"), "output")
|
||||
|
||||
def test_harness_context_uses_canonical_lca_vector(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
tmp_path = Path(tmp)
|
||||
study_path = _write_study_assets(tmp_path)
|
||||
study = load_study_spec(study_path)
|
||||
window, requests = load_trace_requests(study, study_spec_path=study_path)
|
||||
profile = build_study_workload_profile(study, requests, window)
|
||||
state = StudyState(study_id=study.study_id, trials=[])
|
||||
summary = summarize_window(requests, window)
|
||||
context = build_harness_context(
|
||||
study=study,
|
||||
window_summary=summary,
|
||||
state=state,
|
||||
workload_profile=profile,
|
||||
)
|
||||
block = context["workload_lca_profile"]
|
||||
# The labeled L-C-A block is the canonical 10-dim metric, not ad-hoc.
|
||||
self.assertEqual(block["vector"], profile.vector)
|
||||
self.assertEqual(len(block["vector"]), 10)
|
||||
self.assertIn("RobustScaler", block["metric"])
|
||||
# Without a profile it falls back to the legacy ad-hoc rendering.
|
||||
legacy = build_harness_context(
|
||||
study=study,
|
||||
window_summary=summary,
|
||||
state=state,
|
||||
)["workload_lca_profile"]
|
||||
self.assertNotIn("vector", legacy)
|
||||
|
||||
def _steady_requests(self, count: int, *, input_tokens: int = 100) -> list:
|
||||
return [
|
||||
TraceRequest(
|
||||
row_id=f"r{i}",
|
||||
arrival_s=float(i),
|
||||
sampling_u=1.0,
|
||||
body={},
|
||||
prompt_tokens_hint=input_tokens,
|
||||
completion_tokens_hint=16,
|
||||
metadata={"hash_ids": None},
|
||||
)
|
||||
for i in range(count)
|
||||
]
|
||||
|
||||
def _conv_window(self) -> WindowRecord:
|
||||
return WindowRecord(
|
||||
window_id="conv",
|
||||
trace_path=Path("trace.jsonl"),
|
||||
trace_type="chat",
|
||||
window_start=0.0,
|
||||
window_end=0.0,
|
||||
source_payload={"block_size": 64},
|
||||
)
|
||||
|
||||
def test_convergence_prefix_stops_early_on_stationary_trace(self) -> None:
|
||||
requests = self._steady_requests(60)
|
||||
point = find_convergence_prefix(
|
||||
requests,
|
||||
self._conv_window(),
|
||||
gpu_count=1,
|
||||
length_mode="total",
|
||||
tau=0.9,
|
||||
tau_c=0.9,
|
||||
stable_checks=3,
|
||||
max_checks=20,
|
||||
min_fraction=0.1,
|
||||
)
|
||||
self.assertTrue(point.converged)
|
||||
# A stationary workload should be trustworthy well before the full window.
|
||||
self.assertLess(point.stop_index, len(requests))
|
||||
self.assertLess(point.fraction, 1.0)
|
||||
self.assertTrue(point.checks)
|
||||
|
||||
def test_convergence_prefix_waits_when_cache_warms_late(self) -> None:
|
||||
window = self._conv_window()
|
||||
# First half: no prefix reuse. Second half: every request reuses block 1,
|
||||
# so the C dimension only stabilizes once the reuse regime is exercised.
|
||||
requests = []
|
||||
for i in range(30):
|
||||
requests.append(
|
||||
TraceRequest(
|
||||
row_id=f"cold{i}",
|
||||
arrival_s=float(i),
|
||||
sampling_u=1.0,
|
||||
body={},
|
||||
prompt_tokens_hint=640,
|
||||
completion_tokens_hint=16,
|
||||
metadata={"hash_ids": [10_000 + i]},
|
||||
)
|
||||
)
|
||||
for i in range(30):
|
||||
requests.append(
|
||||
TraceRequest(
|
||||
row_id=f"warm{i}",
|
||||
arrival_s=float(30 + i),
|
||||
sampling_u=1.0,
|
||||
body={},
|
||||
prompt_tokens_hint=640,
|
||||
completion_tokens_hint=16,
|
||||
metadata={"hash_ids": [1, 2, 3, 4, 5]},
|
||||
)
|
||||
)
|
||||
point = find_convergence_prefix(
|
||||
requests,
|
||||
window,
|
||||
gpu_count=1,
|
||||
length_mode="total",
|
||||
tau=0.9,
|
||||
tau_c=0.95,
|
||||
stable_checks=2,
|
||||
max_checks=20,
|
||||
min_fraction=0.1,
|
||||
)
|
||||
# The C family similarity must be low while only the cold half is seen.
|
||||
early = [c for c in point.checks if c["fraction"] <= 0.4]
|
||||
self.assertTrue(early)
|
||||
self.assertTrue(any(c["family_similarity"]["C"] < 0.9 for c in early))
|
||||
|
||||
def test_adaptive_replay_set_truncates_only_when_enabled(self) -> None:
|
||||
from types import SimpleNamespace
|
||||
|
||||
requests = self._steady_requests(60)
|
||||
window = self._conv_window()
|
||||
enabled_study = SimpleNamespace(
|
||||
trace=SimpleNamespace(
|
||||
adaptive_stop=AdaptiveStopSpec(
|
||||
enabled=True,
|
||||
tau=0.9,
|
||||
tau_c=0.9,
|
||||
stable_checks=3,
|
||||
max_checks=20,
|
||||
min_fraction=0.1,
|
||||
),
|
||||
request_mode="chat",
|
||||
),
|
||||
hardware=SimpleNamespace(gpu_count=1),
|
||||
)
|
||||
replay, certificate = _adaptive_replay_set(
|
||||
requests, study=enabled_study, window=window
|
||||
)
|
||||
self.assertIsNotNone(certificate)
|
||||
self.assertTrue(certificate["enabled"])
|
||||
self.assertEqual(len(replay), certificate["stop_index"])
|
||||
self.assertLessEqual(len(replay), len(requests))
|
||||
|
||||
disabled_study = SimpleNamespace(
|
||||
trace=SimpleNamespace(
|
||||
adaptive_stop=AdaptiveStopSpec(enabled=False),
|
||||
request_mode="chat",
|
||||
),
|
||||
hardware=SimpleNamespace(gpu_count=1),
|
||||
)
|
||||
passthrough, no_cert = _adaptive_replay_set(
|
||||
requests, study=disabled_study, window=window
|
||||
)
|
||||
self.assertIsNone(no_cert)
|
||||
self.assertEqual(len(passthrough), len(requests))
|
||||
|
||||
def test_lca_similarity_matrix_separates_different_profiles(self) -> None:
|
||||
window = WindowRecord(
|
||||
window_id="base",
|
||||
|
||||
100
tests/test_prepare_trace_windows.py
Normal file
100
tests/test_prepare_trace_windows.py
Normal file
@@ -0,0 +1,100 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||
_SPEC = importlib.util.spec_from_file_location(
|
||||
"prepare_trace_windows",
|
||||
REPO_ROOT / "scripts" / "prepare_trace_windows.py",
|
||||
)
|
||||
assert _SPEC and _SPEC.loader
|
||||
ptw = importlib.util.module_from_spec(_SPEC)
|
||||
# Register before exec so dataclasses can resolve the module's annotations.
|
||||
sys.modules[_SPEC.name] = ptw
|
||||
_SPEC.loader.exec_module(ptw)
|
||||
|
||||
|
||||
class SessionCoherentSamplingTests(unittest.TestCase):
|
||||
def test_multi_hop_chain_resolves_to_root(self) -> None:
|
||||
root_of: dict[object, object] = {}
|
||||
# turn1 root, turn2 -> turn1, turn3 -> turn2 (multi-hop), streamed in order.
|
||||
self.assertEqual(
|
||||
ptw.resolve_session_root({"chat_id": 1, "parent_chat_id": -1, "turn": 1}, root_of),
|
||||
1,
|
||||
)
|
||||
self.assertEqual(
|
||||
ptw.resolve_session_root({"chat_id": 2, "parent_chat_id": 1, "turn": 2}, root_of),
|
||||
1,
|
||||
)
|
||||
self.assertEqual(
|
||||
ptw.resolve_session_root({"chat_id": 3, "parent_chat_id": 2, "turn": 3}, root_of),
|
||||
1,
|
||||
)
|
||||
|
||||
def test_unknown_parent_falls_back_to_parent_id(self) -> None:
|
||||
root_of: dict[object, object] = {}
|
||||
# parent never seen (fell outside the span): group siblings under the parent.
|
||||
self.assertEqual(
|
||||
ptw.resolve_session_root({"chat_id": 50, "parent_chat_id": 9, "turn": 2}, root_of),
|
||||
9,
|
||||
)
|
||||
self.assertEqual(
|
||||
ptw.resolve_session_root({"chat_id": 51, "parent_chat_id": 9, "turn": 2}, root_of),
|
||||
9,
|
||||
)
|
||||
|
||||
def test_all_turns_of_a_session_share_one_u(self) -> None:
|
||||
root_of: dict[object, object] = {}
|
||||
rows = [
|
||||
{"chat_id": 1, "parent_chat_id": -1, "turn": 1},
|
||||
{"chat_id": 2, "parent_chat_id": 1, "turn": 2},
|
||||
{"chat_id": 3, "parent_chat_id": 2, "turn": 3},
|
||||
]
|
||||
us = {
|
||||
ptw.session_uniform(
|
||||
seed=7,
|
||||
window_id="w",
|
||||
session_root=ptw.resolve_session_root(row, root_of),
|
||||
)
|
||||
for row in rows
|
||||
}
|
||||
self.assertEqual(len(us), 1)
|
||||
only = next(iter(us))
|
||||
self.assertGreaterEqual(only, 0.0)
|
||||
self.assertLess(only, 1.0)
|
||||
|
||||
def test_thresholding_keeps_or_drops_whole_sessions(self) -> None:
|
||||
# Two distinct sessions get distinct scores; a threshold either keeps a
|
||||
# session's every turn or none of them.
|
||||
seed, window_id = 20260325, "chat_w_x"
|
||||
sessions = {
|
||||
"A": [
|
||||
{"chat_id": 10, "parent_chat_id": -1},
|
||||
{"chat_id": 11, "parent_chat_id": 10},
|
||||
],
|
||||
"B": [
|
||||
{"chat_id": 20, "parent_chat_id": -1},
|
||||
{"chat_id": 21, "parent_chat_id": 20},
|
||||
],
|
||||
}
|
||||
root_of: dict[object, object] = {}
|
||||
scored: list[tuple[str, float]] = []
|
||||
for name, rows in sessions.items():
|
||||
for row in rows:
|
||||
root = ptw.resolve_session_root(row, root_of)
|
||||
u = ptw.session_uniform(seed=seed, window_id=window_id, session_root=root)
|
||||
scored.append((name, u))
|
||||
for name in sessions:
|
||||
us = {u for n, u in scored if n == name}
|
||||
self.assertEqual(len(us), 1, f"session {name} turns must share one u")
|
||||
for threshold in (0.0, 0.25, 0.5, 0.75, 1.0):
|
||||
for name in sessions:
|
||||
kept = {u <= threshold for n, u in scored if n == name}
|
||||
self.assertEqual(len(kept), 1, "a session must be kept/dropped as a whole")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user