gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add trace length bucket tuning support

Browse Source
This commit is contained in:
Gahow Wang
2026-04-07 11:03:16 +08:00
parent e9b5e9b957
commit 46ed688ace
12 changed files with 922 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
src/aituner/llm.py
View File

@@ -42,6 +42,11 @@ def build_prompt(
json.dumps(
{
"study_id": study.study_id,
"current_best": {
"trial_id": state.best_trial_id,
"best_sampling_u": state.best_sampling_u,
"best_request_rate": state.best_request_rate,
},
"hardware": {
"gpu_count": study.hardware.gpu_count,
"gpu_model": study.hardware.gpu_model,
@@ -50,6 +55,17 @@ def build_prompt(
"model_id": study.model.model_id,
"served_model_name": study.model.served_model_name,
},
"trace": {
"window_id": study.trace.window_id,
"input_length_filter": (
{
"min_input_tokens": study.trace.input_length_filter.min_input_tokens,
"max_input_tokens": study.trace.input_length_filter.max_input_tokens,
}
if study.trace.input_length_filter is not None
else None
),
},
"engine": {
"engine_name": study.engine.engine_name,
"engine_version": study.engine.engine_version,
@@ -84,6 +100,8 @@ def build_prompt(
"Trial history:",
json.dumps(history, ensure_ascii=False, indent=2),
"",
"The proposal must beat the current incumbent. Do not propose a config that is only likely to be feasible below the current best_sampling_u/request_rate.",
"The evaluator for a new trial will start searching from the current best feasible sampling_u and only look for improvements above it.",
"The proposal should improve the maximum feasible sampling_u under the 95%+ SLO target.",
]
return "\n".join(sections)
@@ -110,8 +128,22 @@ def validate_proposal(proposal: Proposal, study: StudySpec) -> Proposal:
return proposal
def _parse_json_object_text(text: str) -> dict[str, Any]:
try:
payload = json.loads(text)
except json.JSONDecodeError:
start = text.find("{")
end = text.rfind("}")
if start < 0 or end < start:
raise
payload = json.loads(text[start : end + 1])
if not isinstance(payload, dict):
raise SpecError("proposal payload must be a JSON object")
return payload
def parse_proposal_text(text: str, study: StudySpec) -> Proposal:
payload = json.loads(text)
payload = _parse_json_object_text(text)
proposal = Proposal.from_dict(payload)
return validate_proposal(proposal, study)

49
src/aituner/spec.py
View File

@@ -142,6 +142,42 @@ class EngineLaunchSpec:
)
@dataclass(frozen=True)
class InputLengthFilterSpec:
min_input_tokens: int | None = None
max_input_tokens: int | None = None
@classmethod
def from_dict(cls, data: Mapping[str, Any], *, context: str) -> "InputLengthFilterSpec":
min_input_tokens = data.get("min_input_tokens")
max_input_tokens = data.get("max_input_tokens")
spec = cls(
min_input_tokens=(
_require_int(min_input_tokens, context=f"{context}.min_input_tokens")
if min_input_tokens is not None
else None
),
max_input_tokens=(
_require_int(max_input_tokens, context=f"{context}.max_input_tokens")
if max_input_tokens is not None
else None
),
)
if spec.min_input_tokens is None and spec.max_input_tokens is None:
raise SpecError(
f"{context} must define at least one of min_input_tokens/max_input_tokens."
)
if (
spec.min_input_tokens is not None
and spec.max_input_tokens is not None
and spec.min_input_tokens > spec.max_input_tokens
):
raise SpecError(
f"{context}.min_input_tokens must be <= {context}.max_input_tokens."
)
return spec
@dataclass(frozen=True)
class TraceSpec:
windows_path: str
@@ -150,6 +186,7 @@ class TraceSpec:
u_field: str
timestamp_field: str
max_concurrency: int
input_length_filter: InputLengthFilterSpec | None = None
max_requests_per_probe: int | None = None
synthetic_prompt_cap_tokens: int | None = None
replay_time_scale: float = 1.0
@@ -171,6 +208,17 @@ class TraceSpec:
max_concurrency=_require_int(
data.get("max_concurrency", 64), context="trace.max_concurrency"
),
input_length_filter=(
InputLengthFilterSpec.from_dict(
_require_mapping(
data.get("input_length_filter"),
context="trace.input_length_filter",
),
context="trace.input_length_filter",
)
if data.get("input_length_filter") is not None
else None
),
max_requests_per_probe=int(max_requests) if max_requests is not None else None,
synthetic_prompt_cap_tokens=(
int(synthetic_prompt_cap) if synthetic_prompt_cap is not None else None
@@ -454,6 +502,7 @@ class TrialSummary:
class StudyState:
study_id: str
best_trial_id: str | None = None
best_sampling_u: float | None = None
best_request_rate: float | None = None
next_trial_index: int = 1
trials: list[TrialSummary] = field(default_factory=list)

21
src/aituner/store.py
View File

@@ -32,6 +32,7 @@ class StudyStore:
return StudyState(
study_id=str(payload["study_id"]),
best_trial_id=payload.get("best_trial_id"),
best_sampling_u=payload.get("best_sampling_u"),
best_request_rate=payload.get("best_request_rate"),
next_trial_index=int(payload.get("next_trial_index", 1)),
trials=trials,
@@ -64,7 +65,18 @@ class StudyStore:
study_id=study.study_id,
trial_id=trial_id,
config_patch=proposal.config_patch,
search=study.search,
search=replace(
study.search,
low=min(
study.search.high,
max(
study.search.low,
float(state.best_sampling_u)
if isinstance(state.best_sampling_u, (int, float))
else study.search.low,
),
),
),
study_spec_path=str((self.study_root(study.study_id) / "study_spec.source").resolve()),
artifact_dir=str(trial_root),
probe_log_path=str(trial_root / "probe_history.json"),
@@ -89,6 +101,7 @@ class StudyStore:
by_id = {item.trial_id: item for item in state.trials}
trials_dir = self.study_root(study_id) / "trials"
best_trial_id = state.best_trial_id
best_sampling_u = state.best_sampling_u
best_rate = state.best_request_rate
for trial_dir in sorted(trials_dir.glob("trial-*")):
result_path = trial_dir / "result.json"
@@ -112,7 +125,13 @@ class StudyStore:
and (best_rate is None or summary.best_request_rate > best_rate)
):
best_rate = float(summary.best_request_rate)
best_sampling_u = (
float(summary.best_sampling_u)
if isinstance(summary.best_sampling_u, (int, float))
else None
)
best_trial_id = trial_id
state.best_sampling_u = best_sampling_u
state.best_request_rate = best_rate
state.best_trial_id = best_trial_id
self.save_state(state)

21
src/aituner/trace.py
View File

@@ -132,6 +132,25 @@ def _downsample_requests(
return [requests[idx] for idx in indexes]
def _matches_input_length_filter(study: StudySpec, *, prompt_tokens_hint: int | None) -> bool:
length_filter = study.trace.input_length_filter
if length_filter is None:
return True
if prompt_tokens_hint is None:
return False
if (
length_filter.min_input_tokens is not None
and prompt_tokens_hint < length_filter.min_input_tokens
):
return False
if (
length_filter.max_input_tokens is not None
and prompt_tokens_hint > length_filter.max_input_tokens
):
return False
return True
def load_trace_requests(study: StudySpec, *, study_spec_path: Path) -> tuple[WindowRecord, list[TraceRequest]]:
window = resolve_window_record(study, study_spec_path=study_spec_path)
time_scale = float(study.trace.replay_time_scale)
@@ -163,6 +182,8 @@ def load_trace_requests(study: StudySpec, *, study_spec_path: Path) -> tuple[Win
if isinstance(sampling_u, bool) or not isinstance(sampling_u, (int, float)):
raise TraceError(f"trace row {idx} is missing numeric {study.trace.u_field}")
prompt_tokens_hint = _coerce_prompt_tokens(row)
if not _matches_input_length_filter(study, prompt_tokens_hint=prompt_tokens_hint):
continue
try:
messages = _coerce_messages(row)
except TraceError:

9
src/aituner/worker.py
View File

@@ -177,14 +177,19 @@ def _replay_requests(
if early_stopped:
break
if futures_by_request:
timeout = None
timeout = 0.5
if next_index < len(requests):
timeout = max(0.0, requests[next_index].arrival_s - elapsed)
timeout = min(timeout, max(0.0, requests[next_index].arrival_s - elapsed))
if max_elapsed_s is not None:
remaining_elapsed = max(0.0, max_elapsed_s - elapsed)
timeout = min(timeout, remaining_elapsed)
done, _ = wait(
list(futures_by_request),
timeout=timeout,
return_when=FIRST_COMPLETED,
)
if not done:
continue
for future in done:
request = futures_by_request.pop(future)
outcome = future.result()