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.
- {litserve-0.2.6.dev0/src/litserve.egg-info → litserve-0.2.6.dev1}/PKG-INFO +1 -1
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/__about__.py +1 -1
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/__init__.py +3 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/api.py +2 -63
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/loops.py +326 -18
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/server.py +12 -6
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/specs/openai.py +9 -2
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/specs/openai_embedding.py +11 -5
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/utils.py +44 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1/src/litserve.egg-info}/PKG-INFO +1 -1
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/LICENSE +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/MANIFEST.in +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/README.md +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/requirements.txt +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/setup.cfg +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/setup.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/__main__.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/callbacks/__init__.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/callbacks/base.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/callbacks/defaults/__init__.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/callbacks/defaults/metric_callback.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/connector.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/docker_builder.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/loggers.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/middlewares.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/python_client.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/schema/__init__.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/schema/image.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/specs/__init__.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/specs/base.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/test_examples/__init__.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/test_examples/openai_embedding_spec_example.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/test_examples/openai_spec_example.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/test_examples/simple_example.py +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve.egg-info/SOURCES.txt +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve.egg-info/dependency_links.txt +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve.egg-info/entry_points.txt +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve.egg-info/not-zip-safe +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve.egg-info/requires.txt +0 -0
- {litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve.egg-info/top_level.txt +0 -0
|
@@ -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.
|
|
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
|
|
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
|
-
|
|
41
|
-
|
|
42
|
-
|
|
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
|
|
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(
|
|
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(
|
|
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(
|
|
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
|
|
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,
|
|
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.
|
|
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
|
-
|
|
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)
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/callbacks/defaults/metric_callback.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{litserve-0.2.6.dev0 → litserve-0.2.6.dev1}/src/litserve/test_examples/openai_spec_example.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|