Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

format encoded response #85

Merged
merged 4 commits into from
May 14, 2024
Merged

format encoded response #85

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
f6730c8
format encoded response
aniketmaurya May 11, 2024
f0a8142
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 11, 2024
4ea8a2f
fix test
aniketmaurya May 11, 2024
814f263
Merge remote-tracking branch 'origin/aniket/stream-dump-dict' into an…
aniketmaurya May 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ Explore various examples that show different models deployed with LitServe:
| Example | description | Run |
|---|---|---|
| [Hello world](#implement-a-server) | Hello world model | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/litserve-hello-world"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/></a> |
| [Llama 3 (8B)](https://lightning.ai/lightning-ai/studios/deploy-a-private-llama-3-8b-api) | **(LLM)** Deploy Llama 3 | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-a-private-llama-3-8b-api"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/></a> |
| [Llama 3 (8B)](https://lightning.ai/lightning-ai/studios/deploy-a-private-llama-3-8b-api) | **(LLM)** Deploy Llama 3 | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-a-private-llama-3-8b-api"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/></a> |
| [ANY Hugging face model](https://lightning.ai/lightning-ai/studios/deploy-any-hugging-face-model-instantly) | **(Text)** Deploy any Hugging Face model | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-any-hugging-face-model-instantly"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/></a> |
| [Hugging face BERT model](https://lightning.ai/lightning-ai/studios/deploy-a-hugging-face-bert-model) | **(Text)** Deploy model for tasks like text generation and more | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-a-hugging-face-bert-model"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/></a>
| [Open AI CLIP](https://lightning.ai/lightning-ai/studios/deploy-open-ai-clip-with-litserve) | **(Multimodal)** Deploy Open AI CLIP for tasks like image understanding | <a target="_blank" href="https://lightning.ai/lightning-ai/studios/deploy-open-ai-clip-with-litserve"><img src="https://pl-bolts-doc-images.s3.us-east-2.amazonaws.com/app-2/studio-badge.svg" alt="Open In Studio"/></a>
Expand Down
10 changes: 10 additions & 0 deletions src/litserve/api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
import json
from abc import ABC, abstractmethod

from pydantic import BaseModel


def no_batch_unbatch_message_no_stream(obj, data):
return f"""
Expand Down Expand Up @@ -113,6 +116,13 @@ def encode_response(self, output):
"""
pass

def format_encoded_response(self, data):
if isinstance(data, dict):
return json.dumps(data)
if isinstance(data, BaseModel):
return data.model_dump_json()
return data

@property
def stream(self):
return self._stream
Expand Down
4 changes: 3 additions & 1 deletion src/litserve/server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,7 @@ def run_single_loop(lit_api, request_queue: Queue, request_buffer):
pipe_s.send((pickle.dumps(e), LitAPIStatus.ERROR))


def run_streaming_loop(lit_api, request_queue: Queue, request_buffer):
def run_streaming_loop(lit_api: LitAPI, request_queue: Queue, request_buffer):
while True:
try:
uid = request_queue.get(timeout=1.0)
Expand All @@ -171,6 +171,7 @@ def run_streaming_loop(lit_api, request_queue: Queue, request_buffer):
y_enc_gen = lit_api.encode_response(y_gen)
for y_enc in y_enc_gen:
with contextlib.suppress(BrokenPipeError):
y_enc = lit_api.format_encoded_response(y_enc)
pipe_s.send((y_enc, LitAPIStatus.OK))
with contextlib.suppress(BrokenPipeError):
pipe_s.send(("", LitAPIStatus.FINISH_STREAMING))
Expand Down Expand Up @@ -205,6 +206,7 @@ def run_batched_streaming_loop(lit_api, request_queue: Queue, request_buffer, ma
for y_batch in y_enc_iter:
for y_enc, pipe_s in zip(y_batch, pipes):
with contextlib.suppress(BrokenPipeError):
y_enc = lit_api.format_encoded_response(y_enc)
pipe_s.send((y_enc, LitAPIStatus.OK))

for pipe_s in pipes:
Expand Down
2 changes: 2 additions & 0 deletions tests/test_lit_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,6 +197,7 @@ def fake_encode(output):
fake_stream_api.decode_request = MagicMock(side_effect=lambda x: x["prompt"])
fake_stream_api.predict = MagicMock(side_effect=fake_predict)
fake_stream_api.encode_response = MagicMock(side_effect=fake_encode)
fake_stream_api.format_encoded_response = MagicMock(side_effect=lambda x: x)

_, requests_queue, request_buffer = loop_args
request_buffer = Manager().dict()
Expand Down Expand Up @@ -249,6 +250,7 @@ def fake_encode(output_iter):
fake_stream_api.predict = MagicMock(side_effect=fake_predict)
fake_stream_api.encode_response = MagicMock(side_effect=fake_encode)
fake_stream_api.unbatch = MagicMock(side_effect=lambda inputs: inputs)
fake_stream_api.format_encoded_response = MagicMock(side_effect=lambda x: x)

_, requests_queue, request_buffer = loop_args
request_buffer = Manager().dict()
Expand Down
Loading