litserve 0.2.6.dev0__tar.gz → 0.2.6.dev2__tar.gz

{litserve-0.2.6.dev0/src/litserve.egg-info → litserve-0.2.6.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: litserve
-Version: 0.2.6.dev0
+Version: 0.2.6.dev2
 Summary: Lightweight AI server.
 Home-page: https://github.com/Lightning-AI/litserve
 Download-URL: https://github.com/Lightning-AI/litserve
@@ -29,6 +29,10 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: fastapi>=0.100
 Requires-Dist: uvicorn[standard]>=0.29.0
+Provides-Extra: perf
+Requires-Dist: jsonargparse; extra == "perf"
+Requires-Dist: tenacity; extra == "perf"
+Requires-Dist: uvloop; extra == "perf"
 Provides-Extra: test
 Requires-Dist: asgi-lifespan; extra == "test"
 Requires-Dist: coverage[toml]>=7.5.3; extra == "test"
@@ -47,10 +51,6 @@ Requires-Dist: python-multipart; extra == "test"
 Requires-Dist: requests; extra == "test"
 Requires-Dist: torch>2.0.0; extra == "test"
 Requires-Dist: transformers; extra == "test"
-Provides-Extra: perf
-Requires-Dist: jsonargparse; extra == "perf"
-Requires-Dist: tenacity; extra == "perf"
-Requires-Dist: uvloop; extra == "perf"
 
 <div align='center'>
 

{litserve-0.2.6.dev0 → litserve-0.2.6.dev2}/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.6.dev0"
+__version__ = "0.2.6.dev2"
 __author__ = "Lightning-AI et al."
 __author_email__ = "community@lightning.ai"
 __license__ = "Apache-2.0"

{litserve-0.2.6.dev0 → litserve-0.2.6.dev2}/src/litserve/__init__.py

@@ -18,6 +18,9 @@ from litserve.callbacks import Callback
 from litserve.loggers import Logger
 from litserve.server import LitServer, Request, Response
 from litserve.specs import OpenAIEmbeddingSpec, OpenAISpec
+from litserve.utils import configure_logging
+
+configure_logging()
 
 __all__ = [
     "LitAPI",

{litserve-0.2.6.dev0 → litserve-0.2.6.dev2}/src/litserve/api.py

@@ -11,7 +11,6 @@
 # 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 inspect
 import json
 import warnings
 from abc import ABC, abstractmethod
@@ -57,10 +56,9 @@ class LitAPI(ABC):
 
         return inputs
 
-    @abstractmethod
     def predict(self, x, **kwargs):
         """Run the model on the input and return or yield the output."""
-        pass
+        raise NotImplementedError("predict is not implemented")
 
     def _unbatch_no_stream(self, output):
         if isinstance(output, str):
@@ -113,75 +111,16 @@ class LitAPI(ABC):
     def device(self, value):
         self._device = value
 
-    def _sanitize(self, max_batch_size: int, spec: Optional[LitSpec]):
+    def pre_setup(self, max_batch_size: int, spec: Optional[LitSpec]):
+        self.max_batch_size = max_batch_size
         if self.stream:
             self._default_unbatch = self._unbatch_stream
         else:
             self._default_unbatch = self._unbatch_no_stream
 
-        # we will sanitize regularly if no spec
-        # in case, we have spec then:
-        # case 1: spec implements a streaming API
-        # Case 2: spec implements a non-streaming API
         if spec:
-            # TODO: Implement sanitization
             self._spec = spec
-            return
-
-        original = self.unbatch.__code__ is LitAPI.unbatch.__code__
-        if (
-            self.stream
-            and max_batch_size > 1
-            and not all([
-                inspect.isgeneratorfunction(self.predict),
-                inspect.isgeneratorfunction(self.encode_response),
-                (original or inspect.isgeneratorfunction(self.unbatch)),
-            ])
-        ):
-            raise ValueError(
-                """When `stream=True` with max_batch_size > 1, `lit_api.predict`, `lit_api.encode_response` and
-                `lit_api.unbatch` must generate values using `yield`.
-
-             Example:
-
-                def predict(self, inputs):
-                    ...
-                    for i in range(max_token_length):
-                        yield prediction
-
-                def encode_response(self, outputs):
-                    for output in outputs:
-                        encoded_output = ...
-                        yield encoded_output
-
-                def unbatch(self, outputs):
-                    for output in outputs:
-                        unbatched_output = ...
-                        yield unbatched_output
-             """
-            )
-
-        if self.stream and not all([
-            inspect.isgeneratorfunction(self.predict),
-            inspect.isgeneratorfunction(self.encode_response),
-        ]):
-            raise ValueError(
-                """When `stream=True` both `lit_api.predict` and
-             `lit_api.encode_response` must generate values using `yield`.
-
-             Example:
-
-                def predict(self, inputs):
-                    ...
-                    for i in range(max_token_length):
-                        yield prediction
-
-                def encode_response(self, outputs):
-                    for output in outputs:
-                        encoded_output = ...
-                        yield encoded_output
-             """
-            )
+            spec.pre_setup(self)
 
     def set_logger_queue(self, queue: Queue):
         """Set the queue for logging events."""

{litserve-0.2.6.dev0 → litserve-0.2.6.dev2}/src/litserve/loops.py

@@ -18,8 +18,9 @@ import multiprocessing as mp
 import sys
 import time
 from abc import ABC
+from dataclasses import dataclass
 from queue import Empty, Queue
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from fastapi import HTTPException
 from starlette.formparsers import MultiPartParser
@@ -37,10 +38,11 @@ try:
     asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 except ImportError:
-    print(
-        "uvloop is not installed. Falling back to the default asyncio event loop. "
-        "Please install uvloop for better performance using `pip install uvloop`."
-    )
+    if sys.platform != "win32":
+        print(
+            "uvloop is not installed. Falling back to the default asyncio event loop. "
+            "Please install uvloop for better performance using `pip install uvloop`."
+        )
 
 logger = logging.getLogger(__name__)
 
@@ -439,6 +441,9 @@ class _BaseLoop(ABC):
 
     """
 
+    def pre_setup(self, lit_api: LitAPI, spec: Optional[LitSpec]):
+        pass
+
     def __call__(
         self,
         lit_api: LitAPI,
@@ -485,7 +490,108 @@ class _BaseLoop(ABC):
         raise NotImplementedError
 
 
-class SingleLoop(_BaseLoop):
+class LitLoop(_BaseLoop):
+    def __init__(self):
+        self._context = {}
+
+    def get_batch_requests(self, lit_api: LitAPI, request_queue: Queue, max_batch_size: int, batch_timeout: float):
+        batches, timed_out_uids = collate_requests(
+            lit_api,
+            request_queue,
+            max_batch_size,
+            batch_timeout,
+        )
+        return batches, timed_out_uids
+
+    def get_request(self, request_queue: Queue, timeout: float = 1.0):
+        try:
+            return request_queue.get(timeout=timeout)
+        except Empty:
+            return None
+
+    def populate_context(self, lit_spec: LitSpec, request: Any):
+        if lit_spec and hasattr(lit_spec, "populate_context"):
+            lit_spec.populate_context(self._context, request)
+
+    def put_response(
+        self, response_queues: List[Queue], response_queue_id: int, uid: str, response_data: Any, status: LitAPIStatus
+    ) -> None:
+        response_queues[response_queue_id].put((uid, (response_data, status)))
+
+    def put_error_response(
+        self, response_queues: List[Queue], response_queue_id: int, uid: str, error: Exception
+    ) -> None:
+        response_queues[response_queue_id].put((uid, (error, LitAPIStatus.ERROR)))
+
+
+class DefaultLoop(LitLoop):
+    def pre_setup(self, lit_api: LitAPI, spec: Optional[LitSpec]):
+        # we will sanitize regularly if no spec
+        # in case, we have spec then:
+        # case 1: spec implements a streaming API
+        # Case 2: spec implements a non-streaming API
+        if spec:
+            # TODO: Implement sanitization
+            lit_api._spec = spec
+            return
+
+        original = lit_api.unbatch.__code__ is LitAPI.unbatch.__code__
+        if (
+            lit_api.stream
+            and lit_api.max_batch_size > 1
+            and not all([
+                inspect.isgeneratorfunction(lit_api.predict),
+                inspect.isgeneratorfunction(lit_api.encode_response),
+                (original or inspect.isgeneratorfunction(lit_api.unbatch)),
+            ])
+        ):
+            raise ValueError(
+                """When `stream=True` with max_batch_size > 1, `lit_api.predict`, `lit_api.encode_response` and
+                `lit_api.unbatch` must generate values using `yield`.
+
+             Example:
+
+                def predict(self, inputs):
+                    ...
+                    for i in range(max_token_length):
+                        yield prediction
+
+                def encode_response(self, outputs):
+                    for output in outputs:
+                        encoded_output = ...
+                        yield encoded_output
+
+                def unbatch(self, outputs):
+                    for output in outputs:
+                        unbatched_output = ...
+                        yield unbatched_output
+             """
+            )
+
+        if lit_api.stream and not all([
+            inspect.isgeneratorfunction(lit_api.predict),
+            inspect.isgeneratorfunction(lit_api.encode_response),
+        ]):
+            raise ValueError(
+                """When `stream=True` both `lit_api.predict` and
+             `lit_api.encode_response` must generate values using `yield`.
+
+             Example:
+
+                def predict(self, inputs):
+                    ...
+                    for i in range(max_token_length):
+                        yield prediction
+
+                def encode_response(self, outputs):
+                    for output in outputs:
+                        encoded_output = ...
+                        yield encoded_output
+             """
+            )
+
+
+class SingleLoop(DefaultLoop):
     def __call__(
         self,
         lit_api: LitAPI,
@@ -503,7 +609,7 @@ class SingleLoop(_BaseLoop):
         run_single_loop(lit_api, lit_spec, request_queue, response_queues, callback_runner)
 
 
-class BatchedLoop(_BaseLoop):
+class BatchedLoop(DefaultLoop):
     def __call__(
         self,
         lit_api: LitAPI,
@@ -529,7 +635,7 @@ class BatchedLoop(_BaseLoop):
         )
 
 
-class StreamingLoop(_BaseLoop):
+class StreamingLoop(DefaultLoop):
     def __call__(
         self,
         lit_api: LitAPI,
@@ -547,7 +653,7 @@ class StreamingLoop(_BaseLoop):
         run_streaming_loop(lit_api, lit_spec, request_queue, response_queues, callback_runner)
 
 
-class BatchedStreamingLoop(_BaseLoop):
+class BatchedStreamingLoop(DefaultLoop):
     def __call__(
         self,
         lit_api: LitAPI,
@@ -573,6 +679,248 @@ class BatchedStreamingLoop(_BaseLoop):
         )
 
 
+def notify_timed_out_requests(
+    response_queues: List[Queue],
+    timed_out_uids: List[Tuple[int, str]],
+):
+    for response_queue_id, uid in timed_out_uids:
+        logger.error(f"Request {uid} was waiting in the queue for too long and has been timed out.")
+        response_queues[response_queue_id].put((uid, (HTTPException(504, "Request timed out"), LitAPIStatus.ERROR)))
+
+
+@dataclass
+class Output:
+    """Outputs from a single step of the loop."""
+
+    uid: str
+    output: Any
+    status: LitAPIStatus
+
+
+class ContinuousBatchingLoop(LitLoop):
+    def __init__(self, max_sequence_length: int = 2048):
+        """Runs continuous batching loop. This loop handles adding new requests, processing them in batches, and
+        managing the state of active sequences.
+
+        The loop requires the following methods to be implemented in the LitAPI:
+          - setup: sets up the model on the device
+          - decode_request: decodes the client request into a format that can be processed by the model
+          - step: generates a new token for each sequence
+          - encode_response: encodes the response into a format that can be sent to the client
+          - has_finished: checks if the sequence has finished generating
+
+        Args:
+            max_sequence_length (int): The maximum sequence length allowed for any active sequence.
+
+        """
+        super().__init__()
+        self.active_sequences: Dict[str, Dict] = {}  # uid -> {input, current_length, generated_sequence}
+        self.max_sequence_length = max_sequence_length
+        self.response_queue_ids: Dict[str, int] = {}  # uid -> response_queue_id
+
+    def pre_setup(self, lit_api: LitAPI, spec: Optional[LitSpec]):
+        if not lit_api.stream:
+            raise ValueError(
+                "Continuous batching loop requires streaming to be enabled. Please set LitServe(..., stream=True)"
+            )
+
+        if not hasattr(lit_api, "step") and not hasattr(lit_api, "predict"):
+            raise ValueError("""Using the default step method with Continuous batching loop requires the lit_api to
+have a `predict` method which accepts decoded request inputs and a list of generated_sequence.
+Please implement the has_finished method in the lit_api.
+
+    class ExampleAPI(LitAPI):
+        ...
+        def predict(self, inputs, generated_sequence):
+            # implement predict logic
+            # return list of new tokens
+            ...
+        """)
+
+        if not hasattr(lit_api, "step") and not hasattr(lit_api, "has_finished"):
+            raise ValueError("""Using the default step method with Continuous batching loop
+requires the lit_api to have a has_finished method. Please implement the has_finished method in the lit_api.
+
+    class ExampleAPI(LitAPI):
+        ...
+        def has_finished(self, uid: str, token: str, max_sequence_length: int) -> bool:
+            # implement has_finished logic
+            return False
+        """)
+
+    def add_request(self, uid: str, request: Any, lit_api: LitAPI, lit_spec: Optional[LitSpec]) -> None:
+        """Add a new sequence to active sequences and perform any action before prediction such as filling the cache."""
+        if hasattr(lit_api, "add_request"):
+            lit_api.add_request(uid, request)
+        decoded_request = lit_api.decode_request(request)
+        self.active_sequences[uid] = {"input": decoded_request, "current_length": 0, "generated_sequence": []}
+
+    def mark_completed(self, uid: str) -> None:
+        """Mark a request as completed and remove it from the tracked state."""
+        logger.debug(f"Marking sequence {uid} as completed")
+        del self.active_sequences[uid]
+        del self.response_queue_ids[uid]
+
+    def has_capacity(self, lit_api: LitAPI) -> bool:
+        """Check if we can add more sequences based on current batch."""
+        capacity = len(self.active_sequences) < lit_api.max_batch_size
+        if not capacity:
+            logger.info(
+                f"No capacity: {len(self.active_sequences)} active sequences, max batch size: {lit_api.max_batch_size}"
+            )
+        return capacity
+
+    def step(self, prev_outputs: Optional[List[Output]], lit_api: LitAPI, lit_spec: Optional[LitSpec]) -> List[Output]:
+        """Process one token generation step for all active sequences."""
+        if hasattr(lit_api, "step"):
+            return lit_api.step(prev_outputs)
+
+        if not self.active_sequences:
+            return []
+
+        # Batch forward pass for all active sequences
+        inputs = [seq["input"] for seq in self.active_sequences.values()]
+        generated = [seq["generated_sequence"] for seq in self.active_sequences.values()]
+
+        try:
+            # Assume lit_api.predict handles batched token generation
+            new_tokens: List[Any] = lit_api.predict(inputs, generated)
+
+            responses: List[Output] = []
+
+            # Process each sequence's new token
+            for uid, token in zip(self.active_sequences.keys(), new_tokens):
+                seq = self.active_sequences[uid]
+                seq["generated_sequence"].append(token)
+                seq["current_length"] += 1
+
+                step_output = Output(uid, token, LitAPIStatus.OK)
+                responses.append(step_output)
+
+                # Check completion conditions
+                is_finished = lit_api.has_finished(uid, token, self.max_sequence_length)
+
+                if is_finished:
+                    # Encode final response for completed sequence
+                    step_output = Output(uid, "", LitAPIStatus.FINISH_STREAMING)
+                    responses.append(step_output)
+
+            return responses
+
+        except Exception as e:
+            logger.exception("Error during batch token generation")
+            # On error, terminate all active sequences
+            responses = [(uid, (e, LitAPIStatus.ERROR)) for uid in self.active_sequences]
+            self.active_sequences.clear()
+            return responses
+
+    def prefill(
+        self,
+        pending_requests: List[Tuple[str, Any]],
+        lit_api: LitAPI,
+        lit_spec: Optional[LitSpec],
+        request_queue: Queue,
+        max_batch_size: int,
+        response_queues: List[Queue],
+    ) -> List[Tuple[str, Any]]:
+        """Fill available capacity with pending and new requests."""
+        # First process existing pending requests
+        while pending_requests and self.has_capacity(lit_api):
+            response_queue_id, uid, input = pending_requests.pop(0)
+            self.add_request(uid, input, lit_api, lit_spec)
+            self.response_queue_ids[uid] = response_queue_id
+
+        # Then check for new requests if we still have capacity
+        if self.has_capacity(lit_api):
+            new_batches, timed_out_uids = self.get_batch_requests(
+                lit_api, request_queue, max_batch_size, batch_timeout=0.0001
+            )
+            notify_timed_out_requests(response_queues, timed_out_uids)
+
+            if new_batches:
+                # Add new requests to pending_requests and try to process them
+                for response_queue_id, uid, input in new_batches:
+                    logger.debug(f"New request: {uid}, {input}")
+                    if self.has_capacity(lit_api):
+                        self.add_request(uid, input, lit_api, lit_spec)
+                        self.response_queue_ids[uid] = response_queue_id
+                    else:
+                        pending_requests.append((response_queue_id, uid, input))
+
+        return pending_requests
+
+    def run(
+        self,
+        lit_api: LitAPI,
+        lit_spec: Optional[LitSpec],
+        device: str,
+        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,
+    ):
+        """Main loop that processes batches of requests."""
+        pending_requests = self.prefill(
+            [],
+            lit_api,
+            lit_spec,
+            request_queue,
+            max_batch_size,
+            response_queues,
+        )
+        try:
+            prev_outputs = None
+            while pending_requests or self.active_sequences:
+                # Process one step for all active sequences
+                responses = self.step(prev_outputs, lit_api, lit_spec)
+                logger.debug(f"Responses from step(): {responses}")
+                if len(responses) == 0:
+                    raise HTTPException(500, "No responses from step()")
+                if responses and not isinstance(responses[0], Output):
+                    raise HTTPException(500, "Expected StepOutput from step()")
+
+                prev_outputs = responses
+
+                # Send responses for all sequences (both streaming and completed)
+                for step_output in responses:
+                    logger.debug(f"Processing response: {step_output}")
+                    status = step_output.status
+                    response_data = lit_api.encode_response(step_output.output)
+                    uid = step_output.uid
+                    response_queue_id = self.response_queue_ids[uid]
+
+                    response_data = lit_api.format_encoded_response(response_data)
+                    if status == LitAPIStatus.ERROR:
+                        self.put_error_response(response_queues, response_queue_id, uid, response_data)
+                        self.mark_completed(uid)
+                    elif status == LitAPIStatus.FINISH_STREAMING:
+                        self.put_response(response_queues, response_queue_id, uid, response_data, status)
+                        self.mark_completed(uid)
+                    else:
+                        self.put_response(response_queues, response_queue_id, uid, response_data, status)
+
+                # Fill available capacity with both pending and new requests
+                pending_requests = self.prefill(
+                    pending_requests,
+                    lit_api,
+                    lit_spec,
+                    request_queue,
+                    max_batch_size,
+                    response_queues,
+                )
+
+        except Exception as e:
+            logger.exception(f"Error in continuous batching loop: {e}")
+            # Handle any errors by sending error responses for all tracked requests
+            for uid, response_queue_id in self.response_queue_ids.items():
+                self.put_error_response(response_queues, response_queue_id, uid, e)
+            self.response_queue_ids.clear()
+
+
 def inference_worker(
     lit_api: LitAPI,
     lit_spec: Optional[LitSpec],
@@ -606,15 +954,7 @@ def inference_worker(
         logging.info(f"LitServe will use {lit_spec.__class__.__name__} spec")
 
     if loop == "auto":
-        loop = (
-            BatchedStreamingLoop()
-            if stream and max_batch_size > 1
-            else StreamingLoop()
-            if stream
-            else BatchedLoop()
-            if max_batch_size > 1
-            else SingleLoop()
-        )
+        loop = get_default_loop(stream, max_batch_size)
 
     loop(
         lit_api,
@@ -629,3 +969,15 @@ def inference_worker(
         workers_setup_status,
         callback_runner,
     )
+
+
+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()
+    )

{litserve-0.2.6.dev0 → litserve-0.2.6.dev2}/src/litserve/server.py

@@ -40,7 +40,7 @@ from litserve import LitAPI
 from litserve.callbacks.base import Callback, CallbackRunner, EventTypes
 from litserve.connector import _Connector
 from litserve.loggers import Logger, _LoggerConnector
-from litserve.loops import _BaseLoop, inference_worker
+from litserve.loops import LitLoop, get_default_loop, inference_worker
 from litserve.middlewares import MaxSizeMiddleware, RequestCountMiddleware
 from litserve.python_client import client_template
 from litserve.specs import OpenAISpec
@@ -113,7 +113,7 @@ class LitServer:
         spec: Optional[LitSpec] = None,
         max_payload_size=None,
         track_requests: bool = False,
-        loop: Optional[Union[str, _BaseLoop]] = "auto",
+        loop: Optional[Union[str, LitLoop]] = "auto",
         callbacks: Optional[Union[List[Callback], Callback]] = None,
         middlewares: Optional[list[Union[Callable, tuple[Callable, dict]]]] = None,
         loggers: Optional[Union[Logger, List[Logger]]] = None,
@@ -154,6 +154,8 @@ class LitServer:
 
         if isinstance(loop, str) and loop != "auto":
             raise ValueError("loop must be an instance of _BaseLoop or 'auto'")
+        if loop == "auto":
+            loop = get_default_loop(stream, max_batch_size)
 
         if middlewares is None:
             middlewares = []
@@ -198,7 +200,7 @@ class LitServer:
                 "but the max_batch_size parameter was not set."
             )
 
-        self._loop = loop
+        self._loop: LitLoop = loop
         self.api_path = api_path
         self.healthcheck_path = healthcheck_path
         self.info_path = info_path
@@ -206,7 +208,8 @@ class LitServer:
         self.timeout = timeout
         lit_api.stream = stream
         lit_api.request_timeout = self.timeout
-        lit_api._sanitize(max_batch_size, spec=spec)
+        lit_api.pre_setup(max_batch_size, spec=spec)
+        self._loop.pre_setup(lit_api, spec=spec)
         self.app = FastAPI(lifespan=self.lifespan)
         self.app.response_queue_id = None
         self.response_queue_id = None
@@ -427,9 +430,12 @@ class LitServer:
 
             await event.wait()
             response, status = self.response_buffer.pop(uid)
-
-            if status == LitAPIStatus.ERROR:
+            if status == LitAPIStatus.ERROR and isinstance(response, HTTPException):
+                logger.error("Error in request: %s", response)
                 raise response
+            if status == LitAPIStatus.ERROR:
+                logger.error("Error in request: %s", response)
+                raise HTTPException(status_code=500)
             self._callback_runner.trigger_event(EventTypes.ON_RESPONSE, litserver=self)
             return response
 

{litserve-0.2.6.dev0 → litserve-0.2.6.dev2}/src/litserve/specs/base.py

@@ -15,7 +15,7 @@ from abc import abstractmethod
 from typing import TYPE_CHECKING, Callable, List
 
 if TYPE_CHECKING:
-    from litserve import LitServer
+    from litserve import LitAPI, LitServer
 
 
 class LitSpec:
@@ -26,6 +26,9 @@ class LitSpec:
 
         self._server: LitServer = None
 
+    def pre_setup(self, lit_api: "LitAPI"):
+        pass
+
     def setup(self, server: "LitServer"):
         self._server = server
 

{litserve-0.2.6.dev0 → litserve-0.2.6.dev2}/src/litserve/specs/openai.py

@@ -30,7 +30,7 @@ from litserve.specs.base import LitSpec
 from litserve.utils import LitAPIStatus, azip
 
 if typing.TYPE_CHECKING:
-    from litserve import LitServer
+    from litserve import LitAPI, LitServer
 
 logger = logging.getLogger(__name__)
 
@@ -262,18 +262,18 @@ class OpenAISpec(LitSpec):
         self.add_endpoint("/v1/chat/completions", self.chat_completion, ["POST"])
         self.add_endpoint("/v1/chat/completions", self.options_chat_completions, ["OPTIONS"])
 
-    def setup(self, server: "LitServer"):
+    def pre_setup(self, lit_api: "LitAPI"):
         from litserve import LitAPI
 
-        super().setup(server)
-
-        lit_api = self._server.lit_api
         if not inspect.isgeneratorfunction(lit_api.predict):
             raise ValueError(LITAPI_VALIDATION_MSG.format("predict is not a generator"))
 
         is_encode_response_original = lit_api.encode_response.__code__ is LitAPI.encode_response.__code__
         if not is_encode_response_original and not inspect.isgeneratorfunction(lit_api.encode_response):
             raise ValueError(LITAPI_VALIDATION_MSG.format("encode_response is not a generator"))
+
+    def setup(self, server: "LitServer"):
+        super().setup(server)
         print("OpenAI spec setup complete")
 
     def populate_context(self, context, request):
@@ -381,8 +381,11 @@ class OpenAISpec(LitSpec):
             usage_infos = []
             # iterate over n choices
             for i, (response, status) in enumerate(streaming_response):
-                if status == LitAPIStatus.ERROR:
+                if status == LitAPIStatus.ERROR and isinstance(response, HTTPException):
                     raise response
+                elif status == LitAPIStatus.ERROR:
+                    logger.error("Error in streaming response: %s", response)
+                    raise HTTPException(status_code=500)
                 encoded_response = json.loads(response)
                 logger.debug(encoded_response)
                 chat_msg = ChoiceDelta(**encoded_response)
@@ -425,8 +428,12 @@ class OpenAISpec(LitSpec):
             tool_calls = None
             usage = None
             async for response, status in streaming_response:
-                if status == LitAPIStatus.ERROR:
+                if status == LitAPIStatus.ERROR and isinstance(response, HTTPException):
                     raise response
+                if status == LitAPIStatus.ERROR:
+                    logger.error("Error in OpenAI non-streaming response: %s", response)
+                    raise HTTPException(status_code=500)
+
                 # data from LitAPI.encode_response
                 encoded_response = json.loads(response)
                 logger.debug(encoded_response)

{litserve-0.2.6.dev0 → litserve-0.2.6.dev2}/src/litserve/specs/openai_embedding.py

@@ -18,7 +18,8 @@ import time
 import uuid
 from typing import List, Literal, Optional, Union
 
-from fastapi import Request, Response, status
+from fastapi import HTTPException, Request, Response, status
+from fastapi import status as status_code
 from pydantic import BaseModel
 
 from litserve.specs.base import LitSpec
@@ -28,10 +29,11 @@ logger = logging.getLogger(__name__)
 
 
 class EmbeddingRequest(BaseModel):
-    input: Union[str, List[str]]
+    input: Union[str, List[str], List[int], List[List[int]]]
     model: str
     dimensions: Optional[int] = None
-    encoding_format: Literal["float"] = "float"
+    encoding_format: Literal["float", "base64"] = "float"
+    user: Optional[str] = None
 
     def ensure_list(self):
         return self.input if isinstance(self.input, list) else [self.input]
@@ -147,7 +149,7 @@ class OpenAIEmbeddingSpec(LitSpec):
                 f"{EMBEDDING_API_EXAMPLE}"
             )
 
-    async def embeddings(self, request: EmbeddingRequest):
+    async def embeddings(self, request: EmbeddingRequest) -> EmbeddingResponse:
         response_queue_id = self.response_queue_id
         logger.debug("Received embedding request: %s", request)
         uid = uuid.uuid4()
@@ -159,8 +161,12 @@ class OpenAIEmbeddingSpec(LitSpec):
 
         response, status = self._server.response_buffer.pop(uid)
 
-        if status == LitAPIStatus.ERROR:
+        if status == LitAPIStatus.ERROR and isinstance(response, HTTPException):
+            logger.error("Error in embedding request: %s", response)
             raise response
+        if status == LitAPIStatus.ERROR:
+            logger.error("Error in embedding request: %s", response)
+            raise HTTPException(status_code=status_code.HTTP_500_INTERNAL_SERVER_ERROR)
 
         logger.debug(response)
 

{litserve-0.2.6.dev0 → litserve-0.2.6.dev2}/src/litserve/utils.py

@@ -15,6 +15,7 @@ import asyncio
 import dataclasses
 import logging
 import pickle
+import sys
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, AsyncIterator
 
@@ -87,3 +88,46 @@ class WorkerSetupStatus:
     READY: str = "ready"
     ERROR: str = "error"
     FINISHED: str = "finished"
+
+
+def configure_logging(
+    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", stream=sys.stdout
+):
+    """Configure logging for the entire library with sensible defaults.
+
+    Args:
+        level (int): Logging level (default: logging.INFO)
+        format (str): Log message format string
+        stream (file-like): Output stream for logs
+
+    """
+    # Create a library-wide handler
+    handler = logging.StreamHandler(stream)
+
+    # Set formatter with user-configurable format
+    formatter = logging.Formatter(format)
+    handler.setFormatter(formatter)
+
+    # Configure root library logger
+    library_logger = logging.getLogger("litserve")
+    library_logger.setLevel(level)
+    library_logger.addHandler(handler)
+
+    # Prevent propagation to root logger to avoid duplicate logs
+    library_logger.propagate = False
+
+
+def set_log_level(level):
+    """Allow users to set the global logging level for the library."""
+    logging.getLogger("litserve").setLevel(level)
+
+
+def add_log_handler(handler):
+    """Allow users to add custom log handlers.
+
+    Example usage:
+    file_handler = logging.FileHandler('library_logs.log')
+    add_log_handler(file_handler)
+
+    """
+    logging.getLogger("litserve").addHandler(handler)

{litserve-0.2.6.dev0 → litserve-0.2.6.dev2/src/litserve.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: litserve
-Version: 0.2.6.dev0
+Version: 0.2.6.dev2
 Summary: Lightweight AI server.
 Home-page: https://github.com/Lightning-AI/litserve
 Download-URL: https://github.com/Lightning-AI/litserve
@@ -29,6 +29,10 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: fastapi>=0.100
 Requires-Dist: uvicorn[standard]>=0.29.0
+Provides-Extra: perf
+Requires-Dist: jsonargparse; extra == "perf"
+Requires-Dist: tenacity; extra == "perf"
+Requires-Dist: uvloop; extra == "perf"
 Provides-Extra: test
 Requires-Dist: asgi-lifespan; extra == "test"
 Requires-Dist: coverage[toml]>=7.5.3; extra == "test"
@@ -47,10 +51,6 @@ Requires-Dist: python-multipart; extra == "test"
 Requires-Dist: requests; extra == "test"
 Requires-Dist: torch>2.0.0; extra == "test"
 Requires-Dist: transformers; extra == "test"
-Provides-Extra: perf
-Requires-Dist: jsonargparse; extra == "perf"
-Requires-Dist: tenacity; extra == "perf"
-Requires-Dist: uvloop; extra == "perf"
 
 <div align='center'>