diff --git a/freerelay/core/execution/executor.py b/freerelay/core/execution/executor.py index 085e053..dc76332 100644 --- a/freerelay/core/execution/executor.py +++ b/freerelay/core/execution/executor.py @@ -151,11 +151,7 @@ async def execute_hedged( provider_keys = [(p, k) for p, k, _ in providers] circuits = {p.name: c for p, _, c in providers} try: - response = await hedged_execute(provider_keys, request) - # Record success on the circuit breaker of whichever provider won - for _name, circuit in circuits.items(): - await circuit.record_success() - return response + result = await hedged_execute(provider_keys, request) except Exception as e: # Record failure on all involved circuit breakers for _name, circuit in circuits.items(): @@ -163,6 +159,26 @@ async def execute_hedged( await circuit.record_failure(status) raise + # Only the winning provider's circuit breaker should observe a + # success. The loser's provider either timed out, raised, or + # completed after the winner — calling record_success on it + # would mask the failure that triggered the hedge in the first + # place and prevent the circuit breaker from ever tripping on + # that provider. + winner_circuit = circuits.get(result.winner_name) + if winner_circuit is not None: + await winner_circuit.record_success() + for loser_name in result.loser_names: + loser_circuit = circuits.get(loser_name) + if loser_circuit is not None: + # The loser completed but slower than the winner, so its + # response was discarded — not strictly a "failure" but + # also not a success the breaker should credit. Record a + # non-status failure (None) to nudge the breaker toward + # opening without the aggressive weight of a real 5xx. + await loser_circuit.record_failure(None) + return result.response + async def execute_stream( self, provider: BaseProvider, diff --git a/freerelay/core/execution/hedging.py b/freerelay/core/execution/hedging.py index 39a40bc..5057581 100644 --- a/freerelay/core/execution/hedging.py +++ b/freerelay/core/execution/hedging.py @@ -21,10 +21,32 @@ logger = logging.getLogger("freerelay.hedging") +class HedgedResult: + """Result of a hedged execution, identifying which provider won.""" + + __slots__ = ("response", "winner_name", "loser_names") + + def __init__( + self, + response: ChatCompletionResponse, + winner_name: str, + loser_names: tuple[str, ...], + ) -> None: + self.response = response + self.winner_name = winner_name + self.loser_names = loser_names + + def __repr__(self) -> str: + return ( + f"HedgedResult(winner={self.winner_name!r}, " + f"losers={self.loser_names!r})" + ) + + async def hedged_execute( providers: list[tuple[BaseProvider, str]], request: ChatCompletionRequest, -) -> ChatCompletionResponse: +) -> HedgedResult: """ Fire the same request at up to 2 providers in parallel. Return the first response. Cancel all others. @@ -34,10 +56,13 @@ async def hedged_execute( request: The chat completion request. Returns: - Response from the fastest provider. + HedgedResult with the winning response, the winner's name, and + the names of any losing providers (so callers can update their + circuit breakers accordingly). Raises: - Exception: If all providers fail. + ValueError: If no providers were given. + Exception: If all providers fail. The first error is re-raised. """ if not providers: raise ValueError("No providers for hedged execution") @@ -55,22 +80,53 @@ async def hedged_execute( return_when=asyncio.FIRST_COMPLETED, ) - # Cancel all pending tasks immediately + # Cancel all pending tasks immediately. Use CancelledError-specific + # suppression: when a task is cancelled, awaiting it raises + # CancelledError. A bare Exception in the suppress tuple would also + # swallow the original failure if the task had raised before we got + # here, which would mask real provider errors in our logs. for task in pending: task.cancel() - with contextlib.suppress(asyncio.CancelledError, Exception): + with contextlib.suppress(asyncio.CancelledError): await task - # Check all done tasks for a successful result - # done may contain 1 or 2 tasks (both completed before cancellation) + # Walk done tasks once. If more than one finished before we could + # cancel (e.g. both providers returned on the same scheduler tick), + # the first task we pick with no exception is the winner; the rest + # are losers that need their result fetched (or exception ignored) + # so the asyncio task machinery doesn't warn about un-awaited tasks. + winner_task: asyncio.Task[ChatCompletionResponse] | None = None + winner_provider: BaseProvider | None = None + loser_providers: list[BaseProvider] = [] errors: list[BaseException] = [] + for task in done: - if task.exception() is None: - winner = tasks[task] - logger.info("Hedged winner: %s", winner.name) - # Cancel any remaining done tasks that weren't checked yet - return task.result() - errors.append(task.exception()) # type: ignore[arg-type] + exc = task.exception() + if exc is None and winner_task is None: + winner_task = task + winner_provider = tasks[task] + elif exc is None: + loser_providers.append(tasks[task]) + else: + errors.append(exc) + + if winner_task is not None: + logger.info("Hedged winner: %s", winner_provider.name) + # Drain any leftover successful losers so their task objects + # aren't GC'd in a 'pending result' state — that's a common + # asyncio footgun that produces "Task was destroyed but it is + # pending!" warnings at shutdown. + for loser in loser_providers: + for t in done: + if tasks.get(t) is loser: + with contextlib.suppress(asyncio.CancelledError, Exception): + t.result() + break + return HedgedResult( + response=winner_task.result(), + winner_name=winner_provider.name, + loser_names=tuple(p.name for p in loser_providers), + ) # All done tasks failed logger.warning("All hedged providers failed (%d errors)", len(errors)) diff --git a/tests/unit/test_hedging_winner_only_success.py b/tests/unit/test_hedging_winner_only_success.py new file mode 100644 index 0000000..41b59c8 --- /dev/null +++ b/tests/unit/test_hedging_winner_only_success.py @@ -0,0 +1,121 @@ +""" +Tests — Hedged execution only records circuit-breaker success on the winner. + +The hedged execution path fires a request at two providers in parallel +and returns whichever responds first. The losing provider either gets +cancelled mid-flight or completes after the winner — either way it did +NOT successfully serve the user, so its circuit breaker must NOT be +reset to CLOSED. Only the winner's circuit should record success. +""" + +from __future__ import annotations + +import asyncio +from unittest.mock import MagicMock + +import pytest + +from freerelay.core.execution.executor import Executor +from freerelay.core.models.openai import ( + ChatCompletionRequest, + ChatCompletionResponse, + Choice, + Message, +) +from freerelay.core.resilience.circuit_breaker import CircuitBreaker, CircuitState + + +def _make_provider(name: str, delay: float, fail: bool = False) -> MagicMock: + provider = MagicMock() + provider.name = name + + async def _complete(request, api_key): + await asyncio.sleep(delay) + if fail: + raise RuntimeError(f"{name} boom") + return ChatCompletionResponse( + id=f"resp-{name}", + created=0, + model=f"model-{name}", + choices=[ + Choice( + index=0, + message=Message(role="assistant", content=f"{name}-wins"), + finish_reason="stop", + ) + ], + ) + + provider.complete = _complete + return provider + + +def _make_request() -> ChatCompletionRequest: + req = MagicMock(spec=ChatCompletionRequest) + req.model = "test-model" + req.messages = [] + return req + + +@pytest.mark.asyncio +async def test_hedged_winner_records_success_only_on_winner() -> None: + """Provider A is faster; A's circuit becomes CLOSED but B's stays OPEN + because B was cancelled mid-flight (its request never completed).""" + a_breaker = CircuitBreaker( + provider_name="a", failure_threshold=3, failure_window=60, recovery_timeout=0.5 + ) + b_breaker = CircuitBreaker( + provider_name="b", failure_threshold=3, failure_window=60, recovery_timeout=0.5 + ) + # Pre-open both circuits with 3 failures to mimic a struggling provider. + for _ in range(3): + await a_breaker.record_failure(500) + await b_breaker.record_failure(500) + assert a_breaker.state == CircuitState.OPEN + assert b_breaker.state == CircuitState.OPEN + + # Wait past the recovery window so the auto-transition can fire when + # record_success is called. + await asyncio.sleep(0.6) + + a = _make_provider("A", delay=0.05) + b = _make_provider("B", delay=0.50) # loser (will be cancelled) + executor = Executor(enable_hedging=True, max_retries=0) + result = await executor.execute_hedged( + [(a, "key-a", a_breaker), (b, "key-b", b_breaker)], + _make_request(), + ) + + assert result.choices[0].message.content == "A-wins" + # The winner's circuit went OPEN → HALF_OPEN → CLOSED with failures cleared. + assert a_breaker.state == CircuitState.CLOSED, ( + f"winner A should have success recorded and circuit reset to CLOSED, got {a_breaker.state}" + ) + assert a_breaker.get_score() == 1.0 + # The loser's circuit was NOT marked successful. + # It auto-transitioned to HALF_OPEN during the sleep, but the loser's task + # was cancelled, so no failure was recorded — the circuit should remain + # HALF_OPEN (a probe slot is still open for it to retry). It must NOT + # be CLOSED, because that would mean we silently treated the loser as healthy. + assert b_breaker.state != CircuitState.CLOSED, ( + f"loser B must NOT be marked successful; got {b_breaker.state}" + ) + + +@pytest.mark.asyncio +async def test_hedged_all_fail_records_failure_on_both() -> None: + """When every provider fails, every circuit breaker should see the failure.""" + a = _make_provider("A", delay=0.01, fail=True) + b = _make_provider("B", delay=0.02, fail=True) + circuit_a = CircuitBreaker("A", failure_threshold=5) + circuit_b = CircuitBreaker("B", failure_threshold=5) + executor = Executor(enable_hedging=True, max_retries=0) + with pytest.raises(RuntimeError): + await executor.execute_hedged( + [(a, "key-a", circuit_a), (b, "key-b", circuit_b)], + _make_request(), + ) + + # Both should still be CLOSED (threshold=5, only 1 failure each) + assert circuit_a.state == CircuitState.CLOSED + assert circuit_b.state == CircuitState.CLOSED