fix: proper SSE handling for vllm (#877)

fix: proper SSE handling
bentoml · Feb 2, 2024 · aff5dc8 · aff5dc8
1 parent 7866068
commit aff5dc8
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/openllm-python/src/openllm/_llm.py b/openllm-python/src/openllm/_llm.py
@@ -108,11 +108,13 @@ async def generate_iterator(self, prompt, prompt_token_ids=None, stop=None, stop
       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
@@ -151,7 +151,9 @@ def __init__(self, llm):
   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')