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

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

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: litserve
-Version: 0.2.6.dev2
+Version: 0.2.7.dev0
 Summary: Lightweight AI server.
 Home-page: https://github.com/Lightning-AI/litserve
 Download-URL: https://github.com/Lightning-AI/litserve
@@ -29,6 +29,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: fastapi>=0.100
 Requires-Dist: uvicorn[standard]>=0.29.0
+Requires-Dist: pyzmq>=22.0.0
 Provides-Extra: perf
 Requires-Dist: jsonargparse; extra == "perf"
 Requires-Dist: tenacity; extra == "perf"
@@ -51,6 +52,20 @@ Requires-Dist: python-multipart; extra == "test"
 Requires-Dist: requests; extra == "test"
 Requires-Dist: torch>2.0.0; extra == "test"
 Requires-Dist: transformers; extra == "test"
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: download-url
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: project-url
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 <div align='center'>
 

{litserve-0.2.6.dev2 → litserve-0.2.7.dev0}/requirements.txt

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

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

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

@@ -137,3 +137,21 @@ class LitAPI(ABC):
             )
             return
         self._logger_queue.put((key, value))
+
+    def has_active_requests(self) -> bool:
+        raise NotImplementedError("has_active_requests is not implemented")
+
+    def has_capacity(self) -> bool:
+        raise NotImplementedError("has_capacity is not implemented")
+
+    def health(self) -> bool:
+        """Check the additional health status of the API.
+
+        This method is used in the /health endpoint of the server to determine the health status.
+        Users can extend this method to include additional health checks specific to their application.
+
+        Returns:
+            bool: True if the API is healthy, False otherwise.
+
+        """
+        return True

litserve-0.2.7.dev0/src/litserve/loops/__init__.py

@@ -0,0 +1,38 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 multiprocessing as mp
+
+from litserve.loops.base import LitLoop, _BaseLoop
+from litserve.loops.continuous_batching_loop import ContinuousBatchingLoop, Output
+from litserve.loops.loops import (
+    get_default_loop,
+    inference_worker,
+)
+from litserve.loops.simple_loops import BatchedLoop, SingleLoop
+from litserve.loops.streaming_loops import BatchedStreamingLoop, StreamingLoop
+
+mp.allow_connection_pickling()
+
+__all__ = [
+    "_BaseLoop",
+    "LitLoop",
+    "get_default_loop",
+    "inference_worker",
+    "Output",
+    "SingleLoop",
+    "BatchedLoop",
+    "StreamingLoop",
+    "BatchedStreamingLoop",
+    "ContinuousBatchingLoop",
+]

litserve-0.2.7.dev0/src/litserve/loops/base.py

@@ -0,0 +1,364 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 inspect
+import logging
+import pickle
+import signal
+import sys
+import time
+from abc import ABC
+from queue import Empty, Queue
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from starlette.formparsers import MultiPartParser
+
+from litserve import LitAPI
+from litserve.callbacks import CallbackRunner
+from litserve.specs.base import LitSpec
+from litserve.utils import LitAPIStatus
+from litserve.zmq_queue import Producer
+
+logger = logging.getLogger(__name__)
+# FastAPI writes form files to disk over 1MB by default, which prevents serialization by multiprocessing
+MultiPartParser.max_file_size = sys.maxsize
+
+
+def _inject_context(context: Union[List[dict], dict], func, *args, **kwargs):
+    sig = inspect.signature(func)
+    if "context" in sig.parameters:
+        return func(*args, **kwargs, context=context)
+    return func(*args, **kwargs)
+
+
+def collate_requests(
+    lit_api: LitAPI, request_queue: Queue, max_batch_size: int, batch_timeout: float
+) -> Tuple[List, List]:
+    payloads = []
+    timed_out_uids = []
+    entered_at = time.monotonic()
+    end_time = entered_at + batch_timeout
+    apply_timeout = lit_api.request_timeout not in (-1, False)
+
+    if batch_timeout == 0:
+        while len(payloads) < 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:
+                    timed_out_uids.append((response_queue_id, uid))
+                else:
+                    payloads.append((response_queue_id, uid, x_enc))
+            except Empty:
+                break
+        return payloads, timed_out_uids
+
+    while time.monotonic() < end_time and len(payloads) < max_batch_size:
+        remaining_time = end_time - time.monotonic()
+        if remaining_time <= 0:
+            break
+
+        try:
+            response_queue_id, uid, timestamp, x_enc = request_queue.get(timeout=min(remaining_time, 0.001))
+            if apply_timeout and time.monotonic() - timestamp > lit_api.request_timeout:
+                timed_out_uids.append((response_queue_id, uid))
+            else:
+                payloads.append((response_queue_id, uid, x_enc))
+
+        except Empty:
+            continue
+
+    return payloads, timed_out_uids
+
+
+class _BaseLoop(ABC):
+    """Loop runs an inference engine that executes a specific set of hooks, implemented in the LitAPI, in a predefined
+    order.
+
+    For a default loop, LitAPI must implement the following hooks:
+      - decode_request
+      - batch
+      - predict
+      - unbatch
+      - encode_response
+
+    To implement a custom loop, subclass this class and implement the `run` method. The `run` method should execute the
+    hooks in the desired order.
+
+    `__call__` method is the entry point for the worker process. It calls the `run` method in a loop until the worker is
+    terminated.
+
+    Example:
+
+    ```python
+    class TestLoop(_BaseLoop):
+        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,
+        ):
+            item = request_queue.get()
+            if item is None:
+                return
+
+            response_queue_id, uid, timestamp, x_enc = item
+            # Expects LitAPI to implement the load_cache method
+            lit_api.load_cache(x_enc)
+            x = lit_api.decode_request(x_enc)
+            response = lit_api.predict(x)
+            response_enc = lit_api.encode_response(response)
+            response_queues[response_queue_id].put((uid, (response_enc, LitAPIStatus.OK)))
+    ```
+
+    """
+
+    def pre_setup(self, lit_api: LitAPI, spec: Optional[LitSpec]):
+        pass
+
+    async def schedule_task(
+        self,
+        lit_api: LitAPI,
+        lit_spec: Optional[LitSpec],
+        request_queue: Queue,
+        max_batch_size: int,
+        batch_timeout: float,
+        response_queues: List[Queue],
+    ):
+        pass
+
+    def __call__(
+        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,
+    ):
+        if asyncio.iscoroutinefunction(self.run):
+            event_loop = asyncio.new_event_loop()
+
+            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, response_queues
+                )
+                _ = event_loop.create_task(future)
+                while True:
+                    try:
+                        await self.run(
+                            lit_api,
+                            lit_spec,
+                            device,
+                            worker_id,
+                            request_queue,
+                            response_queues,
+                            max_batch_size,
+                            batch_timeout,
+                            stream,
+                            workers_setup_status,
+                            callback_runner,
+                        )
+                        await asyncio.sleep(0)
+                    except Exception as e:
+                        logger.exception("An error occurred in the loop: %s", e)
+
+            event_loop.run_until_complete(_wrapper())
+        else:
+            while True:
+                self.run(
+                    lit_api,
+                    lit_spec,
+                    device,
+                    worker_id,
+                    request_queue,
+                    response_queues,
+                    max_batch_size,
+                    batch_timeout,
+                    stream,
+                    workers_setup_status,
+                    callback_runner,
+                )
+
+    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,
+    ):
+        raise NotImplementedError
+
+
+class LitLoop(_BaseLoop):
+    def __init__(self):
+        self.producer: Optional[Producer] = None
+        self._context = {}
+        self._setup_signal_handlers()
+
+    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, block: bool = True, timeout: Optional[float] = None):
+        try:
+            return request_queue.get(block=block, 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:
+        if self.producer:
+            self.producer.put((uid, (response_data, status)), consumer_id=response_queue_id)
+        else:
+            response_queues[response_queue_id].put((uid, (response_data, status)), block=False)
+
+    def put_error_response(
+        self, response_queues: List[Queue], response_queue_id: int, uid: str, error: Exception
+    ) -> None:
+        error = pickle.dumps(error)
+        self.put_response(response_queues, response_queue_id, uid, error, LitAPIStatus.ERROR)
+
+    def __del__(self):
+        if self.producer:
+            self.producer.close()
+
+    def _setup_signal_handlers(self):
+        def cleanup_handler(signum=None, frame=None):
+            logging.debug("Worker process received shutdown signal")
+            if self.producer:
+                self.producer.close()
+            sys.exit(0)
+
+        signal.signal(signal.SIGINT, cleanup_handler)
+        signal.signal(signal.SIGTERM, cleanup_handler)
+
+
+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 not lit_api.stream and any([
+            inspect.isgeneratorfunction(lit_api.predict),
+            inspect.isgeneratorfunction(lit_api.encode_response),
+        ]):
+            raise ValueError(
+                """When `stream=False`, `lit_api.predict`, `lit_api.encode_response` must not be
+                generator functions.
+
+                Correct usage:
+
+                    def predict(self, inputs):
+                        ...
+                        return {"output": output}
+
+                Incorrect usage:
+
+                    def predict(self, inputs):
+                        ...
+                        for i in range(max_token_length):
+                            yield prediction
+                """
+            )
+        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
+             """
+            )