Repository: vllm-project/vllm

Issue:
[BUG]: Streaming `logprob` & `echo` combo.
I'm trying to start writing a `logprob` & `echo` support for chat request. 

Unfortunately, running test like #1992 when `echo` is setted as `true` server doesn't respond.

Seeing furtherer I checked that the **bug** begging in #2449 (sha: dd7e8f5f643167e3f13045cf75cbead54cb2ccfe).
Previous commit #2463 (sha: d2a68364c473a3167a1c2b90f947bb611322a867) worked ok.

## LOG:
```
vllm-openai-main  | INFO 02-01 04:31:38 async_llm_engine.py:385] Received request cmpl-dc7fb40d1b534a879768966f3dc50d39: prompt: None, prefix_pos: None,sampling params: SamplingParams(n=1, best_of=1, presence_penalty=0.0, frequency_penalty=0.0, repetition_penalty=1.0, temperature=1.0, top_p=1.0, top_k=-1, min_p=0.0, use_beam_search=False, length_penalty=1.0, early_stopping=False, stop=[], stop_token_ids=[], include_stop_str_in_output=False, ignore_eos=True, max_tokens=20, logprobs=0, prompt_logprobs=0, skip_special_tokens=True, spaces_between_special_tokens=True), prompt token ids: [2, 12375, 351, 5, 232, 651, 11, 2760, 116, 50118, 6557, 45117, 35, 50118].
vllm-openai-main  | ERROR:    Exception in ASGI application
vllm-openai-main  | Traceback (most recent call last):
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/starlette/responses.py", line 259, in __call__
vllm-openai-main  |     await wrap(partial(self.listen_for_disconnect, receive))
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/starlette/responses.py", line 255, in wrap
vllm-openai-main  |     await func()
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/starlette/responses.py", line 232, in listen_for_disconnect
vllm-openai-main  |     message = await receive()
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/uvicorn/protocols/http/httptools_impl.py", line 587, in receive
vllm-openai-main  |     await self.message_event.wait()
vllm-openai-main  |   File "/usr/lib/python3.10/asyncio/locks.py", line 214, in wait
vllm-openai-main  |     await fut
vllm-openai-main  | asyncio.exceptions.CancelledError: Cancelled by cancel scope 7f20005d5150
vllm-openai-main  | 
vllm-openai-main  | During handling of the above exception, another exception occurred:
vllm-openai-main  | 
vllm-openai-main  | Traceback (most recent call last):
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/uvicorn/protocols/http/httptools_impl.py", line 426, in run_asgi
vllm-openai-main  |     result = await app(  # type: ignore[func-returns-value]
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/uvicorn/middleware/proxy_headers.py", line 84, in __call__
vllm-openai-main  |     return await self.app(scope, receive, send)
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/fastapi/applications.py", line 1054, in __call__
vllm-openai-main  |     await super().__call__(scope, receive, send)
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/starlette/applications.py", line 116, in __call__
vllm-openai-main  |     await self.middleware_stack(scope, receive, send)
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/starlette/middleware/errors.py", line 186, in __call__
vllm-openai-main  |     raise exc
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/starlette/middleware/errors.py", line 164, in __call__
vllm-openai-main  |     await self.app(scope, receive, _send)
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/starlette/middleware/cors.py", line 83, in __call__
vllm-openai-main  |     await self.app(scope, receive, send)
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/aioprometheus/asgi/middleware.py", line 184, in __call__
vllm-openai-main  |     await self.asgi_callable(scope, receive, wrapped_send)
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/starlette/middleware/exceptions.py", line 62, in __call__
vllm-openai-main  |     await wrap_app_handling_exceptions(self.app, conn)(scope, receive, send)
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/starlette/_exception_handler.py", line 55, in wrapped_app
vllm-openai-main  |     raise exc
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/starlette/_exception_handler.py", line 44, in wrapped_app
vllm-openai-main  |     await app(scope, receive, sender)
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/starlette/routing.py", line 746, in __call__
vllm-openai-main  |     await route.handle(scope, receive, send)
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/starlette/routing.py", line 288, in handle
vllm-openai-main  |     await self.app(scope, receive, send)
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/starlette/routing.py", line 75, in app
vllm-openai-main  |     await wrap_app_handling_exceptions(app, request)(scope, receive, send)
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/starlette/_exception_handler.py", line 55, in wrapped_app
vllm-openai-main  |     raise exc
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/starlette/_exception_handler.py", line 44, in wrapped_app
vllm-openai-main  |     await app(scope, receive, sender)
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/starlette/routing.py", line 73, in app
vllm-openai-main  |     await response(scope, receive, send)
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/starlette/responses.py", line 252, in __call__
vllm-openai-main  |     async with anyio.create_task_group() as task_group:
vllm-openai-main  |   File "/usr/local/lib/python3.10/dist-packages/anyio/_backends/_asyncio.py", line 678, in __aexit__
vllm-openai-main  |     raise BaseExceptionGroup(
vllm-openai-main  | exceptiongroup.ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)
vllm-openai-main  | INFO 02-01 04:31:38 async_llm_engine.py:111] Finished request cmpl-dc7fb40d1b534a879768966f3dc50d39.
```


Discussion:
same question


Resolving patch:
diff --git a/vllm/entrypoints/openai/serving_chat.py b/vllm/entrypoints/openai/serving_chat.py
--- a/vllm/entrypoints/openai/serving_chat.py
+++ b/vllm/entrypoints/openai/serving_chat.py
@@ -63,8 +63,9 @@ async def create_chat_completion(
 
         request_id = f"cmpl-{random_uuid()}"
         try:
-            token_ids = self._validate_prompt_and_tokenize(request,
-                                                           prompt=prompt)
+            # Tokenize/detokenize depending on prompt format (string/token list)
+            prompt_ids, prompt_text = self._validate_prompt_and_tokenize(
+                request, prompt=prompt)
             sampling_params = request.to_sampling_params()
             lora_request = self._maybe_get_lora(request)
             guided_decode_logits_processor = (
@@ -78,8 +79,8 @@ async def create_chat_completion(
         except ValueError as e:
             return self.create_error_response(str(e))
 
-        result_generator = self.engine.generate(prompt, sampling_params,
-                                                request_id, token_ids,
+        result_generator = self.engine.generate(prompt_text, sampling_params,
+                                                request_id, prompt_ids,
                                                 lora_request)
         # Streaming response
         if request.stream:
diff --git a/vllm/entrypoints/openai/serving_completion.py b/vllm/entrypoints/openai/serving_completion.py
--- a/vllm/entrypoints/openai/serving_completion.py
+++ b/vllm/entrypoints/openai/serving_completion.py
@@ -136,23 +136,24 @@ async def create_completion(self, request: CompletionRequest,
 
             for i, prompt in enumerate(prompts):
                 if prompt_is_tokens:
-                    input_ids = self._validate_prompt_and_tokenize(
+                    prompt_formats = self._validate_prompt_and_tokenize(
                         request,
                         prompt_ids=prompt,
                         truncate_prompt_tokens=sampling_params.
                         truncate_prompt_tokens)
                 else:
-                    input_ids = self._validate_prompt_and_tokenize(
+                    prompt_formats = self._validate_prompt_and_tokenize(
                         request,
                         prompt=prompt,
                         truncate_prompt_tokens=sampling_params.
                         truncate_prompt_tokens)
+                prompt_ids, prompt_text = prompt_formats
 
                 generators.append(
-                    self.engine.generate(prompt,
+                    self.engine.generate(prompt_text,
                                          sampling_params,
                                          f"{request_id}-{i}",
-                                         prompt_token_ids=input_ids,
+                                         prompt_token_ids=prompt_ids,
                                          lora_request=lora_request))
         except ValueError as e:
             # TODO: Use a vllm-specific Validation Error
@@ -326,7 +327,8 @@ def request_output_to_completion_response(
                     output_text = prompt_text
                 elif request.echo and request.max_tokens > 0:
                     token_ids = prompt_token_ids + output.token_ids
-                    top_logprobs = prompt_logprobs + output.logprobs
+                    top_logprobs = (prompt_logprobs + output.logprobs
+                                    if request.logprobs else None)
                     output_text = prompt_text + output.text
                 else:
                     token_ids = output.token_ids
@@ -334,6 +336,9 @@ def request_output_to_completion_response(
                     output_text = output.text
 
                 if request.logprobs is not None:
+                    assert top_logprobs is not None, (
+                        "top_logprobs must be provided when logprobs "
+                        "is requested")
                     logprobs = self._create_logprobs(
                         token_ids=token_ids,
                         top_logprobs=top_logprobs,
diff --git a/vllm/entrypoints/openai/serving_engine.py b/vllm/entrypoints/openai/serving_engine.py
--- a/vllm/entrypoints/openai/serving_engine.py
+++ b/vllm/entrypoints/openai/serving_engine.py
@@ -2,7 +2,7 @@
 import json
 from dataclasses import dataclass
 from http import HTTPStatus
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 from pydantic import conint
 
@@ -99,27 +99,32 @@ def _create_logprobs(
         last_token_len = 0
         if num_output_top_logprobs:
             logprobs.top_logprobs = []
+
         for i, token_id in enumerate(token_ids):
             step_top_logprobs = top_logprobs[i]
-            if step_top_logprobs is not None:
-                token_logprob = step_top_logprobs[token_id].logprob
+            if step_top_logprobs is None:
+                token = self.tokenizer.decode(token_id)
+                logprobs.tokens.append(token)
+                logprobs.token_logprobs.append(None)
+                logprobs.top_logprobs.append(None)
             else:
-                token_logprob = None
-            token = step_top_logprobs[token_id].decoded_token
-            logprobs.tokens.append(token)
-            logprobs.token_logprobs.append(token_logprob)
+                token_logprob = step_top_logprobs[token_id].logprob
+                token = step_top_logprobs[token_id].decoded_token
+                logprobs.tokens.append(token)
+                logprobs.token_logprobs.append(token_logprob)
+
+                if num_output_top_logprobs:
+                    logprobs.top_logprobs.append({
+                        p.decoded_token: p.logprob
+                        for i, p in step_top_logprobs.items()
+                    } if step_top_logprobs else None)
+
             if len(logprobs.text_offset) == 0:
                 logprobs.text_offset.append(initial_text_offset)
             else:
                 logprobs.text_offset.append(logprobs.text_offset[-1] +
                                             last_token_len)
             last_token_len = len(token)
-
-            if num_output_top_logprobs:
-                logprobs.top_logprobs.append({
-                    p.decoded_token: p.logprob
-                    for i, p in step_top_logprobs.items()
-                } if step_top_logprobs else None)
         return logprobs
 
     def create_error_response(
@@ -164,12 +169,12 @@ def _maybe_get_lora(self, request) -> Optional[LoRARequest]:
         raise ValueError("The model `{request.model}` does not exist.")
 
     def _validate_prompt_and_tokenize(
-            self,
-            request: Union[ChatCompletionRequest, CompletionRequest],
-            prompt: Optional[str] = None,
-            prompt_ids: Optional[List[int]] = None,
-            truncate_prompt_tokens: Optional[conint(ge=1)] = None
-    ) -> List[int]:
+        self,
+        request: Union[ChatCompletionRequest, CompletionRequest],
+        prompt: Optional[str] = None,
+        prompt_ids: Optional[List[int]] = None,
+        truncate_prompt_tokens: Optional[conint(ge=1)] = None
+    ) -> Tuple[List[int], str]:
         if not (prompt or prompt_ids):
             raise ValueError("Either prompt or prompt_ids should be provided.")
         if (prompt and prompt_ids):
@@ -187,6 +192,8 @@ def _validate_prompt_and_tokenize(
         else:
             input_ids = prompt_ids
 
+        input_text = prompt if prompt is not None else self.tokenizer.decode(
+            prompt_ids)
         token_num = len(input_ids)
 
         if request.max_tokens is None:
@@ -201,4 +208,4 @@ def _validate_prompt_and_tokenize(
                 f"{request.max_tokens} in the completion). "
                 f"Please reduce the length of the messages or completion.", )
         else:
-            return input_ids
+            return input_ids, input_text


Test patch:
diff --git a/tests/entrypoints/test_openai_server.py b/tests/entrypoints/test_openai_server.py
--- a/tests/entrypoints/test_openai_server.py
+++ b/tests/entrypoints/test_openai_server.py
@@ -742,5 +742,36 @@ async def test_guided_grammar(server, client: openai.AsyncOpenAI):
     assert content.strip() == ground_truth
 
 
+@pytest.mark.parametrize(
+    # first test base model, then test loras
+    "model_name",
+    [MODEL_NAME, "zephyr-lora", "zephyr-lora2"],
+)
+async def test_echo_logprob_completion(server, client: openai.AsyncOpenAI,
+                                       model_name: str):
+    tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME)
+    # test using text and token IDs
+    for prompt in ("Hello, my name is", [0, 0, 0, 0, 0]):
+        completion = await client.completions.create(model=model_name,
+                                                     prompt=prompt,
+                                                     max_tokens=5,
+                                                     temperature=0.0,
+                                                     echo=True,
+                                                     logprobs=1)
+
+        prompt_text = tokenizer.decode(prompt) if isinstance(prompt,
+                                                             list) else prompt
+        assert (completion.choices[0].text is not None
+                and re.search(r"^" + prompt_text, completion.choices[0].text))
+        logprobs = completion.choices[0].logprobs
+        assert logprobs is not None
+        assert len(logprobs.text_offset) > 5
+        assert (len(logprobs.token_logprobs) > 5
+                and logprobs.token_logprobs[0] is None)
+        assert (len(logprobs.top_logprobs) > 5
+                and logprobs.top_logprobs[0] is None)
+        assert len(logprobs.tokens) > 5
+
+
 if __name__ == "__main__":
     pytest.main([__file__])