Fix/partial gen from waiting queue miss metadata (#17610)

This commit is contained in:
JD
2026-02-13 19:04:08 -08:00
committed by GitHub
parent 45a4697d45
commit f6c18c3a85
2 changed files with 40 additions and 13 deletions

View File

@@ -2164,6 +2164,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
return
state = self.rid_to_state[recv_obj.rid]
state.finished = True
state.finished_time = time.time()
abort_message = recv_obj.abort_message or "Abort in waiting queue"
finish_reason = {
@@ -2172,7 +2173,12 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
}
if recv_obj.finished_reason:
finish_reason = recv_obj.finished_reason
meta_info = {"id": recv_obj.rid, "finish_reason": finish_reason}
meta_info = {
"id": recv_obj.rid,
"finish_reason": finish_reason,
"weight_version": self.server_args.weight_version,
"e2e_latency": state.finished_time - state.created_time,
}
is_stream = getattr(state.obj, "stream", False)
if getattr(state.obj, "return_logprob", False):
self.add_logprob_to_meta_info(

View File

@@ -165,6 +165,8 @@ class TestAbortAllWithRetraction(CustomTestCase):
"max_new_tokens": 4000,
"ignore_eos": True,
},
"return_logprob": True,
"top_logprobs_num": 3,
},
)
return response.json()
@@ -185,21 +187,40 @@ class TestAbortAllWithRetraction(CustomTestCase):
)
abort_in_queue_count = 0
abort_in_queue_with_none_empty_text = 0
abort_in_queue_with_partial_gen = 0
for future in as_completed(futures):
self.assertEqual(
future.result()["meta_info"]["finish_reason"]["type"], "abort"
)
if (
future.result()["meta_info"]["finish_reason"]["message"]
== "Abort in waiting queue"
):
result = future.result()
meta_info = result["meta_info"]
finish_reason = meta_info.get("finish_reason", {})
self.assertEqual(finish_reason.get("type"), "abort")
if finish_reason.get("message") == "Abort in waiting queue":
abort_in_queue_count += 1
if len(future.result()["output_ids"]) > 0:
abort_in_queue_with_none_empty_text += 1
assert abort_in_queue_count > 0
assert abort_in_queue_with_none_empty_text > 0
output_ids = result.get("output_ids", [])
if len(output_ids) > 0:
abort_in_queue_with_partial_gen += 1
self.assertEqual(
meta_info.get("completion_tokens"), len(output_ids)
)
self.assertGreater(len(result.get("text", "")), 0)
self.assertIsNotNone(meta_info.get("weight_version"))
self.assertGreater(meta_info.get("e2e_latency"), 0)
for logprob_key in [
"output_token_logprobs",
"output_top_logprobs",
]:
self.assertEqual(
len(meta_info.get(logprob_key, [])),
len(output_ids),
f"Length of '{logprob_key}' should match output_ids length",
)
self.assertGreater(abort_in_queue_count, 0)
self.assertGreater(abort_in_queue_with_partial_gen, 0)
print("Finished test_abort_all_with_retraction")