From f6730c8b5a86205fd88236e0a12094a7524dd583 Mon Sep 17 00:00:00 2001 From: Aniket Maurya Date: Sat, 11 May 2024 14:20:46 +0100 Subject: [PATCH 1/3] format encoded response --- src/litserve/api.py | 10 ++++++++++ src/litserve/server.py | 4 +++- tests/test_lit_server.py | 1 + 3 files changed, 14 insertions(+), 1 deletion(-) diff --git a/src/litserve/api.py b/src/litserve/api.py index 18750da4..6c440f46 100644 --- a/src/litserve/api.py +++ b/src/litserve/api.py @@ -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""" @@ -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 diff --git a/src/litserve/server.py b/src/litserve/server.py index 907ffc08..327e4cb2 100644 --- a/src/litserve/server.py +++ b/src/litserve/server.py @@ -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) @@ -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)) @@ -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: diff --git a/tests/test_lit_server.py b/tests/test_lit_server.py index 5ee839cc..19e65d69 100644 --- a/tests/test_lit_server.py +++ b/tests/test_lit_server.py @@ -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() From f0a81424b1a5d9e79f10e83873ec481cf0fb66f7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 11 May 2024 13:21:20 +0000 Subject: [PATCH 2/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 3bda2d4f..8e18a6d8 100644 --- a/README.md +++ b/README.md @@ -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 | Open In Studio | -| [Llama 3 (8B)](https://lightning.ai/lightning-ai/studios/deploy-a-private-llama-3-8b-api) | **(LLM)** Deploy Llama 3 | Open In Studio | +| [Llama 3 (8B)](https://lightning.ai/lightning-ai/studios/deploy-a-private-llama-3-8b-api) | **(LLM)** Deploy Llama 3 | Open In Studio | | [ANY Hugging face model](https://lightning.ai/lightning-ai/studios/deploy-any-hugging-face-model-instantly) | **(Text)** Deploy any Hugging Face model | Open In Studio | | [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 | Open In Studio | [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 | Open In Studio From 4ea8a2f0481d7b0942a9356bdb0bf719f76b70b1 Mon Sep 17 00:00:00 2001 From: Aniket Maurya Date: Sat, 11 May 2024 14:26:25 +0100 Subject: [PATCH 3/3] fix test --- tests/test_lit_server.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_lit_server.py b/tests/test_lit_server.py index 19e65d69..09d681bf 100644 --- a/tests/test_lit_server.py +++ b/tests/test_lit_server.py @@ -250,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()