litserve 0.2.8.dev0__tar.gz → 0.2.9.dev0__tar.gz

{litserve-0.2.8.dev0/src/litserve.egg-info → litserve-0.2.9.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: litserve
-Version: 0.2.8.dev0
+Version: 0.2.9.dev0
 Summary: Lightweight AI server.
 Home-page: https://github.com/Lightning-AI/litserve
 Download-URL: https://github.com/Lightning-AI/litserve
@@ -30,7 +30,6 @@ License-File: LICENSE
 Requires-Dist: fastapi>=0.100
 Requires-Dist: uvicorn[standard]>=0.29.0
 Requires-Dist: pyzmq>=22.0.0
-Requires-Dist: starlette
 Provides-Extra: test
 Requires-Dist: asgi-lifespan; extra == "test"
 Requires-Dist: coverage[toml]>=7.5.3; extra == "test"
@@ -71,19 +70,14 @@ Dynamic: summary
 
 <div align='center'>
 
-# Easily serve AI models Lightning fast ⚡    
+# Deploy AI models and inference pipelines - ⚡ fast    
 
 <img alt="Lightning" src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/ls_banner2.png" width="800px" style="max-width: 100%;">
 
-&nbsp;
-
-<strong>Lightning-fast serving engine for AI models.</strong>    
-Easy. Flexible. Enterprise-scale.    
+&nbsp; 
 </div>
 
-----
-
-**LitServe** is an easy-to-use, flexible serving engine for AI models built on FastAPI. It augments FastAPI with features like batching, streaming, and GPU autoscaling eliminate the need to rebuild a FastAPI server per model.  
+**LitServe** lets you build high-performance AI inference pipelines on top of FastAPI - no boilerplate. Define one or more models, connect vector DBs, stream responses, batch requests, and autoscale on GPUs out of the box.
 
 LitServe is at least [2x faster](#performance) than plain FastAPI due to AI-specific multi-worker handling.    
 
@@ -93,12 +87,13 @@ LitServe is at least [2x faster](#performance) than plain FastAPI due to AI-spec
 ✅ (2x)+ faster serving  ✅ Easy to use               ✅ LLMs, non LLMs and more
 ✅ Bring your own model  ✅ PyTorch/JAX/TF/...        ✅ Built on FastAPI       
 ✅ GPU autoscaling       ✅ Batching, Streaming       ✅ Self-host or ⚡️ managed
-✅ Compound AI           ✅ Integrate with vLLM, etc  ✅ Serverless             
+✅ Inference pipeline    ✅ Integrate with vLLM, etc  ✅ Serverless             
    
 </pre>
 
 <div align='center'>
 
+[![PyPI Downloads](https://static.pepy.tech/badge/litserve)](https://pepy.tech/projects/litserve)
 [![Discord](https://img.shields.io/discord/1077906959069626439?label=Get%20help%20on%20Discord)](https://discord.gg/WajDThKAur)
 ![cpu-tests](https://github.com/Lightning-AI/litserve/actions/workflows/ci-testing.yml/badge.svg)
 [![codecov](https://codecov.io/gh/Lightning-AI/litserve/graph/badge.svg?token=SmzX8mnKlA)](https://codecov.io/gh/Lightning-AI/litserve)
@@ -136,16 +131,16 @@ pip install litserve
 ```
     
 ### Define a server    
-This toy example with 2 models (AI compound system) shows LitServe's flexibility ([see real examples](#examples)):    
+This toy example with 2 models (inference pipeline) shows LitServe's flexibility ([see real examples](#featured-examples)):    
 
 ```python
 # server.py
 import litserve as ls
 
-# (STEP 1) - DEFINE THE API (compound AI system)
+# (STEP 1) - DEFINE THE API ("inference" pipeline)
 class SimpleLitAPI(ls.LitAPI):
     def setup(self, device):
-        # setup is called once at startup. Build a compound AI system (1+ models), connect DBs, load data, etc...
+        # setup is called once at startup. Defines elements of the pipeline: models, connect DBs, load data, etc...
         self.model1 = lambda x: x**2
         self.model2 = lambda x: x**3
 
@@ -154,11 +149,11 @@ class SimpleLitAPI(ls.LitAPI):
         return request["input"] 
 
     def predict(self, x):
-        # Easily build compound systems. Run inference and return the output.
-        squared = self.model1(x)
-        cubed = self.model2(x)
-        output = squared + cubed
-        return {"output": output}
+        # Run the inference pipeline and return the output
+        a = self.model1(x)
+        b = self.model2(x)
+        c = a + b
+        return {"output": c}
 
     def encode_response(self, output):
         # Convert the model output to a response payload.
@@ -189,12 +184,7 @@ python server.py
 ```
 
 ### Test the server
-Run the auto-generated test client:        
-```bash
-python client.py    
-```
-
-Or use this terminal command:
+Simulate an http request (run this on any terminal):
 ```bash
 curl -X POST http://127.0.0.1:8000/predict -H "Content-Type: application/json" -d '{"input": 4.0}'
 ```
@@ -218,7 +208,7 @@ litgpt serve microsoft/phi-2
 &nbsp;
 
 # Featured examples    
-Use LitServe to deploy any model or AI service: (Compound AI, Gen AI, classic ML, embeddings, LLMs, vision, audio, etc...)       
+Here are examples of inference pipelines for common model types and use cases.      
   
 <pre>
 <strong>Toy model:</strong>      <a target="_blank" href="#define-a-server">Hello world</a>
@@ -247,13 +237,17 @@ https://github.com/user-attachments/assets/ff83dab9-0c9f-4453-8dcb-fb9526726344
 
 Self-hosting is ideal for hackers, students, and DIY developers while fully managed hosting is ideal for enterprise developers needing easy autoscaling, security, release management, and 99.995% uptime and observability.
 
-To host on [Lightning AI](https://lightning.ai/deploy), simply add the `--cloud` arg, login and choose the cloud of your choice.
+*Note:* Lightning offers a generous free tier for developers.
+
+To host on [Lightning AI](https://lightning.ai/deploy), simply run the command, login and choose the cloud of your choice.
 ```bash
-lightning serve api server.py --cloud
+lightning serve server.py
 ```
 
 &nbsp;
 
+## Features
+
 <div align='center'>
 
 | [Feature](https://lightning.ai/docs/litserve/features)               | Self Managed                      | [Fully Managed on Lightning](https://lightning.ai/deploy)         |

{litserve-0.2.8.dev0 → litserve-0.2.9.dev0}/README.md

@@ -1,18 +1,13 @@
 <div align='center'>
 
-# Easily serve AI models Lightning fast ⚡    
+# Deploy AI models and inference pipelines - ⚡ fast    
 
 <img alt="Lightning" src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/ls_banner2.png" width="800px" style="max-width: 100%;">
 
-&nbsp;
-
-<strong>Lightning-fast serving engine for AI models.</strong>    
-Easy. Flexible. Enterprise-scale.    
+&nbsp; 
 </div>
 
-----
-
-**LitServe** is an easy-to-use, flexible serving engine for AI models built on FastAPI. It augments FastAPI with features like batching, streaming, and GPU autoscaling eliminate the need to rebuild a FastAPI server per model.  
+**LitServe** lets you build high-performance AI inference pipelines on top of FastAPI - no boilerplate. Define one or more models, connect vector DBs, stream responses, batch requests, and autoscale on GPUs out of the box.
 
 LitServe is at least [2x faster](#performance) than plain FastAPI due to AI-specific multi-worker handling.    
 
@@ -22,12 +17,13 @@ LitServe is at least [2x faster](#performance) than plain FastAPI due to AI-spec
 ✅ (2x)+ faster serving  ✅ Easy to use               ✅ LLMs, non LLMs and more
 ✅ Bring your own model  ✅ PyTorch/JAX/TF/...        ✅ Built on FastAPI       
 ✅ GPU autoscaling       ✅ Batching, Streaming       ✅ Self-host or ⚡️ managed
-✅ Compound AI           ✅ Integrate with vLLM, etc  ✅ Serverless             
+✅ Inference pipeline    ✅ Integrate with vLLM, etc  ✅ Serverless             
    
 </pre>
 
 <div align='center'>
 
+[![PyPI Downloads](https://static.pepy.tech/badge/litserve)](https://pepy.tech/projects/litserve)
 [![Discord](https://img.shields.io/discord/1077906959069626439?label=Get%20help%20on%20Discord)](https://discord.gg/WajDThKAur)
 ![cpu-tests](https://github.com/Lightning-AI/litserve/actions/workflows/ci-testing.yml/badge.svg)
 [![codecov](https://codecov.io/gh/Lightning-AI/litserve/graph/badge.svg?token=SmzX8mnKlA)](https://codecov.io/gh/Lightning-AI/litserve)
@@ -65,16 +61,16 @@ pip install litserve
 ```
     
 ### Define a server    
-This toy example with 2 models (AI compound system) shows LitServe's flexibility ([see real examples](#examples)):    
+This toy example with 2 models (inference pipeline) shows LitServe's flexibility ([see real examples](#featured-examples)):    
 
 ```python
 # server.py
 import litserve as ls
 
-# (STEP 1) - DEFINE THE API (compound AI system)
+# (STEP 1) - DEFINE THE API ("inference" pipeline)
 class SimpleLitAPI(ls.LitAPI):
     def setup(self, device):
-        # setup is called once at startup. Build a compound AI system (1+ models), connect DBs, load data, etc...
+        # setup is called once at startup. Defines elements of the pipeline: models, connect DBs, load data, etc...
         self.model1 = lambda x: x**2
         self.model2 = lambda x: x**3
 
@@ -83,11 +79,11 @@ class SimpleLitAPI(ls.LitAPI):
         return request["input"] 
 
     def predict(self, x):
-        # Easily build compound systems. Run inference and return the output.
-        squared = self.model1(x)
-        cubed = self.model2(x)
-        output = squared + cubed
-        return {"output": output}
+        # Run the inference pipeline and return the output
+        a = self.model1(x)
+        b = self.model2(x)
+        c = a + b
+        return {"output": c}
 
     def encode_response(self, output):
         # Convert the model output to a response payload.
@@ -118,12 +114,7 @@ python server.py
 ```
 
 ### Test the server
-Run the auto-generated test client:        
-```bash
-python client.py    
-```
-
-Or use this terminal command:
+Simulate an http request (run this on any terminal):
 ```bash
 curl -X POST http://127.0.0.1:8000/predict -H "Content-Type: application/json" -d '{"input": 4.0}'
 ```
@@ -147,7 +138,7 @@ litgpt serve microsoft/phi-2
 &nbsp;
 
 # Featured examples    
-Use LitServe to deploy any model or AI service: (Compound AI, Gen AI, classic ML, embeddings, LLMs, vision, audio, etc...)       
+Here are examples of inference pipelines for common model types and use cases.      
   
 <pre>
 <strong>Toy model:</strong>      <a target="_blank" href="#define-a-server">Hello world</a>
@@ -176,13 +167,17 @@ https://github.com/user-attachments/assets/ff83dab9-0c9f-4453-8dcb-fb9526726344
 
 Self-hosting is ideal for hackers, students, and DIY developers while fully managed hosting is ideal for enterprise developers needing easy autoscaling, security, release management, and 99.995% uptime and observability.
 
-To host on [Lightning AI](https://lightning.ai/deploy), simply add the `--cloud` arg, login and choose the cloud of your choice.
+*Note:* Lightning offers a generous free tier for developers.
+
+To host on [Lightning AI](https://lightning.ai/deploy), simply run the command, login and choose the cloud of your choice.
 ```bash
-lightning serve api server.py --cloud
+lightning serve server.py
 ```
 
 &nbsp;
 
+## Features
+
 <div align='center'>
 
 | [Feature](https://lightning.ai/docs/litserve/features)               | Self Managed                      | [Fully Managed on Lightning](https://lightning.ai/deploy)         |

{litserve-0.2.8.dev0 → litserve-0.2.9.dev0}/requirements.txt

@@ -1,4 +1,3 @@
 fastapi >=0.100
 uvicorn[standard] >=0.29.0
 pyzmq >=22.0.0
-starlette

{litserve-0.2.8.dev0 → litserve-0.2.9.dev0}/src/litserve/__about__.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-__version__ = "0.2.8.dev0"
+__version__ = "0.2.9.dev0"
 __author__ = "Lightning-AI et al."
 __author_email__ = "community@lightning.ai"
 __license__ = "Apache-2.0"

{litserve-0.2.8.dev0 → litserve-0.2.9.dev0}/src/litserve/api.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import asyncio
 import json
 import warnings
 from abc import ABC, abstractmethod
@@ -30,6 +31,45 @@ class LitAPI(ABC):
     _logger_queue: Optional[Queue] = None
     request_timeout: Optional[float] = None
 
+    def __init__(self, max_batch_size: int = 1, batch_timeout: float = 0.0, enable_async: bool = False):
+        """Initialize a LitAPI instance.
+
+        Args:
+            max_batch_size: Maximum number of requests to process in a batch.
+            batch_timeout: Maximum time to wait for a batch to fill before processing.
+            enable_async: Enable async support.
+
+        """
+
+        if max_batch_size <= 0:
+            raise ValueError("max_batch_size must be greater than 0")
+
+        if batch_timeout < 0:
+            raise ValueError("batch_timeout must be greater than or equal to 0")
+        self.max_batch_size = max_batch_size
+        self.batch_timeout = batch_timeout
+        self.enable_async = enable_async
+        self._validate_async_methods()
+
+    def _validate_async_methods(self):
+        """Validate that async methods are properly implemented when enable_async is True."""
+        if self.enable_async:
+            # check if LitAPI methods are coroutines
+            for method in ["decode_request", "predict", "encode_response"]:
+                if not asyncio.iscoroutinefunction(getattr(self, method)):
+                    raise ValueError("""LitAPI(enable_async=True) requires all methods to be coroutines.
+
+Please either set enable_async=False or implement the following methods as coroutines:
+Example:
+    class MyLitAPI(LitAPI):
+        async def decode_request(self, request, **kwargs):
+            return request
+        async def predict(self, x, **kwargs):
+            return x
+        async def encode_response(self, output, **kwargs):
+            return output
+""")
+
     @abstractmethod
     def setup(self, device):
         """Setup the model so it can be called in `predict`."""
@@ -116,8 +156,10 @@ class LitAPI(ABC):
     def device(self, value):
         self._device = value
 
-    def pre_setup(self, max_batch_size: int, spec: Optional[LitSpec]):
-        self.max_batch_size = max_batch_size
+    def pre_setup(self, spec: Optional[LitSpec]):
+        if self.batch_timeout > self.request_timeout and self.request_timeout not in (False, -1):
+            raise ValueError("batch_timeout must be less than request_timeout")
+
         if self.stream:
             self._default_unbatch = self._unbatch_stream
         else:

{litserve-0.2.8.dev0 → litserve-0.2.9.dev0}/src/litserve/connector.py

@@ -22,20 +22,15 @@ from typing import List, Optional, Union
 class _Connector:
     def __init__(self, accelerator: str = "auto", devices: Union[List[int], int, str] = "auto"):
         accelerator = self._sanitize_accelerator(accelerator)
-        if accelerator == "cpu":
-            self._accelerator = "cpu"
-        elif accelerator == "cuda":
-            self._accelerator = "cuda"
-        elif accelerator == "mps":
-            self._accelerator = "mps"
-
+        if accelerator in ("cpu", "cuda", "mps"):
+            self._accelerator = accelerator
         elif accelerator == "auto":
             self._accelerator = self._choose_auto_accelerator()
         elif accelerator == "gpu":
             self._accelerator = self._choose_gpu_accelerator_backend()
 
         if devices == "auto":
-            self._devices = self._auto_device_count(self._accelerator)
+            self._devices = self._accelerator_device_count()
         else:
             self._devices = devices
 
@@ -43,7 +38,7 @@ class _Connector:
 
     def check_devices_and_accelerators(self):
         """Check if the devices are in a valid fomra and raise an error if they are not."""
-        if self._accelerator in ["cuda", "mps"]:
+        if self._accelerator in ("cuda", "mps"):
             if not isinstance(self._devices, int) and not (
                 isinstance(self._devices, list) and all(isinstance(device, int) for device in self._devices)
             ):
@@ -68,7 +63,7 @@ class _Connector:
             accelerator = accelerator.lower()
 
         if accelerator not in ["auto", "cpu", "mps", "cuda", "gpu", None]:
-            raise ValueError("accelerator must be one of 'auto', 'cpu', 'mps', 'cuda', or 'gpu'")
+            raise ValueError(f"accelerator must be one of 'auto', 'cpu', 'mps', 'cuda', or 'gpu'. Found: {accelerator}")
 
         if accelerator is None:
             return "auto"
@@ -80,8 +75,8 @@ class _Connector:
             return gpu_backend
         return "cpu"
 
-    def _auto_device_count(self, accelerator) -> int:
-        if accelerator == "cuda":
+    def _accelerator_device_count(self) -> int:
+        if self._accelerator == "cuda":
             return check_cuda_with_nvidia_smi()
         return 1
 

{litserve-0.2.8.dev0 → litserve-0.2.9.dev0}/src/litserve/loops/base.py

@@ -14,7 +14,9 @@
 import asyncio
 import inspect
 import logging
+import os
 import pickle
+import signal
 import sys
 import time
 from abc import ABC
@@ -43,17 +45,25 @@ def _inject_context(context: Union[List[dict], dict], func, *args, **kwargs):
     return func(*args, **kwargs)
 
 
+async def _async_inject_context(context: Union[List[dict], dict], func, *args, **kwargs):
+    sig = inspect.signature(func)
+    if "context" in sig.parameters:
+        return await func(*args, **kwargs, context=context)
+    return await func(*args, **kwargs)
+
+
 def collate_requests(
-    lit_api: LitAPI, request_queue: Queue, max_batch_size: int, batch_timeout: float
+    lit_api: LitAPI,
+    request_queue: Queue,
 ) -> Tuple[List, List]:
     payloads = []
     timed_out_uids = []
     entered_at = time.monotonic()
-    end_time = entered_at + batch_timeout
+    end_time = entered_at + lit_api.batch_timeout
     apply_timeout = lit_api.request_timeout not in (-1, False)
 
-    if batch_timeout == 0:
-        while len(payloads) < max_batch_size:
+    if lit_api.batch_timeout == 0:
+        while len(payloads) < lit_api.max_batch_size:
             try:
                 response_queue_id, uid, timestamp, x_enc = request_queue.get_nowait()
                 if apply_timeout and time.monotonic() - timestamp > lit_api.request_timeout:
@@ -64,7 +74,7 @@ def collate_requests(
                 break
         return payloads, timed_out_uids
 
-    while time.monotonic() < end_time and len(payloads) < max_batch_size:
+    while time.monotonic() < end_time and len(payloads) < lit_api.max_batch_size:
         remaining_time = end_time - time.monotonic()
         if remaining_time <= 0:
             break
@@ -111,8 +121,6 @@ class _BaseLoop(ABC):
             worker_id: int,
             request_queue: Queue,
             response_queues: List[Queue],
-            max_batch_size: int,
-            batch_timeout: float,
             stream: bool,
             workers_setup_status: Dict[int, str],
             callback_runner: CallbackRunner,
@@ -140,8 +148,6 @@ class _BaseLoop(ABC):
         lit_api: LitAPI,
         lit_spec: Optional[LitSpec],
         request_queue: Queue,
-        max_batch_size: int,
-        batch_timeout: float,
         response_queues: List[Queue],
     ):
         pass
@@ -154,8 +160,6 @@ class _BaseLoop(ABC):
         worker_id: int,
         request_queue: Queue,
         transport: MessageTransport,
-        max_batch_size: int,
-        batch_timeout: float,
         stream: bool,
         workers_setup_status: Dict[int, str],
         callback_runner: CallbackRunner,
@@ -165,7 +169,7 @@ class _BaseLoop(ABC):
 
             async def _wrapper():
                 logger.info("Running LitLoop in a asyncio event loop")
-                future = self.schedule_task(lit_api, lit_spec, request_queue, max_batch_size, batch_timeout, transport)
+                future = self.schedule_task(lit_api, lit_spec, request_queue, transport)
                 _ = event_loop.create_task(future)
                 while True:
                     try:
@@ -176,8 +180,6 @@ class _BaseLoop(ABC):
                             worker_id,
                             request_queue,
                             transport,
-                            max_batch_size,
-                            batch_timeout,
                             stream,
                             workers_setup_status,
                             callback_runner,
@@ -196,8 +198,6 @@ class _BaseLoop(ABC):
                     worker_id,
                     request_queue,
                     transport,
-                    max_batch_size,
-                    batch_timeout,
                     stream,
                     workers_setup_status,
                     callback_runner,
@@ -211,8 +211,6 @@ class _BaseLoop(ABC):
         worker_id: int,
         request_queue: Queue,
         transport: MessageTransport,
-        max_batch_size: int,
-        batch_timeout: float,
         stream: bool,
         workers_setup_status: Dict[int, str],
         callback_runner: CallbackRunner,
@@ -223,13 +221,24 @@ class _BaseLoop(ABC):
 class LitLoop(_BaseLoop):
     def __init__(self):
         self._context = {}
+        self._server_pid = os.getpid()
 
-    def get_batch_requests(self, lit_api: LitAPI, request_queue: Queue, max_batch_size: int, batch_timeout: float):
+    def kill(self):
+        try:
+            print(f"Stop Server Requested - Kill parent pid [{self._server_pid}] from [{os.getpid()}]")
+            os.kill(self._server_pid, signal.SIGTERM)
+        except PermissionError:
+            # Access Denied because pid already killed...
+            return
+
+    def get_batch_requests(
+        self,
+        lit_api: LitAPI,
+        request_queue: Queue,
+    ):
         batches, timed_out_uids = collate_requests(
             lit_api,
             request_queue,
-            max_batch_size,
-            batch_timeout,
         )
         return batches, timed_out_uids
 

{litserve-0.2.8.dev0 → litserve-0.2.9.dev0}/src/litserve/loops/continuous_batching_loop.py

@@ -120,9 +120,9 @@ requires the lit_api to have a has_finished method. Please implement the has_fin
         lit_api: LitAPI,
         lit_spec: Optional[LitSpec],
         request_queue: Queue,
+        response_queues: List[Queue] = None,
         max_batch_size: Optional[int] = None,
         batch_timeout: Optional[float] = None,
-        response_queues: List[Queue] = None,
     ) -> List[Tuple[str, Any]]:
         """Fill available capacity with pending and new requests."""
         # First process existing pending requests
@@ -151,8 +151,6 @@ requires the lit_api to have a has_finished method. Please implement the has_fin
         lit_api: LitAPI,
         lit_spec: Optional[LitSpec],
         request_queue: Queue,
-        max_batch_size: int,
-        batch_timeout: float,
         response_queues: List[Queue],
     ):
         logger.info("Running prefill in background")
@@ -164,9 +162,9 @@ requires the lit_api to have a has_finished method. Please implement the has_fin
                     lit_api,
                     lit_spec,
                     request_queue,
-                    max_batch_size,
-                    batch_timeout,
                     response_queues,
+                    max_batch_size=lit_api.max_batch_size,
+                    batch_timeout=lit_api.batch_timeout,
                 )
                 await asyncio.sleep(0)
         except Exception as e:
@@ -187,8 +185,6 @@ requires the lit_api to have a has_finished method. Please implement the has_fin
         worker_id: int,
         request_queue: Queue,
         transport: MessageTransport,
-        max_batch_size: int,
-        batch_timeout: float,
         stream: bool,
         workers_setup_status: Dict[int, str],
         callback_runner: CallbackRunner,

{litserve-0.2.8.dev0 → litserve-0.2.9.dev0}/src/litserve/loops/loops.py

@@ -27,16 +27,22 @@ from litserve.utils import WorkerSetupStatus
 logger = logging.getLogger(__name__)
 
 
-def get_default_loop(stream: bool, max_batch_size: int) -> _BaseLoop:
-    return (
-        BatchedStreamingLoop()
-        if stream and max_batch_size > 1
-        else StreamingLoop()
-        if stream
-        else BatchedLoop()
-        if max_batch_size > 1
-        else SingleLoop()
-    )
+def get_default_loop(stream: bool, max_batch_size: int, enable_async: bool = False) -> _BaseLoop:
+    if enable_async:
+        if stream:
+            raise ValueError("Async streaming is not supported. Please use enable_async=False with streaming.")
+        if max_batch_size > 1:
+            raise ValueError("Async batching is not supported. Please use enable_async=False with batching.")
+        return SingleLoop()  # Only SingleLoop supports async currently
+
+    if stream:
+        if max_batch_size > 1:
+            return BatchedStreamingLoop()
+        return StreamingLoop()
+
+    if max_batch_size > 1:
+        return BatchedLoop()
+    return SingleLoop()
 
 
 def inference_worker(
@@ -46,8 +52,6 @@ def inference_worker(
     worker_id: int,
     request_queue: Queue,
     transport: MessageTransport,
-    max_batch_size: int,
-    batch_timeout: float,
     stream: bool,
     workers_setup_status: Dict[int, str],
     callback_runner: CallbackRunner,
@@ -72,7 +76,7 @@ def inference_worker(
         logging.info(f"LitServe will use {lit_spec.__class__.__name__} spec")
 
     if loop == "auto":
-        loop = get_default_loop(stream, max_batch_size)
+        loop = get_default_loop(stream, lit_api.max_batch_size, lit_api.enable_async)
 
     loop(
         lit_api,
@@ -81,8 +85,6 @@ def inference_worker(
         worker_id,
         request_queue,
         transport,
-        max_batch_size,
-        batch_timeout,
         stream,
         workers_setup_status,
         callback_runner,