litserve 0.2.6.dev0__tar.gz → 0.2.6.dev1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (40) hide show
  1. {litserve-0.2.6.dev0/src/litserve.egg-info → litserve-0.2.6.dev1}/PKG-INFO +1 -1
  2. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/__about__.py +1 -1
  3. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/__init__.py +3 -0
  4. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/api.py +2 -63
  5. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/loops.py +326 -18
  6. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/server.py +12 -6
  7. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/specs/openai.py +9 -2
  8. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/specs/openai_embedding.py +11 -5
  9. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/utils.py +44 -0
  10. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1/src/litserve.egg-info}/PKG-INFO +1 -1
  11. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/LICENSE +0 -0
  12. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/MANIFEST.in +0 -0
  13. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/README.md +0 -0
  14. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/requirements.txt +0 -0
  15. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/setup.cfg +0 -0
  16. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/setup.py +0 -0
  17. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/__main__.py +0 -0
  18. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/callbacks/__init__.py +0 -0
  19. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/callbacks/base.py +0 -0
  20. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/callbacks/defaults/__init__.py +0 -0
  21. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/callbacks/defaults/metric_callback.py +0 -0
  22. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/connector.py +0 -0
  23. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/docker_builder.py +0 -0
  24. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/loggers.py +0 -0
  25. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/middlewares.py +0 -0
  26. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/python_client.py +0 -0
  27. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/schema/__init__.py +0 -0
  28. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/schema/image.py +0 -0
  29. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/specs/__init__.py +0 -0
  30. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/specs/base.py +0 -0
  31. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/test_examples/__init__.py +0 -0
  32. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/test_examples/openai_embedding_spec_example.py +0 -0
  33. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/test_examples/openai_spec_example.py +0 -0
  34. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/test_examples/simple_example.py +0 -0
  35. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve.egg-info/SOURCES.txt +0 -0
  36. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve.egg-info/dependency_links.txt +0 -0
  37. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve.egg-info/entry_points.txt +0 -0
  38. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve.egg-info/not-zip-safe +0 -0
  39. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve.egg-info/requires.txt +0 -0
  40. {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: litserve
3
- Version: 0.2.6.dev0
3
+ Version: 0.2.6.dev1
4
4
  Summary: Lightweight AI server.
5
5
  Home-page: https://github.com/Lightning-AI/litserve
6
6
  Download-URL: https://github.com/Lightning-AI/litserve
@@ -11,7 +11,7 @@
11
11
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
- __version__ = "0.2.6.dev0"
14
+ __version__ = "0.2.6.dev1"
15
15
  __author__ = "Lightning-AI et al."
16
16
  __author_email__ = "community@lightning.ai"
17
17
  __license__ = "Apache-2.0"
@@ -18,6 +18,9 @@ from litserve.callbacks import Callback
18
18
  from litserve.loggers import Logger
19
19
  from litserve.server import LitServer, Request, Response
20
20
  from litserve.specs import OpenAIEmbeddingSpec, OpenAISpec
21
+ from litserve.utils import configure_logging
22
+
23
+ configure_logging()
21
24
 
22
25
  __all__ = [
23
26
  "LitAPI",
@@ -11,7 +11,6 @@
11
11
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
- import inspect
15
14
  import json
16
15
  import warnings
17
16
  from abc import ABC, abstractmethod
@@ -113,75 +112,15 @@ class LitAPI(ABC):
113
112
  def device(self, value):
114
113
  self._device = value
115
114
 
116
- def _sanitize(self, max_batch_size: int, spec: Optional[LitSpec]):
115
+ def pre_setup(self, max_batch_size: int, spec: Optional[LitSpec]):
116
+ self.max_batch_size = max_batch_size
117
117
  if self.stream:
118
118
  self._default_unbatch = self._unbatch_stream
119
119
  else:
120
120
  self._default_unbatch = self._unbatch_no_stream
121
121
 
122
- # we will sanitize regularly if no spec
123
- # in case, we have spec then:
124
- # case 1: spec implements a streaming API
125
- # Case 2: spec implements a non-streaming API
126
122
  if spec:
127
- # TODO: Implement sanitization
128
123
  self._spec = spec
129
- return
130
-
131
- original = self.unbatch.__code__ is LitAPI.unbatch.__code__
132
- if (
133
- self.stream
134
- and max_batch_size > 1
135
- and not all([
136
- inspect.isgeneratorfunction(self.predict),
137
- inspect.isgeneratorfunction(self.encode_response),
138
- (original or inspect.isgeneratorfunction(self.unbatch)),
139
- ])
140
- ):
141
- raise ValueError(
142
- """When `stream=True` with max_batch_size > 1, `lit_api.predict`, `lit_api.encode_response` and
143
- `lit_api.unbatch` must generate values using `yield`.
144
-
145
- Example:
146
-
147
- def predict(self, inputs):
148
- ...
149
- for i in range(max_token_length):
150
- yield prediction
151
-
152
- def encode_response(self, outputs):
153
- for output in outputs:
154
- encoded_output = ...
155
- yield encoded_output
156
-
157
- def unbatch(self, outputs):
158
- for output in outputs:
159
- unbatched_output = ...
160
- yield unbatched_output
161
- """
162
- )
163
-
164
- if self.stream and not all([
165
- inspect.isgeneratorfunction(self.predict),
166
- inspect.isgeneratorfunction(self.encode_response),
167
- ]):
168
- raise ValueError(
169
- """When `stream=True` both `lit_api.predict` and
170
- `lit_api.encode_response` must generate values using `yield`.
171
-
172
- Example:
173
-
174
- def predict(self, inputs):
175
- ...
176
- for i in range(max_token_length):
177
- yield prediction
178
-
179
- def encode_response(self, outputs):
180
- for output in outputs:
181
- encoded_output = ...
182
- yield encoded_output
183
- """
184
- )
185
124
 
186
125
  def set_logger_queue(self, queue: Queue):
187
126
  """Set the queue for logging events."""
@@ -18,8 +18,9 @@ import multiprocessing as mp
18
18
  import sys
19
19
  import time
20
20
  from abc import ABC
21
+ from dataclasses import dataclass
21
22
  from queue import Empty, Queue
22
- from typing import Dict, List, Optional, Tuple, Union
23
+ from typing import Any, Dict, List, Optional, Tuple, Union
23
24
 
24
25
  from fastapi import HTTPException
25
26
  from starlette.formparsers import MultiPartParser
@@ -37,10 +38,11 @@ try:
37
38
  asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
38
39
 
39
40
  except ImportError:
40
- print(
41
- "uvloop is not installed. Falling back to the default asyncio event loop. "
42
- "Please install uvloop for better performance using `pip install uvloop`."
43
- )
41
+ if sys.platform != "win32":
42
+ print(
43
+ "uvloop is not installed. Falling back to the default asyncio event loop. "
44
+ "Please install uvloop for better performance using `pip install uvloop`."
45
+ )
44
46
 
45
47
  logger = logging.getLogger(__name__)
46
48
 
@@ -439,6 +441,9 @@ class _BaseLoop(ABC):
439
441
 
440
442
  """
441
443
 
444
+ def pre_setup(self, lit_api: LitAPI, spec: Optional[LitSpec]):
445
+ pass
446
+
442
447
  def __call__(
443
448
  self,
444
449
  lit_api: LitAPI,
@@ -485,7 +490,109 @@ class _BaseLoop(ABC):
485
490
  raise NotImplementedError
486
491
 
487
492
 
488
- class SingleLoop(_BaseLoop):
493
+ class LitLoop(_BaseLoop):
494
+ def __init__(self):
495
+ self._context = {}
496
+
497
+ def get_batch_requests(self, lit_api: LitAPI, request_queue: Queue, max_batch_size: int, batch_timeout: float):
498
+ if max_batch_size <= 1:
499
+ raise ValueError("max_batch_size must be greater than 1")
500
+
501
+ batches, timed_out_uids = collate_requests(
502
+ lit_api,
503
+ request_queue,
504
+ max_batch_size,
505
+ batch_timeout,
506
+ )
507
+ return batches, timed_out_uids
508
+
509
+ def get_request(self, request_queue: Queue, timeout: float = 1.0):
510
+ response_queue_id, uid, timestamp, x_enc = request_queue.get(timeout=timeout)
511
+ return response_queue_id, uid, timestamp, x_enc
512
+
513
+ def populate_context(self, lit_spec: LitSpec, request: Any):
514
+ if lit_spec and hasattr(lit_spec, "populate_context"):
515
+ lit_spec.populate_context(self._context, request)
516
+
517
+ def put_response(
518
+ self, response_queues: List[Queue], response_queue_id: int, uid: str, response_data: Any, status: LitAPIStatus
519
+ ) -> None:
520
+ response_queues[response_queue_id].put((uid, (response_data, status)))
521
+
522
+ def put_error_response(
523
+ self, response_queues: List[Queue], response_queue_id: int, uid: str, error: Exception
524
+ ) -> None:
525
+ response_queues[response_queue_id].put((uid, (error, LitAPIStatus.ERROR)))
526
+
527
+
528
+ class DefaultLoop(LitLoop):
529
+ def pre_setup(self, lit_api: LitAPI, spec: Optional[LitSpec]):
530
+ # we will sanitize regularly if no spec
531
+ # in case, we have spec then:
532
+ # case 1: spec implements a streaming API
533
+ # Case 2: spec implements a non-streaming API
534
+ if spec:
535
+ # TODO: Implement sanitization
536
+ lit_api._spec = spec
537
+ return
538
+
539
+ original = lit_api.unbatch.__code__ is LitAPI.unbatch.__code__
540
+ if (
541
+ lit_api.stream
542
+ and lit_api.max_batch_size > 1
543
+ and not all([
544
+ inspect.isgeneratorfunction(lit_api.predict),
545
+ inspect.isgeneratorfunction(lit_api.encode_response),
546
+ (original or inspect.isgeneratorfunction(lit_api.unbatch)),
547
+ ])
548
+ ):
549
+ raise ValueError(
550
+ """When `stream=True` with max_batch_size > 1, `lit_api.predict`, `lit_api.encode_response` and
551
+ `lit_api.unbatch` must generate values using `yield`.
552
+
553
+ Example:
554
+
555
+ def predict(self, inputs):
556
+ ...
557
+ for i in range(max_token_length):
558
+ yield prediction
559
+
560
+ def encode_response(self, outputs):
561
+ for output in outputs:
562
+ encoded_output = ...
563
+ yield encoded_output
564
+
565
+ def unbatch(self, outputs):
566
+ for output in outputs:
567
+ unbatched_output = ...
568
+ yield unbatched_output
569
+ """
570
+ )
571
+
572
+ if lit_api.stream and not all([
573
+ inspect.isgeneratorfunction(lit_api.predict),
574
+ inspect.isgeneratorfunction(lit_api.encode_response),
575
+ ]):
576
+ raise ValueError(
577
+ """When `stream=True` both `lit_api.predict` and
578
+ `lit_api.encode_response` must generate values using `yield`.
579
+
580
+ Example:
581
+
582
+ def predict(self, inputs):
583
+ ...
584
+ for i in range(max_token_length):
585
+ yield prediction
586
+
587
+ def encode_response(self, outputs):
588
+ for output in outputs:
589
+ encoded_output = ...
590
+ yield encoded_output
591
+ """
592
+ )
593
+
594
+
595
+ class SingleLoop(DefaultLoop):
489
596
  def __call__(
490
597
  self,
491
598
  lit_api: LitAPI,
@@ -503,7 +610,7 @@ class SingleLoop(_BaseLoop):
503
610
  run_single_loop(lit_api, lit_spec, request_queue, response_queues, callback_runner)
504
611
 
505
612
 
506
- class BatchedLoop(_BaseLoop):
613
+ class BatchedLoop(DefaultLoop):
507
614
  def __call__(
508
615
  self,
509
616
  lit_api: LitAPI,
@@ -529,7 +636,7 @@ class BatchedLoop(_BaseLoop):
529
636
  )
530
637
 
531
638
 
532
- class StreamingLoop(_BaseLoop):
639
+ class StreamingLoop(DefaultLoop):
533
640
  def __call__(
534
641
  self,
535
642
  lit_api: LitAPI,
@@ -547,7 +654,7 @@ class StreamingLoop(_BaseLoop):
547
654
  run_streaming_loop(lit_api, lit_spec, request_queue, response_queues, callback_runner)
548
655
 
549
656
 
550
- class BatchedStreamingLoop(_BaseLoop):
657
+ class BatchedStreamingLoop(DefaultLoop):
551
658
  def __call__(
552
659
  self,
553
660
  lit_api: LitAPI,
@@ -573,6 +680,203 @@ class BatchedStreamingLoop(_BaseLoop):
573
680
  )
574
681
 
575
682
 
683
+ def notify_timed_out_requests(
684
+ response_queues: List[Queue],
685
+ timed_out_uids: List[Tuple[int, str]],
686
+ ):
687
+ for response_queue_id, uid in timed_out_uids:
688
+ logger.error(f"Request {uid} was waiting in the queue for too long and has been timed out.")
689
+ response_queues[response_queue_id].put((uid, (HTTPException(504, "Request timed out"), LitAPIStatus.ERROR)))
690
+
691
+
692
+ @dataclass
693
+ class Output:
694
+ """Outputs from a single step of the loop."""
695
+
696
+ uid: str
697
+ output: Any
698
+ status: LitAPIStatus
699
+
700
+
701
+ class ContinuousBatchingLoop(LitLoop):
702
+ def __init__(self, max_sequence_length: int = 2048):
703
+ super().__init__()
704
+ self.active_sequences: Dict[str, Dict] = {} # uid -> {input, current_length, generated_tokens}
705
+ self.max_sequence_length = max_sequence_length
706
+ self.response_queue_ids: Dict[str, int] = {} # uid -> response_queue_id
707
+
708
+ def add_request(self, uid: str, request: Any, lit_api: LitAPI, lit_spec: Optional[LitSpec]) -> None:
709
+ """Add a new sequence to active sequences."""
710
+ decoded_request = lit_api.decode_request(request)
711
+ self.active_sequences[uid] = {"input": decoded_request, "current_length": 0, "generated_tokens": []}
712
+
713
+ def mark_completed(self, uid: str) -> None:
714
+ """Mark a sequence as completed."""
715
+ logger.info(f"Marking sequence {uid} as completed")
716
+ del self.active_sequences[uid]
717
+ del self.response_queue_ids[uid]
718
+
719
+ def has_capacity(self, lit_api: LitAPI) -> bool:
720
+ """Check if we can add more sequences based on current batch."""
721
+ capacity = len(self.active_sequences) < lit_api.max_batch_size
722
+ if not capacity:
723
+ logger.info(
724
+ f"No capacity: {len(self.active_sequences)} active sequences, max batch size: {lit_api.max_batch_size}"
725
+ )
726
+ return capacity
727
+
728
+ def step(
729
+ self, prev_outputs: Optional[List[Output]], lit_api: LitAPI, lit_spec: Optional[LitSpec]
730
+ ) -> List[Tuple[str, Tuple[Any, LitAPIStatus]]]:
731
+ """Process one token generation step for all active sequences."""
732
+ if not self.active_sequences:
733
+ return []
734
+
735
+ # Batch forward pass for all active sequences
736
+ inputs = [seq["input"] for seq in self.active_sequences.values()]
737
+ generated = [seq["generated_tokens"] for seq in self.active_sequences.values()]
738
+
739
+ try:
740
+ # Assume lit_api.predict handles batched token generation
741
+ new_tokens = lit_api.predict(inputs, generated)
742
+
743
+ responses = []
744
+
745
+ # Process each sequence's new token
746
+ for uid, token in zip(self.active_sequences.keys(), new_tokens):
747
+ seq = self.active_sequences[uid]
748
+ seq["generated_tokens"].append(token)
749
+ seq["current_length"] += 1
750
+
751
+ # Check completion conditions
752
+ is_finished = lit_api.is_finished(uid, token, self.max_sequence_length)
753
+
754
+ if is_finished:
755
+ # Encode final response for completed sequence
756
+ response = lit_api.encode_response(seq["generated_tokens"])
757
+ step_output = Output(uid, response, LitAPIStatus.FINISH_STREAMING)
758
+ responses.append(step_output)
759
+
760
+ return responses
761
+
762
+ except Exception as e:
763
+ logger.exception("Error during batch token generation")
764
+ # On error, terminate all active sequences
765
+ responses = [(uid, (e, LitAPIStatus.ERROR)) for uid in self.active_sequences]
766
+ self.active_sequences.clear()
767
+ return responses
768
+
769
+ def prefill(
770
+ self,
771
+ pending_requests: List[Tuple[str, Any]],
772
+ lit_api: LitAPI,
773
+ lit_spec: Optional[LitSpec],
774
+ request_queue: Queue,
775
+ max_batch_size: int,
776
+ response_queues: List[Queue],
777
+ ) -> List[Tuple[str, Any]]:
778
+ """Fill available capacity with pending and new requests."""
779
+ # First process existing pending requests
780
+ while pending_requests and self.has_capacity(lit_api):
781
+ response_queue_id, uid, input = pending_requests.pop(0)
782
+ self.add_request(uid, input, lit_api, lit_spec)
783
+ self.response_queue_ids[uid] = response_queue_id
784
+
785
+ # Then check for new requests if we still have capacity
786
+ if self.has_capacity(lit_api):
787
+ new_batches, timed_out_uids = self.get_batch_requests(
788
+ lit_api, request_queue, max_batch_size, batch_timeout=0.0001
789
+ )
790
+ notify_timed_out_requests(response_queues, timed_out_uids)
791
+
792
+ if new_batches:
793
+ # Add new requests to pending_requests and try to process them
794
+ for response_queue_id, uid, input in new_batches:
795
+ logger.info(f"New request: {uid}, {input}")
796
+ if self.has_capacity(lit_api):
797
+ self.add_request(uid, input, lit_api, lit_spec)
798
+ self.response_queue_ids[uid] = response_queue_id
799
+ else:
800
+ pending_requests.append((response_queue_id, uid, input))
801
+
802
+ return pending_requests
803
+
804
+ def run(
805
+ self,
806
+ lit_api: LitAPI,
807
+ lit_spec: Optional[LitSpec],
808
+ device: str,
809
+ worker_id: int,
810
+ request_queue: Queue,
811
+ response_queues: List[Queue],
812
+ max_batch_size: int,
813
+ batch_timeout: float,
814
+ stream: bool,
815
+ workers_setup_status: Dict[int, str],
816
+ callback_runner: CallbackRunner,
817
+ ):
818
+ if not lit_api.stream:
819
+ raise ValueError(
820
+ "Continuous batching loop requires streaming to be enabled. Please set LitServe(..., stream=True)"
821
+ )
822
+
823
+ """Main loop that processes batches of requests."""
824
+ pending_requests = self.prefill(
825
+ [],
826
+ lit_api,
827
+ lit_spec,
828
+ request_queue,
829
+ max_batch_size,
830
+ response_queues,
831
+ )
832
+ try:
833
+ prev_outputs = None
834
+ while pending_requests or self.active_sequences:
835
+ # Process one step for all active sequences
836
+ responses = self.step(prev_outputs, lit_api, lit_spec)
837
+ logger.debug(f"Responses from step(): {responses}")
838
+ if len(responses) == 0:
839
+ raise HTTPException(500, "No responses from step()")
840
+ if responses and not isinstance(responses[0], Output):
841
+ raise HTTPException(500, "Expected StepOutput from step()")
842
+
843
+ prev_outputs = responses
844
+
845
+ # Send responses for all sequences (both streaming and completed)
846
+ for step_output in responses:
847
+ logger.debug(f"Processing response: {step_output}")
848
+ status = step_output.status
849
+ response_data = step_output.output
850
+ uid = step_output.uid
851
+ response_queue_id = self.response_queue_ids[uid]
852
+
853
+ if status == LitAPIStatus.ERROR:
854
+ self.put_error_response(response_queues, response_queue_id, uid, response_data)
855
+ self.mark_completed(uid)
856
+ elif status == LitAPIStatus.FINISH_STREAMING:
857
+ self.put_response(response_queues, response_queue_id, uid, response_data, status)
858
+ self.mark_completed(uid)
859
+ else:
860
+ self.put_response(response_queues, response_queue_id, uid, response_data, status)
861
+
862
+ # Fill available capacity with both pending and new requests
863
+ pending_requests = self.prefill(
864
+ pending_requests,
865
+ lit_api,
866
+ lit_spec,
867
+ request_queue,
868
+ max_batch_size,
869
+ response_queues,
870
+ )
871
+
872
+ except Exception as e:
873
+ logger.exception("Error in continuous batching loop")
874
+ # Handle any errors by sending error responses for all tracked requests
875
+ for uid, response_queue_id in self.response_queue_ids.items():
876
+ self.put_error_response(response_queues, response_queue_id, uid, e)
877
+ self.response_queue_ids.clear()
878
+
879
+
576
880
  def inference_worker(
577
881
  lit_api: LitAPI,
578
882
  lit_spec: Optional[LitSpec],
@@ -606,15 +910,7 @@ def inference_worker(
606
910
  logging.info(f"LitServe will use {lit_spec.__class__.__name__} spec")
607
911
 
608
912
  if loop == "auto":
609
- loop = (
610
- BatchedStreamingLoop()
611
- if stream and max_batch_size > 1
612
- else StreamingLoop()
613
- if stream
614
- else BatchedLoop()
615
- if max_batch_size > 1
616
- else SingleLoop()
617
- )
913
+ loop = get_default_loop(stream, max_batch_size)
618
914
 
619
915
  loop(
620
916
  lit_api,
@@ -629,3 +925,15 @@ def inference_worker(
629
925
  workers_setup_status,
630
926
  callback_runner,
631
927
  )
928
+
929
+
930
+ def get_default_loop(stream: bool, max_batch_size: int) -> _BaseLoop:
931
+ return (
932
+ BatchedStreamingLoop()
933
+ if stream and max_batch_size > 1
934
+ else StreamingLoop()
935
+ if stream
936
+ else BatchedLoop()
937
+ if max_batch_size > 1
938
+ else SingleLoop()
939
+ )
@@ -40,7 +40,7 @@ from litserve import LitAPI
40
40
  from litserve.callbacks.base import Callback, CallbackRunner, EventTypes
41
41
  from litserve.connector import _Connector
42
42
  from litserve.loggers import Logger, _LoggerConnector
43
- from litserve.loops import _BaseLoop, inference_worker
43
+ from litserve.loops import LitLoop, get_default_loop, inference_worker
44
44
  from litserve.middlewares import MaxSizeMiddleware, RequestCountMiddleware
45
45
  from litserve.python_client import client_template
46
46
  from litserve.specs import OpenAISpec
@@ -113,7 +113,7 @@ class LitServer:
113
113
  spec: Optional[LitSpec] = None,
114
114
  max_payload_size=None,
115
115
  track_requests: bool = False,
116
- loop: Optional[Union[str, _BaseLoop]] = "auto",
116
+ loop: Optional[Union[str, LitLoop]] = "auto",
117
117
  callbacks: Optional[Union[List[Callback], Callback]] = None,
118
118
  middlewares: Optional[list[Union[Callable, tuple[Callable, dict]]]] = None,
119
119
  loggers: Optional[Union[Logger, List[Logger]]] = None,
@@ -154,6 +154,8 @@ class LitServer:
154
154
 
155
155
  if isinstance(loop, str) and loop != "auto":
156
156
  raise ValueError("loop must be an instance of _BaseLoop or 'auto'")
157
+ if loop == "auto":
158
+ loop = get_default_loop(stream, max_batch_size)
157
159
 
158
160
  if middlewares is None:
159
161
  middlewares = []
@@ -198,7 +200,7 @@ class LitServer:
198
200
  "but the max_batch_size parameter was not set."
199
201
  )
200
202
 
201
- self._loop = loop
203
+ self._loop: LitLoop = loop
202
204
  self.api_path = api_path
203
205
  self.healthcheck_path = healthcheck_path
204
206
  self.info_path = info_path
@@ -206,7 +208,8 @@ class LitServer:
206
208
  self.timeout = timeout
207
209
  lit_api.stream = stream
208
210
  lit_api.request_timeout = self.timeout
209
- lit_api._sanitize(max_batch_size, spec=spec)
211
+ lit_api.pre_setup(max_batch_size, spec=spec)
212
+ self._loop.pre_setup(lit_api, spec=spec)
210
213
  self.app = FastAPI(lifespan=self.lifespan)
211
214
  self.app.response_queue_id = None
212
215
  self.response_queue_id = None
@@ -427,9 +430,12 @@ class LitServer:
427
430
 
428
431
  await event.wait()
429
432
  response, status = self.response_buffer.pop(uid)
430
-
431
- if status == LitAPIStatus.ERROR:
433
+ if status == LitAPIStatus.ERROR and isinstance(response, HTTPException):
434
+ logger.error("Error in request: %s", response)
432
435
  raise response
436
+ if status == LitAPIStatus.ERROR:
437
+ logger.error("Error in request: %s", response)
438
+ raise HTTPException(status_code=500)
433
439
  self._callback_runner.trigger_event(EventTypes.ON_RESPONSE, litserver=self)
434
440
  return response
435
441
 
@@ -381,8 +381,11 @@ class OpenAISpec(LitSpec):
381
381
  usage_infos = []
382
382
  # iterate over n choices
383
383
  for i, (response, status) in enumerate(streaming_response):
384
- if status == LitAPIStatus.ERROR:
384
+ if status == LitAPIStatus.ERROR and isinstance(response, HTTPException):
385
385
  raise response
386
+ elif status == LitAPIStatus.ERROR:
387
+ logger.error("Error in streaming response: %s", response)
388
+ raise HTTPException(status_code=500)
386
389
  encoded_response = json.loads(response)
387
390
  logger.debug(encoded_response)
388
391
  chat_msg = ChoiceDelta(**encoded_response)
@@ -425,8 +428,12 @@ class OpenAISpec(LitSpec):
425
428
  tool_calls = None
426
429
  usage = None
427
430
  async for response, status in streaming_response:
428
- if status == LitAPIStatus.ERROR:
431
+ if status == LitAPIStatus.ERROR and isinstance(response, HTTPException):
429
432
  raise response
433
+ if status == LitAPIStatus.ERROR:
434
+ logger.error("Error in OpenAI non-streaming response: %s", response)
435
+ raise HTTPException(status_code=500)
436
+
430
437
  # data from LitAPI.encode_response
431
438
  encoded_response = json.loads(response)
432
439
  logger.debug(encoded_response)
@@ -18,7 +18,8 @@ import time
18
18
  import uuid
19
19
  from typing import List, Literal, Optional, Union
20
20
 
21
- from fastapi import Request, Response, status
21
+ from fastapi import HTTPException, Request, Response, status
22
+ from fastapi import status as status_code
22
23
  from pydantic import BaseModel
23
24
 
24
25
  from litserve.specs.base import LitSpec
@@ -28,10 +29,11 @@ logger = logging.getLogger(__name__)
28
29
 
29
30
 
30
31
  class EmbeddingRequest(BaseModel):
31
- input: Union[str, List[str]]
32
+ input: Union[str, List[str], List[int], List[List[int]]]
32
33
  model: str
33
34
  dimensions: Optional[int] = None
34
- encoding_format: Literal["float"] = "float"
35
+ encoding_format: Literal["float", "base64"] = "float"
36
+ user: Optional[str] = None
35
37
 
36
38
  def ensure_list(self):
37
39
  return self.input if isinstance(self.input, list) else [self.input]
@@ -147,7 +149,7 @@ class OpenAIEmbeddingSpec(LitSpec):
147
149
  f"{EMBEDDING_API_EXAMPLE}"
148
150
  )
149
151
 
150
- async def embeddings(self, request: EmbeddingRequest):
152
+ async def embeddings(self, request: EmbeddingRequest) -> EmbeddingResponse:
151
153
  response_queue_id = self.response_queue_id
152
154
  logger.debug("Received embedding request: %s", request)
153
155
  uid = uuid.uuid4()
@@ -159,8 +161,12 @@ class OpenAIEmbeddingSpec(LitSpec):
159
161
 
160
162
  response, status = self._server.response_buffer.pop(uid)
161
163
 
162
- if status == LitAPIStatus.ERROR:
164
+ if status == LitAPIStatus.ERROR and isinstance(response, HTTPException):
165
+ logger.error("Error in embedding request: %s", response)
163
166
  raise response
167
+ if status == LitAPIStatus.ERROR:
168
+ logger.error("Error in embedding request: %s", response)
169
+ raise HTTPException(status_code=status_code.HTTP_500_INTERNAL_SERVER_ERROR)
164
170
 
165
171
  logger.debug(response)
166
172
 
@@ -15,6 +15,7 @@ import asyncio
15
15
  import dataclasses
16
16
  import logging
17
17
  import pickle
18
+ import sys
18
19
  from contextlib import contextmanager
19
20
  from typing import TYPE_CHECKING, AsyncIterator
20
21
 
@@ -87,3 +88,46 @@ class WorkerSetupStatus:
87
88
  READY: str = "ready"
88
89
  ERROR: str = "error"
89
90
  FINISHED: str = "finished"
91
+
92
+
93
+ def configure_logging(
94
+ level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", stream=sys.stdout
95
+ ):
96
+ """Configure logging for the entire library with sensible defaults.
97
+
98
+ Args:
99
+ level (int): Logging level (default: logging.INFO)
100
+ format (str): Log message format string
101
+ stream (file-like): Output stream for logs
102
+
103
+ """
104
+ # Create a library-wide handler
105
+ handler = logging.StreamHandler(stream)
106
+
107
+ # Set formatter with user-configurable format
108
+ formatter = logging.Formatter(format)
109
+ handler.setFormatter(formatter)
110
+
111
+ # Configure root library logger
112
+ library_logger = logging.getLogger("litserve")
113
+ library_logger.setLevel(level)
114
+ library_logger.addHandler(handler)
115
+
116
+ # Prevent propagation to root logger to avoid duplicate logs
117
+ library_logger.propagate = False
118
+
119
+
120
+ def set_log_level(level):
121
+ """Allow users to set the global logging level for the library."""
122
+ logging.getLogger("litserve").setLevel(level)
123
+
124
+
125
+ def add_log_handler(handler):
126
+ """Allow users to add custom log handlers.
127
+
128
+ Example usage:
129
+ file_handler = logging.FileHandler('library_logs.log')
130
+ add_log_handler(file_handler)
131
+
132
+ """
133
+ logging.getLogger("litserve").addHandler(handler)
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: litserve
3
- Version: 0.2.6.dev0
3
+ Version: 0.2.6.dev1
4
4
  Summary: Lightweight AI server.
5
5
  Home-page: https://github.com/Lightning-AI/litserve
6
6
  Download-URL: https://github.com/Lightning-AI/litserve
File without changes
File without changes
File without changes
File without changes
File without changes