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

Drain inflight requests after early stop

Browse Source
This commit is contained in:
Gahow Wang
2026-04-25 16:57:01 +08:00
parent 2dc2815620
commit 2d7ebe50ee
3 changed files with 59 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
src/aituner/worker.py
View File

@@ -135,6 +135,7 @@ def _replay_requests(
max_lag_s: float | None,
max_elapsed_s: float | None,
evaluate_outcome: Callable[[RequestOutcome], Any],
drain_inflight_on_early_stop: bool = True,
) -> tuple[list[RequestOutcome], bool, str]:
outcomes_by_id: dict[str, RequestOutcome] = {}
lock = threading.Lock()
@@ -209,18 +210,47 @@ def _replay_requests(
if sleep_for > 0:
time.sleep(min(sleep_for, 0.1))
if early_stopped:
for future in list(futures_by_request):
future.cancel()
for request in futures_by_request.values():
outcomes_by_id[request.row_id] = RequestOutcome(
request_id=request.row_id,
success=False,
ttft_ms=None,
tpot_ms=None,
prompt_tokens=request.prompt_tokens_hint,
completion_tokens=request.completion_tokens_hint,
error=early_stop_reason or "probe_early_stopped",
)
if drain_inflight_on_early_stop and futures_by_request:
done, not_done = wait(list(futures_by_request), timeout=timeout_s)
for future in done:
request = futures_by_request[future]
try:
outcomes_by_id[request.row_id] = future.result(timeout=0)
except Exception: # noqa: BLE001
outcomes_by_id[request.row_id] = RequestOutcome(
request_id=request.row_id,
success=False,
ttft_ms=None,
tpot_ms=None,
prompt_tokens=request.prompt_tokens_hint,
completion_tokens=request.completion_tokens_hint,
error=early_stop_reason or "probe_early_stopped",
)
for future in not_done:
future.cancel()
request = futures_by_request[future]
outcomes_by_id[request.row_id] = RequestOutcome(
request_id=request.row_id,
success=False,
ttft_ms=None,
tpot_ms=None,
prompt_tokens=request.prompt_tokens_hint,
completion_tokens=request.completion_tokens_hint,
error=early_stop_reason or "probe_early_stopped",
)
else:
for future in list(futures_by_request):
future.cancel()
for request in futures_by_request.values():
outcomes_by_id[request.row_id] = RequestOutcome(
request_id=request.row_id,
success=False,
ttft_ms=None,
tpot_ms=None,
prompt_tokens=request.prompt_tokens_hint,
completion_tokens=request.completion_tokens_hint,
error=early_stop_reason or "probe_early_stopped",
)
for request in requests:
if request.row_id in submitted_ids:
continue