[Fix] Guard conditional top-logprob keys in the completions echo path (#34776)

Co-authored-by: James Liu <jamesl@modal.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
This commit is contained in:
gilfordting
2026-09-17 19:47:26 -07:00
committed by GitHub
co-authored by James Liu hnyls2002 Liangsheng Yin
parent 740f57a02c
commit db39b7f961
3 changed files with 115 additions and 11 deletions
@@ -74,7 +74,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
) -> tuple[GenerateReqInput, CompletionRequest]: ) -> tuple[GenerateReqInput, CompletionRequest]:
"""Convert OpenAI completion request to internal format""" """Convert OpenAI completion request to internal format"""
# NOTE: with openai API, the prompt's logprobs are always not computed # NOTE: with openai API, the prompt's logprobs are always not computed
if request.echo and request.logprobs: if request.echo and request.logprobs is not None:
logger.warning( logger.warning(
"Echo is not compatible with logprobs. " "Echo is not compatible with logprobs. "
"To compute logprobs of input prompt, please use the native /generate API." "To compute logprobs of input prompt, please use the native /generate API."
@@ -274,6 +274,9 @@ class OpenAIServingCompletion(OpenAIServingBase):
content["meta_info"] content["meta_info"]
) )
finish_reason = content["meta_info"].get("finish_reason", None)
finish_reason_type = finish_reason["type"] if finish_reason else None
is_first_chunk = index not in stream_offsets is_first_chunk = index not in stream_offsets
offset = stream_offsets.get(index, 0) offset = stream_offsets.get(index, 0)
# Handle echo for first chunk # Handle echo for first chunk
@@ -286,11 +289,13 @@ class OpenAIServingCompletion(OpenAIServingBase):
logprobs = None logprobs = None
if request.logprobs is not None: if request.logprobs is not None:
# The first chunk and echo is enabled. # The first chunk and echo is enabled.
if is_first_chunk and request.echo: if is_first_chunk and request.echo and request.logprobs:
input_token_logprobs = content["meta_info"][ input_token_logprobs = content["meta_info"][
"input_token_logprobs" "input_token_logprobs"
] ]
input_top_logprobs = content["meta_info"]["input_top_logprobs"] input_top_logprobs = content["meta_info"].get(
"input_top_logprobs", None
)
else: else:
input_token_logprobs = None input_token_logprobs = None
input_top_logprobs = None input_top_logprobs = None
@@ -343,8 +348,6 @@ class OpenAIServingCompletion(OpenAIServingBase):
else: else:
delta = text[offset:] delta = text[offset:]
stream_offsets[index] = len(content["text"]) stream_offsets[index] = len(content["text"])
finish_reason = content["meta_info"].get("finish_reason", None)
finish_reason_type = finish_reason["type"] if finish_reason else None
# Abort with an explicit error status_code is a system error # Abort with an explicit error status_code is a system error
# (timeout, OOM, validation): emit a streaming error chunk. # (timeout, OOM, validation): emit a streaming error chunk.
@@ -575,9 +578,11 @@ class OpenAIServingCompletion(OpenAIServingBase):
# Handle logprobs # Handle logprobs
logprobs = None logprobs = None
if request.logprobs is not None: if request.logprobs is not None:
if echo: if echo and request.logprobs:
input_token_logprobs = ret_item["meta_info"]["input_token_logprobs"] input_token_logprobs = ret_item["meta_info"]["input_token_logprobs"]
input_top_logprobs = ret_item["meta_info"]["input_top_logprobs"] input_top_logprobs = ret_item["meta_info"].get(
"input_top_logprobs", None
)
else: else:
input_token_logprobs = None input_token_logprobs = None
input_top_logprobs = None input_top_logprobs = None
@@ -2677,8 +2677,9 @@ class ServingChatTestCase(unittest.TestCase):
"status_code": err_code, "status_code": err_code,
"message": err_msg, "message": err_msg,
}, },
"output_token_logprobs": None, "output_token_logprobs": [],
"output_top_logprobs": None, "output_token_logprobs_length": 0,
"output_top_logprobs": [],
}, },
"index": 0, "index": 0,
} }
@@ -2691,6 +2692,8 @@ class ServingChatTestCase(unittest.TestCase):
temperature=0.7, temperature=0.7,
max_tokens=100, max_tokens=100,
stream=True, stream=True,
logprobs=True,
top_logprobs=5,
) )
with patch( with patch(
@@ -270,8 +270,9 @@ class ServingCompletionTestCase(unittest.TestCase):
"status_code": err_code, "status_code": err_code,
"message": err_msg, "message": err_msg,
}, },
"output_token_logprobs": None, "output_token_logprobs": [],
"output_top_logprobs": None, "output_token_logprobs_length": 0,
"output_top_logprobs": [],
}, },
"index": 0, "index": 0,
} }
@@ -283,6 +284,7 @@ class ServingCompletionTestCase(unittest.TestCase):
prompt="Hello world", prompt="Hello world",
max_tokens=100, max_tokens=100,
stream=True, stream=True,
logprobs=5,
) )
adapted_request, _ = self.sc._convert_to_internal_request(req) adapted_request, _ = self.sc._convert_to_internal_request(req)
@@ -318,6 +320,100 @@ class ServingCompletionTestCase(unittest.TestCase):
self.assertGreaterEqual(len(chunks), 2) self.assertGreaterEqual(len(chunks), 2)
self.assertIn("error", chunks[0]) self.assertIn("error", chunks[0])
def test_echo_with_zero_logprobs_streaming(self):
"""logprobs=0 requests token logprobs without top-logprobs, so the
scheduler never fills the top-logprob keys. The echo branch must not
assume they are present."""
async def _mock_generate(*args, **kwargs):
yield {
"text": "Hello world",
"meta_info": {
"id": "cmpl-test",
"prompt_tokens": 2,
"completion_tokens": 2,
"cached_tokens": 0,
"finish_reason": {"type": "stop"},
# top_logprobs_num == 0, so no input/output top-logprob keys.
"input_token_logprobs": [],
"output_token_logprobs": [
(-0.1, 3, "Hello"),
(-0.2, 4, " world"),
],
"output_token_logprobs_length": 2,
},
"index": 0,
}
self.sc.tokenizer_manager.generate_request = _mock_generate
req = CompletionRequest(
model="x",
prompt="Hi",
max_tokens=100,
stream=True,
echo=True,
logprobs=0,
)
adapted_request, _ = self.sc._convert_to_internal_request(req)
async def run_stream():
return [
chunk
async for chunk in self.sc._generate_completion_stream(
adapted_request, req, self.fastapi_request
)
]
loop = get_or_create_event_loop()
chunks = loop.run_until_complete(run_stream())
# Assert on the payload, not just termination: a regression that
# silently drops logprobs still produces a well-formed stream.
self.assertNotIn("error", chunks[0])
self.assertEqual(chunks[-1], "data: [DONE]\n\n")
choice = json.loads(chunks[0][len("data: ") :])["choices"][0]
self.assertTrue(choice["text"].startswith("Hi"))
logprobs = choice["logprobs"]
# logprobs=0 asks for token logprobs but no top-logprobs, and the echoed
# prompt contributes none because input logprobs were never requested.
self.assertEqual(logprobs["tokens"], ["Hello", " world"])
self.assertEqual(logprobs["token_logprobs"], [-0.1, -0.2])
self.assertEqual(logprobs["top_logprobs"], [])
def test_echo_with_zero_logprobs_non_streaming(self):
"""Same contract on the non-streaming path."""
req = CompletionRequest(
model="x",
prompt="Hi",
max_tokens=100,
echo=True,
logprobs=0,
)
ret = [
{
"text": " world",
"meta_info": {
"id": "cmpl-test",
"prompt_tokens": 2,
"completion_tokens": 1,
"cached_tokens": 0,
"finish_reason": {"type": "stop"},
"weight_version": "v1",
"input_token_logprobs": [],
"output_token_logprobs": [(-0.1, 3, " world")],
"output_token_logprobs_length": 1,
},
}
]
response = self.sc._build_completion_response(req, ret, 1234567890)
self.assertEqual(len(response.choices), 1)
self.assertEqual(response.choices[0].logprobs.token_logprobs, [-0.1])
self.assertEqual(response.choices[0].logprobs.top_logprobs, [])
def test_streaming_token_ids_deltas_cover_output_exactly(self): def test_streaming_token_ids_deltas_cover_output_exactly(self):
req = CompletionRequest( req = CompletionRequest(
model="x", model="x",