Feat/evict req on client disconnect streaming case

src/litserve/server.py

@@ -217,7 +217,13 @@ def run_batched_loop(
                 response_queues[response_queue_id].put((uid, (err_pkl, LitAPIStatus.ERROR)))
 
 
-def run_streaming_loop(lit_api: LitAPI, lit_spec: LitSpec, request_queue: Queue, response_queues: List[Queue]):
+def run_streaming_loop(
+    lit_api: LitAPI,
+    lit_spec: LitSpec,
+    request_queue: Queue,
+    response_queues: List[Queue],
+    request_evicted_status: Dict[str, bool],
+):
     while True:
         try:
             response_queue_id, uid, timestamp, x_enc = request_queue.get(timeout=1.0)
@@ -256,6 +262,9 @@ def run_streaming_loop(lit_api: LitAPI, lit_spec: LitSpec, request_queue: Queue,
                 y_gen,
             )
             for y_enc in y_enc_gen:
+                if request_evicted_status.get(uid):
+                    request_evicted_status.pop(uid)
+                    break
                 y_enc = lit_api.format_encoded_response(y_enc)
                 response_queues[response_queue_id].put((uid, (y_enc, LitAPIStatus.OK)))
             response_queues[response_queue_id].put((uid, ("", LitAPIStatus.FINISH_STREAMING)))
@@ -342,6 +351,7 @@ def inference_worker(
     batch_timeout: float,
     stream: bool,
     workers_setup_status: Dict[str, bool] = None,
+    request_evicted_status: Dict[str, bool] = None,
 ):
     lit_api.setup(device)
     lit_api.device = device
@@ -357,7 +367,7 @@ def inference_worker(
         if max_batch_size > 1:
             run_batched_streaming_loop(lit_api, lit_spec, request_queue, response_queues, max_batch_size, batch_timeout)
         else:
-            run_streaming_loop(lit_api, lit_spec, request_queue, response_queues)
+            run_streaming_loop(lit_api, lit_spec, request_queue, response_queues, request_evicted_status)
         return
 
     if max_batch_size > 1:
@@ -384,27 +394,27 @@ def api_key_auth(x_api_key: str = Depends(APIKeyHeader(name="X-API-Key"))):
 
 async def response_queue_to_buffer(
     response_queue: mp.Queue,
-    buffer: Dict[str, Union[Tuple[deque, asyncio.Event], asyncio.Event]],
+    response_buffer: Dict[str, Union[Tuple[deque, asyncio.Event], asyncio.Event]],
     stream: bool,
-    response_executor: ThreadPoolExecutor,
+    threadpool: ThreadPoolExecutor,
 ):
     loop = asyncio.get_running_loop()
     if stream:
         while True:
             try:
-                uid, payload = await loop.run_in_executor(response_executor, response_queue.get)
+                uid, response = await loop.run_in_executor(threadpool, response_queue.get)
             except Empty:
                 await asyncio.sleep(0.0001)
                 continue
-            q, event = buffer[uid]
-            q.append(payload)
+            stream_response_buffer, event = response_buffer[uid]
+            stream_response_buffer.append((uid, response))
             event.set()
 
     else:
         while True:
-            uid, payload = await loop.run_in_executor(response_executor, response_queue.get)
-            event = buffer.pop(uid)
-            buffer[uid] = payload
+            uid, response = await loop.run_in_executor(threadpool, response_queue.get)
+            event = response_buffer.pop(uid)
+            response_buffer[uid] = response
             event.set()
 
 
@@ -498,6 +508,7 @@ def launch_inference_worker(self, num_uvicorn_servers: int):
         manager = mp.Manager()
         self.workers_setup_status = manager.dict()
         self.request_queue = manager.Queue()
+        self.request_evicted_status = manager.dict()
 
         self.response_queues = []
         for _ in range(num_uvicorn_servers):
@@ -535,6 +546,7 @@ def launch_inference_worker(self, num_uvicorn_servers: int):
                     self.batch_timeout,
                     self.stream,
                     self.workers_setup_status,
+                    self.request_evicted_status,
                 ),
             )
             process.start()
@@ -568,26 +580,33 @@ def device_identifiers(self, accelerator, device):
         return [f"{accelerator}:{device}"]
 
     async def data_streamer(self, q: deque, data_available: asyncio.Event, send_status: bool = False):
+        uid = None
         while True:
-            await data_available.wait()
-            while len(q) > 0:
-                data, status = q.popleft()
-                if status == LitAPIStatus.FINISH_STREAMING:
-                    return
-
-                if status == LitAPIStatus.ERROR:
-                    logger.error(
-                        "Error occurred while streaming outputs from the inference worker. "
-                        "Please check the above traceback."
-                    )
+            try:
+                await data_available.wait()
+                while len(q) > 0:
+                    uid, (data, status) = q.popleft()
+                    if status == LitAPIStatus.FINISH_STREAMING:
+                        return
+
+                    if status == LitAPIStatus.ERROR:
+                        logger.error(
+                            "Error occurred while streaming outputs from the inference worker. "
+                            "Please check the above traceback."
+                        )
+                        if send_status:
+                            yield data, status
+                        return
                     if send_status:
                         yield data, status
-                    return
-                if send_status:
-                    yield data, status
-                else:
-                    yield data
-            data_available.clear()
+                    else:
+                        yield data
+                data_available.clear()
+            except asyncio.CancelledError:
+                if uid is not None:
+                    self.request_evicted_status[uid] = True
+                    logger.exception("Streaming request cancelled for the uid=%s", uid)
+                return
 
     def setup_server(self):
         workers_ready = False

tests/conftest.py

@@ -55,6 +55,14 @@ def encode_response(self, output: Generator) -> Generator:
             yield out.lower()
 
 
+class SimpleDelayedStreamAPI(SimpleStreamAPI):
+    def encode_response(self, output: Generator) -> Generator:
+        delay = 0.2
+        for out in output:
+            time.sleep(delay)
+            yield out.lower()
+
+
 class SimpleBatchedStreamAPI(LitAPI):
     def setup(self, device) -> None:
         self.sentence = "LitServe is streaming output"
@@ -98,6 +106,11 @@ def simple_batched_stream_api():
     return SimpleBatchedStreamAPI()
 
 
+@pytest.fixture()
+def simple_delayed_stream_api():
+    return SimpleDelayedStreamAPI()
+
+
 @pytest.fixture()
 def lit_server(simple_litapi):
     server = LitServer(simple_litapi, accelerator="cpu", devices=1, timeout=10)

tests/test_lit_server.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import asyncio
 import inspect
+import logging
 import pickle
 import re
 from asgi_lifespan import LifespanManager
@@ -120,6 +121,27 @@ async def test_stream(simple_stream_api):
             ), "Server returns input prompt and generated output which didn't match."
 
 
+@pytest.mark.asyncio()
+async def test_stream_client_disconnection(simple_delayed_stream_api, caplog):
+    server = LitServer(simple_delayed_stream_api, stream=True, timeout=10)
+
+    with wrap_litserve_start(server) as server, caplog.at_level(logging.DEBUG):
+        async with LifespanManager(server.app) as manager, AsyncClient(app=manager.app, base_url="http://test") as ac:
+            task = asyncio.create_task(ac.post("/predict", json={"prompt": "Hey, How are you doing?"}, timeout=10))
+            await asyncio.sleep(2)
+            task.cancel()  # simulate client disconnection
+            await asyncio.sleep(1)  # wait for the task to stop
+            with pytest.raises(asyncio.CancelledError):
+                await task
+            assert "Streaming request cancelled for the uid=" in caplog.text
+            # TODO: also check if the task actually stopped in the server
+
+            caplog.clear()
+            task = asyncio.create_task(ac.post("/predict", json={"prompt": "Hey, How are you doing?"}, timeout=10))
+            await task
+            assert "Streaming request cancelled for the uid=" not in caplog.text
+
+
 @pytest.mark.asyncio()
 async def test_batched_stream_server(simple_batched_stream_api):
     server = LitServer(simple_batched_stream_api, stream=True, max_batch_size=4, batch_timeout=2, timeout=30)
@@ -177,9 +199,10 @@ def fake_encode(output):
     requests_queue = Queue()
     requests_queue.put((0, "UUID-1234", time.monotonic(), {"prompt": "Hello"}))
     response_queues = [FakeStreamResponseQueue(num_streamed_outputs)]
+    request_evicted_status = {}
 
     with pytest.raises(StopIteration, match="exit loop"):
-        run_streaming_loop(fake_stream_api, fake_stream_api, requests_queue, response_queues)
+        run_streaming_loop(fake_stream_api, fake_stream_api, requests_queue, response_queues, request_evicted_status)
 
     fake_stream_api.predict.assert_called_once_with("Hello")
     fake_stream_api.encode_response.assert_called_once()