Fix/partial gen from waiting queue miss metadata (#17610)
This commit is contained in:
@@ -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(
|
||||
|
||||
@@ -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")
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user