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

Stop waiting on in-flight requests after early stop

Browse Source
This commit is contained in:
Gahow Wang
2026-04-05 00:56:26 +08:00
parent 75a9842f1a
commit e00bedb466
2 changed files with 108 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
src/aituner/worker.py
View File

@@ -150,21 +150,19 @@ def _replay_requests(
sleep_for = max(0.0, requests[next_index].arrival_s - elapsed)
if sleep_for > 0:
time.sleep(min(sleep_for, 0.1))
for future, request in list(futures_by_request.items()):
try:
outcome = future.result(timeout=timeout_s)
except Exception: # noqa: BLE001
outcome = RequestOutcome(
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="probe_early_stopped",
error=early_stop_reason or "probe_early_stopped",
)
outcomes_by_id[outcome.request_id] = outcome
if early_stopped:
for request in requests:
if request.row_id in submitted_ids:
continue
@@ -177,6 +175,21 @@ def _replay_requests(
completion_tokens=request.completion_tokens_hint,
error=early_stop_reason or "probe_early_stopped",
)
else:
for future, request in list(futures_by_request.items()):
try:
outcome = future.result(timeout=timeout_s)
except Exception: # noqa: BLE001
outcome = 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="probe_early_stopped",
)
outcomes_by_id[outcome.request_id] = outcome
finally:
pool.shutdown(wait=False, cancel_futures=True)
ordered = [outcomes_by_id[item.row_id] for item in requests if item.row_id in outcomes_by_id]