diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 6c277d2fe..e35b5e982 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -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( diff --git a/test/registered/scheduler/test_abort.py b/test/registered/scheduler/test_abort.py index ade0a3f2a..9caaa2cac 100644 --- a/test/registered/scheduler/test_abort.py +++ b/test/registered/scheduler/test_abort.py @@ -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")