From fd035c1d0518a35cbed89feaf36be628d58b5f4a Mon Sep 17 00:00:00 2001 From: zhaobosen Date: Sun, 9 Aug 2026 17:13:01 +0800 Subject: [PATCH] fix: fail requests when a model worker stops Prevent unresolved futures from accumulating after worker failures, expose unhealthy engines as 503, and make request timeouts opt-in for deployments that need stalled-request detection. Co-Authored-By: Claude Opus 4.6 --- libs/infinity_emb/infinity_emb/engine.py | 4 + libs/infinity_emb/infinity_emb/env.py | 6 + .../infinity_emb/inference/batch_handler.py | 118 ++++++++++++++++-- .../infinity_emb/infinity_server.py | 36 ++++++ .../unit_test/inference/test_batch_handler.py | 92 ++++++++++++++ 5 files changed, 247 insertions(+), 9 deletions(-) diff --git a/libs/infinity_emb/infinity_emb/engine.py b/libs/infinity_emb/infinity_emb/engine.py index 153e15ba7..7e8fc6a8a 100644 --- a/libs/infinity_emb/infinity_emb/engine.py +++ b/libs/infinity_emb/infinity_emb/engine.py @@ -119,6 +119,10 @@ def is_overloaded(self) -> bool: self._assert_running() return self._batch_handler.is_overloaded() + def is_healthy(self) -> bool: + self._assert_running() + return self._batch_handler.is_healthy() + @property def is_running(self) -> bool: return self.running diff --git a/libs/infinity_emb/infinity_emb/env.py b/libs/infinity_emb/infinity_emb/env.py index 48833e473..ade1c0a98 100644 --- a/libs/infinity_emb/infinity_emb/env.py +++ b/libs/infinity_emb/infinity_emb/env.py @@ -193,6 +193,12 @@ def queue_size(self) -> int: assert size > 0, "INFINITY_QUEUE_SIZE must be a positive number" return size + @cached_property + def request_timeout(self) -> float: + timeout = float(self._optional_infinity_var("request_timeout", default="0")) + assert timeout >= 0, "INFINITY_REQUEST_TIMEOUT must not be negative" + return timeout + @cached_property def max_client_batch_size(self) -> int: size = int(self._optional_infinity_var("max_client_batch_size", default="2048")) diff --git a/libs/infinity_emb/infinity_emb/inference/batch_handler.py b/libs/infinity_emb/infinity_emb/inference/batch_handler.py index 1cb48aa9c..4eda85e1f 100644 --- a/libs/infinity_emb/infinity_emb/inference/batch_handler.py +++ b/libs/infinity_emb/infinity_emb/inference/batch_handler.py @@ -46,6 +46,10 @@ QUEUE_TIMEOUT = 0.5 +class EngineUnhealthyError(RuntimeError): + """Raised when a batching engine can no longer complete accepted requests.""" + + class ShutdownReadOnly: def __init__(self, shutdown: threading.Event) -> None: self._shutdown = shutdown @@ -84,6 +88,7 @@ def __init__( vector_disk_cache_path: str = "", verbose=False, lengths_via_tokenize: bool = False, + request_timeout: float = MANAGER.request_timeout, ) -> None: """ performs the scheduling of the dynamic batching around the model. @@ -93,6 +98,8 @@ def __init__( model (BaseTransformer): the base class of the model to be used max_batch_size (int): max batch size of dynamic batch size max_queue_wait (int, optional): max items to queue in the batch, default 32_000 + request_timeout (float, optional): maximum seconds an accepted request may wait; + disabled when set to 0. batch_delay (float, optional): sleep in seconds, wait time for pre/post methods. Best result: setting to 1/2 the minimal expected time for core_encode method / "gpu inference". @@ -103,6 +110,8 @@ def __init__( """ self._max_queue_wait = max_queue_wait + self._request_timeout = request_timeout + assert request_timeout >= 0, "request_timeout must not be negative" self._lengths_via_tokenize = lengths_via_tokenize self._shutdown = threading.Event() @@ -110,6 +119,10 @@ def __init__( self._queue_prio = CustomFIFOQueue() self._publish_to_model_queue: Queue = Queue(8) self._result_queue: Queue = Queue(8) + self._healthy = True + self._failure_lock = threading.Lock() + self._inflight: dict[asyncio.Future, float] = {} + self._watchdog_task: Optional[asyncio.Task] = None self.max_batch_size = max_batch_size self._verbose = verbose @@ -134,6 +147,7 @@ def __init__( threadpool=ThreadPoolExecutorReadOnly(self._threadpool), input_q=self._publish_to_model_queue, output_q=self._result_queue, + on_failure=self._mark_unhealthy, verbose=self._verbose, batch_delay=batch_delay, ) @@ -324,6 +338,14 @@ async def _schedule(self, list_queueitem: Sequence[AbstractSingle]) -> tuple[lis item=inner, ) new_prioqueue.append(item) + if not self.is_healthy(): + raise EngineUnhealthyError("batching engine is unhealthy") + + for item in new_prioqueue: + future = item.item.future + self._inflight[future] = self.loop.time() + future.add_done_callback(self._inflight.pop) + self._queue_prio.extend(new_prioqueue) result = await asyncio.gather( @@ -336,6 +358,58 @@ def capabilities(self) -> set[ModelCapabilites]: # TODO: try to remove inheritance here and return upon init. return self.model_worker[0].capabilities + def is_healthy(self) -> bool: + return self._healthy + + def _fail_inflight(self, error: EngineUnhealthyError) -> None: + for future in list(self._inflight): + if not future.done(): + future.set_exception(error) + + def _mark_unhealthy(self, error: BaseException) -> None: + if self._shutdown.is_set(): + return + with self._failure_lock: + if not self._healthy: + return + self._healthy = False + self._shutdown.set() + + failure = EngineUnhealthyError(f"batching engine failed: {error}") + logger.error("%s", failure) + self.loop.call_soon_threadsafe(self._fail_inflight, failure) + + def _watch_background_worker(self, future) -> None: + if self._shutdown.is_set() or future.cancelled(): + return + error = future.exception() + if error is None: + error = RuntimeError("batching worker stopped unexpectedly") + self._mark_unhealthy(error) + + def _watch_subscriber(self, task: asyncio.Task) -> None: + if self._shutdown.is_set() or task.cancelled(): + return + error = task.exception() + if error is None: + error = RuntimeError("result subscriber stopped unexpectedly") + self._mark_unhealthy(error) + + async def _watch_requests(self) -> None: + interval = min(max(self._request_timeout / 10, 0.1), 1.0) + while not self._shutdown.is_set(): + await asyncio.sleep(interval) + if not self._inflight: + continue + oldest_request = min(self._inflight.values()) + elapsed = self.loop.time() - oldest_request + if elapsed > self._request_timeout: + self._mark_unhealthy( + TimeoutError( + f"request exceeded the {self._request_timeout}s timeout" + ) + ) + def is_overloaded(self) -> bool: """checks if more items can be queued. @@ -443,6 +517,10 @@ async def _subscribe_to_model( raise e continue results, batch = post_batch + if len(results) != len(batch): + raise ValueError( + f"received {len(results)} results for a batch of {len(batch)} items" + ) for i, item in enumerate(batch): await item.complete(results[i]) @@ -456,15 +534,17 @@ async def spawn(self): logger.info("creating batching engine") self.loop = asyncio.get_event_loop() - self._threadpool.submit( - self._publish_towards_model, - ) + publish_task = self._threadpool.submit(self._publish_towards_model) + publish_task.add_done_callback(self._watch_background_worker) self._push_task = asyncio.create_task( self._subscribe_to_model( ShutdownReadOnly(self._shutdown), self._result_queue, self._threadpool ) ) + self._push_task.add_done_callback(self._watch_subscriber) + if self._request_timeout: + self._watchdog_task = asyncio.create_task(self._watch_requests()) for worker in self.model_worker: worker.spawn() @@ -474,9 +554,14 @@ async def shutdown(self): Blocking event, until shutdown complete. reverses .spawn() """ + was_healthy = self.is_healthy() self._shutdown.set() - await asyncio.to_thread(self._threadpool.shutdown) - # collect task + if self._watchdog_task: + self._watchdog_task.cancel() + if was_healthy: + await asyncio.to_thread(self._threadpool.shutdown) + else: + self._threadpool.shutdown(wait=False, cancel_futures=True) self._push_task.cancel() @@ -490,6 +575,7 @@ def __init__( threadpool: ThreadPoolExecutorReadOnly, input_q: Queue, output_q: Queue, + on_failure, batch_delay: float = 5e-3, verbose=False, ) -> None: @@ -501,6 +587,7 @@ def __init__( self._batch_delay = float(max(1e-4, batch_delay)) self._input_q = input_q self._output_q = output_q + self._on_failure = on_failure self._last_inference = time.perf_counter() self._verbose = verbose self._ready = False @@ -508,10 +595,23 @@ def __init__( def spawn(self): if self._ready: raise ValueError("already spawned") - # start the threads - self._threadpool.submit(self._preprocess_batch) - self._threadpool.submit(self._core_batch) - self._threadpool.submit(self._postprocess_batch) + # Start and supervise each stage: an uncaught exception otherwise only + # terminates the ThreadPoolExecutor task and leaves requests unresolved. + for stage in ( + self._preprocess_batch, + self._core_batch, + self._postprocess_batch, + ): + task = self._threadpool.submit(stage) + task.add_done_callback(self._watch_stage) + + def _watch_stage(self, future) -> None: + if self._shutdown.is_set() or future.cancelled(): + return + error = future.exception() + if error is None: + error = RuntimeError("model worker stage stopped unexpectedly") + self._on_failure(error) @property def capabilities(self) -> set[ModelCapabilites]: diff --git a/libs/infinity_emb/infinity_emb/infinity_server.py b/libs/infinity_emb/infinity_emb/infinity_server.py index 382cc6e8e..d2d2b4f9e 100644 --- a/libs/infinity_emb/infinity_emb/infinity_server.py +++ b/libs/infinity_emb/infinity_emb/infinity_server.py @@ -13,6 +13,7 @@ import infinity_emb from infinity_emb.args import EngineArgs from infinity_emb.engine import AsyncEmbeddingEngine, AsyncEngineArray +from infinity_emb.inference.batch_handler import EngineUnhealthyError from infinity_emb.env import MANAGER from infinity_emb.fastapi_schemas import docs, errors from infinity_emb.log_handler import logger @@ -167,6 +168,11 @@ async def _health() -> dict[str, float]: Returns: dict(unix=float): dict with unix time stamp """ + if not all(engine.is_healthy() for engine in app.engine_array.engines_dict.values()): + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="one or more model engines are unavailable", + ) return {"unix": time.time()} if redirect_slash: @@ -221,6 +227,11 @@ def _resolve_engine(model: str) -> "AsyncEmbeddingEngine": f"Invalid model: {ex}", code=status.HTTP_400_BAD_REQUEST, ) + if not engine.is_healthy(): + raise errors.OpenAIException( + f"model {model} is unavailable", + code=status.HTTP_503_SERVICE_UNAVAILABLE, + ) if engine.is_overloaded(): raise errors.OpenAIException( f"model {model} is currently overloaded", @@ -387,6 +398,11 @@ def url_to_base64(url, modality = "image"): f"{ex.__class__} -> {ex}", code=status.HTTP_400_BAD_REQUEST, ) + except EngineUnhealthyError as ex: + raise errors.OpenAIException( + str(ex), + code=status.HTTP_503_SERVICE_UNAVAILABLE, + ) except Exception as ex: raise errors.OpenAIException( f"InternalServerError: {ex}", @@ -439,6 +455,11 @@ async def _rerank(data: RerankInput): f"ModelNotDeployedError: model=`{data.model}` does not support `rerank`. Reason: {ex}", code=status.HTTP_400_BAD_REQUEST, ) + except EngineUnhealthyError as ex: + raise errors.OpenAIException( + str(ex), + code=status.HTTP_503_SERVICE_UNAVAILABLE, + ) except Exception as ex: raise errors.OpenAIException( f"InternalServerError: {ex}", @@ -483,6 +504,11 @@ async def _classify(data: ClassifyInput): f"ModelNotDeployedError: model=`{data.model}` does not support `classify`. Reason: {ex}", code=status.HTTP_400_BAD_REQUEST, ) + except EngineUnhealthyError as ex: + raise errors.OpenAIException( + str(ex), + code=status.HTTP_503_SERVICE_UNAVAILABLE, + ) except Exception as ex: raise errors.OpenAIException( f"InternalServerError: {ex}", @@ -542,6 +568,11 @@ async def _embeddings_image(data: ImageEmbeddingInput): f"ModelNotDeployedError: model=`{data.model}` does not support `image_embed`. Reason: {ex}", code=status.HTTP_400_BAD_REQUEST, ) + except EngineUnhealthyError as ex: + raise errors.OpenAIException( + str(ex), + code=status.HTTP_503_SERVICE_UNAVAILABLE, + ) except Exception as ex: raise errors.OpenAIException( f"InternalServerError: {ex}", @@ -601,6 +632,11 @@ async def _embeddings_audio(data: AudioEmbeddingInput): f"ModelNotDeployedError: model=`{data.model}` does not support `audio_embed`. Reason: {ex}", code=status.HTTP_400_BAD_REQUEST, ) + except EngineUnhealthyError as ex: + raise errors.OpenAIException( + str(ex), + code=status.HTTP_503_SERVICE_UNAVAILABLE, + ) except Exception as ex: raise errors.OpenAIException( f"InternalServerError: {ex}", diff --git a/libs/infinity_emb/tests/unit_test/inference/test_batch_handler.py b/libs/infinity_emb/tests/unit_test/inference/test_batch_handler.py index 3a36ac34a..6f472b17a 100644 --- a/libs/infinity_emb/tests/unit_test/inference/test_batch_handler.py +++ b/libs/infinity_emb/tests/unit_test/inference/test_batch_handler.py @@ -2,6 +2,7 @@ import copy import random import sys +import threading import time import numpy as np @@ -10,6 +11,7 @@ from infinity_emb.args import EngineArgs from infinity_emb.inference import BatchHandler +from infinity_emb.inference.batch_handler import EngineUnhealthyError from infinity_emb.transformer.embedder.sentence_transformer import ( SentenceTransformerPatched, ) @@ -20,6 +22,39 @@ MODEL_NAME: str = pytest.DEFAULT_BERT_MODEL # type: ignore[assignment] +class FailingCoreModel: + capabilities = {"embed"} + + def encode_pre(self, inputs): + return inputs + + def encode_core(self, features): + raise RuntimeError("simulated core failure") + + def encode_post(self, features): + return features + + def tokenize_lengths(self, sentences): + return [1] * len(sentences) + + +class BlockingCoreModel(FailingCoreModel): + def __init__(self): + self.release = threading.Event() + + def encode_core(self, features): + self.release.wait() + return features + + +class WrongResultModel(FailingCoreModel): + def encode_core(self, features): + return features + + def encode_post(self, features): + return [] + + @pytest.fixture @pytest.mark.anyio @pytest.mark.skipif( @@ -39,6 +74,63 @@ async def load_patched_bh() -> tuple[SentenceTransformerPatched, BatchHandler]: return model, bh +@pytest.mark.anyio +async def test_core_worker_failure_completes_waiting_requests(): + handler = BatchHandler( + model_replicas=[FailingCoreModel()], + max_batch_size=1, + batch_delay=1e-4, + lengths_via_tokenize=True, + ) + await handler.spawn() + + try: + with pytest.raises(EngineUnhealthyError, match="simulated core failure"): + await asyncio.wait_for(handler.embed(["request"]), timeout=2) + assert not handler.is_healthy() + finally: + await handler.shutdown() + + +@pytest.mark.anyio +async def test_request_timeout_completes_waiting_requests(): + model = BlockingCoreModel() + handler = BatchHandler( + model_replicas=[model], + max_batch_size=1, + request_timeout=0.05, + batch_delay=1e-4, + lengths_via_tokenize=True, + ) + await handler.spawn() + + try: + with pytest.raises(EngineUnhealthyError, match="request exceeded"): + await asyncio.wait_for(handler.embed(["request"]), timeout=2) + assert not handler.is_healthy() + finally: + model.release.set() + await handler.shutdown() + + +@pytest.mark.anyio +async def test_result_length_mismatch_completes_waiting_requests(): + handler = BatchHandler( + model_replicas=[WrongResultModel()], + max_batch_size=1, + batch_delay=1e-4, + lengths_via_tokenize=True, + ) + await handler.spawn() + + try: + with pytest.raises(EngineUnhealthyError, match="results for a batch"): + await asyncio.wait_for(handler.embed(["request"]), timeout=2) + assert not handler.is_healthy() + finally: + await handler.shutdown() + + @pytest.mark.performance @pytest.mark.anyio async def test_batch_performance_raw(get_sts_bechmark_dataset, load_patched_bh):