Support multiple LitAPIs for inference process and endpoints

src/litserve/api.py

@@ -17,27 +17,43 @@
 import warnings
 from abc import ABC, abstractmethod
 from queue import Queue
-from typing import Callable, Optional
+from typing import TYPE_CHECKING, Callable, Optional, Union
 
 from pydantic import BaseModel
 
 from litserve.specs.base import LitSpec
 
+if TYPE_CHECKING:
+    from litserve.loops.base import LitLoop
+
 
 class LitAPI(ABC):
     _stream: bool = False
     _default_unbatch: Optional[Callable] = None
-    _spec: Optional[LitSpec] = None
+    _spec: Optional["LitSpec"] = None
     _device: Optional[str] = None
     _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):
+    def __init__(
+        self,
+        max_batch_size: int = 1,
+        batch_timeout: float = 0.0,
+        api_path: str = "/predict",
+        stream: bool = False,
+        loop: Optional[Union[str, "LitLoop"]] = "auto",
+        spec: Optional["LitSpec"] = None,
+        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.
+            api_path: URL path for the prediction endpoint.
+            stream: Whether to enable streaming responses.
+            loop: Inference loop to use, or 'auto' to select based on settings.
+            spec: Specification for the API, such as OpenAISpec or custom specs.
             enable_async: Enable async support.
 
         """
@@ -47,6 +63,36 @@ def __init__(self, max_batch_size: int = 1, batch_timeout: float = 0.0, enable_a
 
         if batch_timeout < 0:
             raise ValueError("batch_timeout must be greater than or equal to 0")
+
+        if isinstance(spec, LitSpec):
+            stream = spec.stream
+
+        if loop is None:
+            loop = "auto"
+
+        if isinstance(loop, str) and loop != "auto":
+            raise ValueError("loop must be an instance of _BaseLoop or 'auto'")
+
+        if not api_path.startswith("/"):
+            raise ValueError(
+                "api_path must start with '/'. "
+                "Please provide a valid api path like '/predict', '/classify', or '/v1/predict'"
+            )
+
+        # Check if the batch and unbatch methods are overridden in the lit_api instance
+        batch_overridden = self.batch.__code__ is not LitAPI.batch.__code__
+        unbatch_overridden = self.unbatch.__code__ is not LitAPI.unbatch.__code__
+
+        if batch_overridden and unbatch_overridden and max_batch_size == 1:
+            warnings.warn(
+                "The LitServer has both batch and unbatch methods implemented, "
+                "but the max_batch_size parameter was not set."
+            )
+
+        self.api_path = api_path
+        self.stream = stream
+        self._loop = loop
+        self._spec = spec
         self.max_batch_size = max_batch_size
         self.batch_timeout = batch_timeout
         self.enable_async = enable_async
@@ -82,7 +128,6 @@ async def predict(self, x, **kwargs):
     @abstractmethod
     def setup(self, device):
         """Setup the model so it can be called in `predict`."""
-        pass
 
     def decode_request(self, request, **kwargs):
         """Convert the request payload to your model input."""
@@ -165,10 +210,7 @@ def device(self):
     def device(self, value):
         self._device = value
 
-    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")
-
+    def pre_setup(self, spec: Optional["LitSpec"]):
         if self.stream:
             self._default_unbatch = self._unbatch_stream
         else:
@@ -212,3 +254,15 @@ def health(self) -> bool:
 
         """
         return True
+
+    @property
+    def loop(self):
+        if self._loop == "auto":
+            from litserve.loops.loops import get_default_loop
+
+            self._loop = get_default_loop(self.stream, self.max_batch_size, self.enable_async)
+        return self._loop
+
+    @loop.setter
+    def loop(self, value: "LitLoop"):
+        self._loop = value

src/litserve/python_client.py

@@ -1,5 +1,4 @@
-client_template = """
-# Copyright The Lightning AI team.
+client_template = """# Copyright The Lightning AI team.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.