diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index ea77294d2..1069ced6a 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -2494,11 +2494,12 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): state.time_stats.set_first_token_time() if state.finished: - if state.time_stats.trace_ctx.tracing_enable: - state.time_stats.trace_ctx.trace_set_root_attrs( - self.convert_to_span_attrs(state, recv_obj, i) - ) - state.time_stats.set_finished_time() + span_attrs = ( + self.convert_to_span_attrs(state, recv_obj, i) + if state.time_stats.trace_ctx.tracing_enable + else None + ) + state.time_stats.set_finished_time(span_attrs=span_attrs) meta_info["e2e_latency"] = state.time_stats.get_e2e_latency() if get_spec().speculative_algorithm: @@ -3658,8 +3659,8 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): [finish_reason] ) - # Latency attributes - span_attrs.update(state.time_stats.convert_to_gen_ai_span_attrs()) + # Latency attributes are added by set_finished_time(), which stamps + # finished_time before deriving them. return span_attrs diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py index 539634de6..4c174950a 100644 --- a/python/sglang/srt/observability/req_time_stats.py +++ b/python/sglang/srt/observability/req_time_stats.py @@ -415,12 +415,18 @@ class APIServerReqTimeStats(ReqTimeStatsBase): convert_time_to_realtime_ns(ts), ) - def set_finished_time(self, ts=None): + def set_finished_time(self, ts=None, span_attrs=None): ts = ts or time.perf_counter() self.finished_time = ts if self.trace_ctx.tracing_enable: - self.trace_ctx.trace_req_finish(convert_time_to_realtime_ns(ts)) + # The latency attrs are derived from finished_time and the root span is + # closed below, so they must be merged in here rather than by the caller. + attrs = dict(span_attrs) if span_attrs else {} + attrs.update(self.convert_to_gen_ai_span_attrs()) + self.trace_ctx.trace_req_finish( + convert_time_to_realtime_ns(ts), attrs=attrs + ) def set_first_token_time(self, ts=None): ts = ts or time.perf_counter() diff --git a/test/registered/unit/observability/test_req_time_stats.py b/test/registered/unit/observability/test_req_time_stats.py index 16b9169ea..40237b4b7 100644 --- a/test/registered/unit/observability/test_req_time_stats.py +++ b/test/registered/unit/observability/test_req_time_stats.py @@ -1,4 +1,4 @@ -"""Unit tests for ReqTimeStats IPC serialization. +"""Unit tests for ReqTimeStats. ReqTimeStatsBase.__setstate__ rebases perf_counter fields onto the receiving process's clock anchor. Rebasing a field that was never stamped (0.0) turns @@ -8,6 +8,11 @@ server never stamps prefill_finished_time locally; if the sentinel arrives at the tokenizer as an epsilon, first-token bookkeeping mistakes it for a real stamp and the TTFT / inter-token-latency histograms record ~node-uptime-sized garbage samples. + +The same "was this stamped?" rule governs the gen_ai.latency.* span attrs: the +ones derived from finished_time are only emitted once it is stamped, and +set_finished_time() closes the trace root span, so they have to be derived +inside that call rather than by the caller. """ import pickle @@ -15,6 +20,7 @@ import unittest from unittest import mock import sglang.srt.observability.req_time_stats as rts +from sglang.srt.observability.trace import SpanAttributes from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -40,5 +46,85 @@ class TestSetstatePreservesUnsetTimeSentinels(CustomTestCase): self.assertAlmostEqual(hop2.wait_queue_entry_time, 123.456 - 9.0) +class TestConvertToGenAiSpanAttrs(CustomTestCase): + def _stats_after_first_token(self) -> rts.APIServerReqTimeStats: + stats = rts.APIServerReqTimeStats() + stats.created_time = 1.0 + stats.api_server_dispatch_finish_time = 1.1 + stats.first_token_time = 1.5 + return stats + + def test_prefill_without_finished_time_omits_e2e_and_decode(self): + attrs = self._stats_after_first_token().convert_to_gen_ai_span_attrs() + self.assertIn(SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_PREFILL, attrs) + self.assertIn(SpanAttributes.GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN, attrs) + self.assertNotIn(SpanAttributes.GEN_AI_LATENCY_E2E, attrs) + self.assertNotIn(SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_DECODE, attrs) + self.assertNotIn(SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_INFERENCE, attrs) + + def test_finished_time_populates_e2e_and_decode(self): + stats = self._stats_after_first_token() + stats.finished_time = 2.0 + attrs = stats.convert_to_gen_ai_span_attrs() + self.assertAlmostEqual(attrs[SpanAttributes.GEN_AI_LATENCY_E2E], 1.0) + self.assertAlmostEqual( + attrs[SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_DECODE], 0.5 + ) + self.assertAlmostEqual( + attrs[SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_PREFILL], 0.4 + ) + self.assertAlmostEqual( + attrs[SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_INFERENCE], 0.9 + ) + + +class TestSetFinishedTimeSpanAttrs(CustomTestCase): + def _tracing_stats(self) -> rts.APIServerReqTimeStats: + stats = rts.APIServerReqTimeStats() + stats.trace_ctx = mock.MagicMock() + stats.trace_ctx.tracing_enable = True + return stats + + def test_passes_caller_attrs_into_trace_req_finish(self): + stats = self._tracing_stats() + + stats.set_finished_time(ts=1.25, span_attrs={"gen_ai.request.id": "rid-1"}) + + self.assertEqual(stats.finished_time, 1.25) + stats.trace_ctx.trace_req_finish.assert_called_once_with( + mock.ANY, attrs={"gen_ai.request.id": "rid-1"} + ) + + def test_merges_latency_attrs_derived_from_finished_time(self): + stats = self._tracing_stats() + stats.created_time = 1.0 + stats.api_server_dispatch_finish_time = 1.1 + stats.first_token_time = 1.5 + + stats.set_finished_time(ts=2.0, span_attrs={"gen_ai.request.id": "rid-1"}) + + attrs = stats.trace_ctx.trace_req_finish.call_args.kwargs["attrs"] + self.assertEqual(attrs["gen_ai.request.id"], "rid-1") + self.assertAlmostEqual(attrs[SpanAttributes.GEN_AI_LATENCY_E2E], 1.0) + self.assertAlmostEqual( + attrs[SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_DECODE], 0.5 + ) + self.assertAlmostEqual( + attrs[SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_INFERENCE], 0.9 + ) + self.assertAlmostEqual( + attrs[SpanAttributes.GEN_AI_LATENCY_TIME_IN_MODEL_PREFILL], 0.4 + ) + + def test_caller_attrs_not_mutated(self): + stats = self._tracing_stats() + stats.created_time = 1.0 + caller_attrs = {"gen_ai.request.id": "rid-1"} + + stats.set_finished_time(ts=2.0, span_attrs=caller_attrs) + + self.assertEqual(caller_attrs, {"gen_ai.request.id": "rid-1"}) + + if __name__ == "__main__": unittest.main()