Skip to content

Commit 344a376

Browse files
committed
refactor(watchdog): extract _RunWatchdog beside _PerfPhaseTimeout
One object owns the whole-run deadline: timer handle, fired flag, and the late-bound session, replacing the nonlocal flag + mutable-holder closure threaded through _run_benchmark_async. Behavior unchanged.
1 parent 4faac58 commit 344a376

1 file changed

Lines changed: 57 additions & 45 deletions

File tree

  • src/inference_endpoint/commands/benchmark

‎src/inference_endpoint/commands/benchmark/execute.py‎

Lines changed: 57 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -716,6 +716,56 @@ def cancel(self) -> None:
716716
self._handle = None
717717

718718

719+
class _RunWatchdog:
720+
"""Whole-run deadline timer for ``settings.timeouts.run_timeout_s``.
721+
722+
Armed before the metrics pipeline starts (so service-launch and
723+
endpoint-connect stalls are bounded) and kept armed through the metrics
724+
drain (so a stuck aggregator drain is bounded too). On fire: stop the
725+
session first — it short-circuits its drain and publishes ENDED promptly,
726+
so the event logger flushes and the aggregator records the buffered
727+
tokenizer-drain samples — then SIGTERM the aggregator, whose handler
728+
writes the INTERRUPTED final snapshot (``publish_final`` is first-wins,
729+
so INTERRUPTED stays authoritative). ``run_benchmark`` raises
730+
``ExecutionError`` after finalization whenever ``fired`` is set, so a
731+
timed-out run always fails loudly even if a still-draining aggregator
732+
finalized COMPLETE first.
733+
"""
734+
735+
def __init__(
736+
self,
737+
loop: asyncio.AbstractEventLoop,
738+
deadline: float | None,
739+
pipe: MetricsPipeline,
740+
) -> None:
741+
self.fired = False
742+
self._session: BenchmarkSession | None = None
743+
self._pipe = pipe
744+
self._handle = (
745+
loop.call_later(max(0.0, deadline - time.monotonic()), self._fire)
746+
if deadline is not None
747+
else None
748+
)
749+
750+
def bind_session(self, session: BenchmarkSession) -> None:
751+
"""Late-bind the session: it is created after the timer is armed."""
752+
self._session = session
753+
754+
def _fire(self) -> None:
755+
self.fired = True
756+
logger.error(
757+
"Run timeout reached; aborting run — report will be marked " "INTERRUPTED."
758+
)
759+
if self._session is not None:
760+
self._session.stop()
761+
self._pipe.terminate_metrics_aggregator()
762+
763+
def cancel(self) -> None:
764+
if self._handle is not None:
765+
self._handle.cancel()
766+
self._handle = None
767+
768+
719769
async def _create_issuer(
720770
ctx: BenchmarkContext, loop: asyncio.AbstractEventLoop
721771
) -> tuple[HttpClientSampleIssuer, HTTPEndpointClient]:
@@ -848,44 +898,7 @@ async def _run_benchmark_async(
848898
# idempotent, so the clean-path shutdown below is a harmless second call.
849899
http_client: HTTPEndpointClient | None = None
850900

851-
# Whole-run watchdog. Armed before the pipeline starts so setup stalls
852-
# (service launch, endpoint connect) are bounded too, and kept armed
853-
# through the metrics drain so run_timeout_s can SIGTERM a stuck
854-
# aggregator drain. Cancelled in the outermost finally.
855-
run_timed_out = False
856-
# The session is created later inside the pipeline scope; bind it through
857-
# a mutable holder so the callback never touches a possibly-unbound local
858-
# (a NameError inside a loop callback is swallowed by the loop's exception
859-
# handler, which would leave the watchdog inert).
860-
session_ref: list[BenchmarkSession] = []
861-
run_timeout_s = config.settings.timeouts.run_timeout_s
862-
863-
def _on_run_timeout() -> None:
864-
nonlocal run_timed_out
865-
run_timed_out = True
866-
logger.error(
867-
"Run timeout (%.1fs) reached; aborting run — report will be "
868-
"marked INTERRUPTED.",
869-
run_timeout_s,
870-
)
871-
# Stop the session first: it short-circuits _drain_inflight and
872-
# run()'s finally publishes ENDED promptly, so the aggregator still
873-
# records the buffered tokenizer-drain samples. Then SIGTERM the
874-
# aggregator: its handler writes the INTERRUPTED final snapshot
875-
# (publish_final is first-wins, so INTERRUPTED stays authoritative;
876-
# even if a still-draining aggregator finalizes as COMPLETE first,
877-
# run_benchmark raises on run_timed_out, so a timed-out run always
878-
# fails loudly). Targeted (not all services): the event logger
879-
# flushes on ENDED, which session.stop() still delivers.
880-
if session_ref:
881-
session_ref[0].stop()
882-
pipe.terminate_metrics_aggregator()
883-
884-
run_watchdog = (
885-
loop.call_later(max(0.0, deadline - time.monotonic()), _on_run_timeout)
886-
if deadline is not None
887-
else None
888-
)
901+
watchdog = _RunWatchdog(loop, deadline, pipe)
889902

890903
try:
891904
tmpfs_dir.mkdir(parents=True, exist_ok=True)
@@ -923,7 +936,7 @@ def _on_run_timeout() -> None:
923936
on_sample_complete=on_sample_complete,
924937
session_id=session_id,
925938
)
926-
session_ref.append(session)
939+
watchdog.bind_session(session)
927940
phases = _build_phases(ctx, perf_strategy=agentic_inference_strategy)
928941

929942
max_duration_ms = (
@@ -963,7 +976,7 @@ def _on_phase_start(phase: PhaseConfig) -> None:
963976

964977
loop.add_signal_handler(signal.SIGINT, session.stop)
965978
try:
966-
if run_timed_out:
979+
if watchdog.fired:
967980
# Deadline elapsed during setup — never start issuing
968981
# load after it. Run the already-stopped session so
969982
# STARTED/ENDED still flow: the event logger exits only
@@ -978,7 +991,7 @@ def _on_phase_start(phase: PhaseConfig) -> None:
978991
)
979992
session_completed_normally = True
980993
except Exception as e:
981-
if run_timed_out:
994+
if watchdog.fired:
982995
# The watchdog already aborted the run; a teardown race
983996
# can surface here as a generic exception. Fall through
984997
# with an empty session result so finalize still writes
@@ -1004,7 +1017,7 @@ def _on_phase_start(phase: PhaseConfig) -> None:
10041017
# Unifies the clean phase-end path and the abort path — both
10051018
# reach this block. A watchdog abort counts as an abort even
10061019
# when session.run returned normally after session.stop().
1007-
profiler.stop(session_completed_normally and not run_timed_out)
1020+
profiler.stop(session_completed_normally and not watchdog.fired)
10081021
# Graceful drain runs on both the clean-finish and session-
10091022
# failure paths (BenchmarkSession.run publishes ENDED in its own
10101023
# finally, so a failed run still has a terminal snapshot worth
@@ -1058,16 +1071,15 @@ def _on_phase_start(phase: PhaseConfig) -> None:
10581071
)
10591072
raise
10601073
finally:
1061-
if run_watchdog is not None:
1062-
run_watchdog.cancel()
1074+
watchdog.cancel()
10631075

10641076
return BenchmarkResult(
10651077
session=result,
10661078
collector=collector,
10671079
report=report,
10681080
tmpfs_dir=tmpfs_dir,
10691081
profiling=profiler.payload(),
1070-
run_timed_out=run_timed_out,
1082+
run_timed_out=watchdog.fired,
10711083
)
10721084

10731085

0 commit comments

Comments
 (0)