Compare commits

3 Commits

Author SHA1 Message Date
51a9e4a007 Add Stop-A: offered-L-C-A convergence early-stop for replay
Phase 2 of the two-stop work. The L-C-A vector is a deterministic function of the
trace's offered metadata, so the convergence of prefix-vs-full L-C-A (the paper's
Fig. 9 curve) can be computed up front rather than monitored live, with identical
result and no per-request overhead.

- lca.find_convergence_prefix: earliest arrival-ordered prefix whose L and A family
  similarities reach tau and the slow C family reaches the stricter tau_c for
  stable_checks consecutive checkpoints. Self-similarity uses the raw log-feature
  vector (same window -> identical per-dim spread; RobustScaler is reserved for the
  cross-window Stop-C). If C never converges it reports the full set, which is the
  C-gate: no early stop on a cold/under-warmed cache. The checkpoint sims double as
  Phase 3 calibration data.
- spec.AdaptiveStopSpec (trace.adaptive_stop), disabled by default until the
  thresholds are calibrated, so existing studies are unaffected.
- worker._adaptive_replay_set truncates each probe's replay to the convergence
  prefix and records a certificate (converged, fraction, family similarity) into
  probe history and probe_details. Offered request_rate at the threshold is
  unchanged; only wall-clock replay shrinks.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-15 14:23:49 +08:00
0f15bbc3f1 Make the offered-load axis session-coherent
Phase 1 of the two-stop work. Subsampling the trace by per-request uniform score
broke multi-turn sessions (a kept turn-2 could lose its turn-1), which lowered the
realized KV-cache hit rate as offered load dropped — so the feasibility boundary
was measured on a workload with a different C than production, contradicting the
paper's scale-stationary L-C-A premise.

prepare_trace_windows now resolves each row's session root via the parent_chat_id
chain in a single streaming pass and assigns sampling_u per session, so thresholding
keeps or drops whole sessions and preserves intra-session prefix reuse. Rows whose
parent fell outside the span fall back to grouping under the parent id.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-15 14:16:06 +08:00
6f8e3c95c1 Unify harness L-C-A on the canonical lca.WorkloadProfile
Phase 0 of the two-stop work. The prompt block labeled `workload_lca_profile`
previously re-derived L-C-A from summarize_window's ad-hoc percentiles, diverging
from the paper's 10-dim RobustScaler vector implemented in lca.py. Make that block
authoritative: build_harness_context now accepts an optional workload_profile and
renders the canonical 10-dim vector + per-family stats when present, falling back
to the legacy rendering only when no profile is supplied (direct unit-test calls).

Real call sites (study prompt/llm-propose/tune, run_baseline_then_llm) build the
profile via lca.build_study_workload_profile and pass it through build_prompt. The
heuristic regime classifiers keep reading window_summary; that is the heuristic
layer, distinct from the similarity metric.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-15 14:12:17 +08:00
10 changed files with 642 additions and 17 deletions

View File

@@ -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

View File

@@ -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)

View File

@@ -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)

View File

@@ -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],

View File

@@ -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"

View File

@@ -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,
)
),
"",

View File

@@ -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")),
)

View File

@@ -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)

View File

@@ -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",

View 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()