Feat: decode request in threadpool

src/litserve/loops.py

@@ -12,14 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import asyncio
+import concurrent.futures
 import inspect
 import logging
 import multiprocessing as mp
+import os
 import pickle
 import sys
 import time
 from queue import Empty, Queue
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
 from fastapi import HTTPException
 from starlette.formparsers import MultiPartParser
@@ -55,6 +57,20 @@ def _inject_context(context: Union[List[dict], dict], func, *args, **kwargs):
     return func(*args, **kwargs)
 
 
+def decode_requests_in_threadpool(
+    executor: concurrent.futures.ThreadPoolExecutor,
+    contexts: List[Union[List[dict], dict]],
+    func: Callable,
+    inputs: List[dict],
+):
+    x = [executor.submit(_inject_context, context, func, input) for input, context in zip(inputs, contexts)]
+
+    for i, _x in enumerate(x):
+        x[i] = _x.result()
+
+    return x
+
+
 def collate_requests(
     lit_api: LitAPI, request_queue: Queue, max_batch_size: int, batch_timeout: float
 ) -> Tuple[List, List]:
@@ -165,7 +181,11 @@ def run_batched_loop(
     max_batch_size: int,
     batch_timeout: float,
     callback_runner: CallbackRunner,
+    concurrent_decode: bool = True,
 ):
+    if concurrent_decode:
+        executor = concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count())
+
     while True:
         batches, timed_out_uids = collate_requests(
             lit_api,
@@ -193,14 +213,17 @@ def run_batched_loop(
                     lit_spec.populate_context(context, input)
 
             callback_runner.trigger_event(EventTypes.BEFORE_DECODE_REQUEST, lit_api=lit_api)
-            x = [
-                _inject_context(
-                    context,
-                    lit_api.decode_request,
-                    input,
-                )
-                for input, context in zip(inputs, contexts)
-            ]
+            if concurrent_decode:
+                x = decode_requests_in_threadpool(executor, contexts, lit_api.decode_request, inputs)
+            else:
+                x = [
+                    _inject_context(
+                        context,
+                        lit_api.decode_request,
+                        input,
+                    )
+                    for input, context in zip(inputs, contexts)
+                ]
             callback_runner.trigger_event(EventTypes.AFTER_DECODE_REQUEST, lit_api=lit_api)
 
             x = lit_api.batch(x)
@@ -300,7 +323,11 @@ def run_batched_streaming_loop(
     max_batch_size: int,
     batch_timeout: float,
     callback_runner: CallbackRunner,
+    concurrent_decode: bool = True,
 ):
+    if concurrent_decode:
+        executor = concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count())
+
     while True:
         batches, timed_out_uids = collate_requests(
             lit_api,
@@ -326,14 +353,17 @@ def run_batched_streaming_loop(
                     lit_spec.populate_context(context, input)
 
             callback_runner.trigger_event(EventTypes.BEFORE_DECODE_REQUEST, lit_api=lit_api)
-            x = [
-                _inject_context(
-                    context,
-                    lit_api.decode_request,
-                    input,
-                )
-                for input, context in zip(inputs, contexts)
-            ]
+            if concurrent_decode:
+                x = decode_requests_in_threadpool(executor, contexts, lit_api.decode_request, inputs)
+            else:
+                x = [
+                    _inject_context(
+                        context,
+                        lit_api.decode_request,
+                        input,
+                    )
+                    for input, context in zip(inputs, contexts)
+                ]
             callback_runner.trigger_event(EventTypes.AFTER_DECODE_REQUEST, lit_api=lit_api)
 
             x = lit_api.batch(x)
@@ -378,6 +408,7 @@ def inference_worker(
     stream: bool,
     workers_setup_status: Dict[str, bool],
     callback_runner: CallbackRunner,
+    concurrent_decode: bool = True,
 ):
     callback_runner.trigger_event(EventTypes.BEFORE_SETUP, lit_api=lit_api)
     lit_api.setup(device)
@@ -394,15 +425,29 @@ def inference_worker(
     if stream:
         if max_batch_size > 1:
             run_batched_streaming_loop(
-                lit_api, lit_spec, request_queue, response_queues, max_batch_size, batch_timeout, callback_runner
+                lit_api,
+                lit_spec,
+                request_queue,
+                response_queues,
+                max_batch_size,
+                batch_timeout,
+                callback_runner,
+                concurrent_decode,
             )
         else:
             run_streaming_loop(lit_api, lit_spec, request_queue, response_queues, callback_runner)
         return
 
     if max_batch_size > 1:
         run_batched_loop(
-            lit_api, lit_spec, request_queue, response_queues, max_batch_size, batch_timeout, callback_runner
+            lit_api,
+            lit_spec,
+            request_queue,
+            response_queues,
+            max_batch_size,
+            batch_timeout,
+            callback_runner,
+            concurrent_decode,
         )
     else:
         run_single_loop(lit_api, lit_spec, request_queue, response_queues, callback_runner)

src/litserve/server.py

@@ -118,6 +118,7 @@ def __init__(
         callbacks: Optional[Union[List[Callback], Callback]] = None,
         middlewares: Optional[list[Union[Callable, tuple[Callable, dict]]]] = None,
         loggers: Optional[Union[Logger, List[Logger]]] = None,
+        concurrent_decode: bool = True,
     ):
         if batch_timeout > timeout and timeout not in (False, -1):
             raise ValueError("batch_timeout must be less than timeout")
@@ -178,6 +179,7 @@ def __init__(
         self.max_payload_size = max_payload_size
         self._connector = _Connector(accelerator=accelerator, devices=devices)
         self._callback_runner = CallbackRunner(callbacks)
+        self._concurrent_decode = concurrent_decode
 
         specs = spec if spec is not None else []
         self._specs = specs if isinstance(specs, Sequence) else [specs]
@@ -246,6 +248,7 @@ def launch_inference_worker(self, num_uvicorn_servers: int):
                     self.stream,
                     self.workers_setup_status,
                     self._callback_runner,
+                    self._concurrent_decode,
                 ),
             )
             process.start()