From aff5dc8ff2a8ecb8df6eb640666d826472e9485e Mon Sep 17 00:00:00 2001 From: Zhao Shenyang Date: Fri, 2 Feb 2024 17:25:58 +0800 Subject: [PATCH] fix: proper SSE handling for vllm (#877) fix: proper SSE handling --- openllm-python/src/openllm/_llm.py | 2 ++ openllm-python/src/openllm/_runners.py | 4 +++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/openllm-python/src/openllm/_llm.py b/openllm-python/src/openllm/_llm.py index afca8ecc..34291695 100644 --- a/openllm-python/src/openllm/_llm.py +++ b/openllm-python/src/openllm/_llm.py @@ -108,11 +108,13 @@ class LLM(t.Generic[M, T]): generator = self.runner.generate_iterator.async_stream( prompt_token_ids, request_id, stop=list(stop), adapter_name=adapter_name, **config.model_dump(flatten=True) ) + generator = bentoml.io.SSE.from_iterator(generator) except Exception as err: raise RuntimeError(f'Failed to start generation task: {err}') from err try: async for out in generator: + out = out.data generated = GenerationOutput.from_runner(out).with_options(prompt=prompt) delta_outputs = [None] * len(generated.outputs) for output in generated.outputs: diff --git a/openllm-python/src/openllm/_runners.py b/openllm-python/src/openllm/_runners.py index b1a120c7..5c43efe5 100644 --- a/openllm-python/src/openllm/_runners.py +++ b/openllm-python/src/openllm/_runners.py @@ -151,7 +151,9 @@ class vLLMRunnable(bentoml.Runnable): async def generate_iterator(self, prompt_token_ids, request_id, stop=None, adapter_name=None, **attrs): _, sampling_params = self.config.model_construct_env(stop=stop, **attrs).inference_options(self.llm) async for request_output in self.model.generate(None, sampling_params, request_id, prompt_token_ids): - yield GenerationOutput.from_vllm(request_output).model_dump_json() + out = GenerationOutput.from_vllm(request_output).model_dump_json() + out = bentoml.io.SSE(out).marshal() + yield out @registry(alias='pt')