litserve 0.2.6.dev1__tar.gz → 0.2.6.dev2__tar.gz

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

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

@@ -56,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):
@@ -121,6 +120,7 @@ class LitAPI(ABC):
 
         if spec:
             self._spec = spec
+            spec.pre_setup(self)
 
     def set_logger_queue(self, queue: Queue):
         """Set the queue for logging events."""

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

@@ -495,9 +495,6 @@ class LitLoop(_BaseLoop):
         self._context = {}
 
     def get_batch_requests(self, lit_api: LitAPI, request_queue: Queue, max_batch_size: int, batch_timeout: float):
-        if max_batch_size <= 1:
-            raise ValueError("max_batch_size must be greater than 1")
-
         batches, timed_out_uids = collate_requests(
             lit_api,
             request_queue,
@@ -507,8 +504,10 @@ class LitLoop(_BaseLoop):
         return batches, timed_out_uids
 
     def get_request(self, request_queue: Queue, timeout: float = 1.0):
-        response_queue_id, uid, timestamp, x_enc = request_queue.get(timeout=timeout)
-        return response_queue_id, uid, timestamp, x_enc
+        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"):
@@ -700,19 +699,65 @@ class Output:
 
 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_tokens}
+        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."""
+        """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_tokens": []}
+        self.active_sequences[uid] = {"input": decoded_request, "current_length": 0, "generated_sequence": []}
 
     def mark_completed(self, uid: str) -> None:
-        """Mark a sequence as completed."""
-        logger.info(f"Marking sequence {uid} as completed")
+        """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]
 
@@ -725,36 +770,39 @@ class ContinuousBatchingLoop(LitLoop):
             )
         return capacity
 
-    def step(
-        self, prev_outputs: Optional[List[Output]], lit_api: LitAPI, lit_spec: Optional[LitSpec]
-    ) -> List[Tuple[str, Tuple[Any, LitAPIStatus]]]:
+    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_tokens"] 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 = lit_api.predict(inputs, generated)
+            new_tokens: List[Any] = lit_api.predict(inputs, generated)
 
-            responses = []
+            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_tokens"].append(token)
+                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.is_finished(uid, token, self.max_sequence_length)
+                is_finished = lit_api.has_finished(uid, token, self.max_sequence_length)
 
                 if is_finished:
                     # Encode final response for completed sequence
-                    response = lit_api.encode_response(seq["generated_tokens"])
-                    step_output = Output(uid, response, LitAPIStatus.FINISH_STREAMING)
+                    step_output = Output(uid, "", LitAPIStatus.FINISH_STREAMING)
                     responses.append(step_output)
 
             return responses
@@ -792,7 +840,7 @@ class ContinuousBatchingLoop(LitLoop):
             if new_batches:
                 # Add new requests to pending_requests and try to process them
                 for response_queue_id, uid, input in new_batches:
-                    logger.info(f"New request: {uid}, {input}")
+                    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
@@ -815,11 +863,6 @@ class ContinuousBatchingLoop(LitLoop):
         workers_setup_status: Dict[int, str],
         callback_runner: CallbackRunner,
     ):
-        if not lit_api.stream:
-            raise ValueError(
-                "Continuous batching loop requires streaming to be enabled. Please set LitServe(..., stream=True)"
-            )
-
         """Main loop that processes batches of requests."""
         pending_requests = self.prefill(
             [],
@@ -846,10 +889,11 @@ class ContinuousBatchingLoop(LitLoop):
                 for step_output in responses:
                     logger.debug(f"Processing response: {step_output}")
                     status = step_output.status
-                    response_data = step_output.output
+                    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)
@@ -870,7 +914,7 @@ class ContinuousBatchingLoop(LitLoop):
                 )
 
         except Exception as e:
-            logger.exception("Error in continuous batching loop")
+            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)

{litserve-0.2.6.dev1 → 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.dev1 → 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):

{litserve-0.2.6.dev1 → 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.dev1
+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'>