fix: proper SSE handling for vllm (#877)

fix: proper SSE handling
This commit is contained in:
Zhao Shenyang
2024-02-02 17:25:58 +08:00
committed by GitHub
parent 78660682f5
commit aff5dc8ff2
2 changed files with 5 additions and 1 deletions

View File

@@ -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:

View File

@@ -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')