Pass finished_before to do_tracing() to avoid exporting duplicate spans

vllm-project · Oct 2, 2024 · 3b2916c · 3b2916c
1 parent cd3cfe5
commit 3b2916c
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 19 deletions.
diff --git a/tests/tracing/test_tracing.py b/tests/tracing/test_tracing.py
@@ -89,17 +89,17 @@ def test_traces(trace_service):
 
     request = trace_service.request
     assert len(request.resource_spans) == 1, (
-        f"Expected 1 resource span, but got {len(trace_service.request.resource_spans)}"
-    )
+        f"Expected 1 resource span, "
+        f"but got {len(request.resource_spans)}")
     assert len(request.resource_spans[0].scope_spans) == 1, (
-        f"Expected 1 scope span, but got {len(trace_service.request.resource_spans[0].scope_spans)}"
-    )
+        f"Expected 1 scope span, "
+        f"but got {len(request.resource_spans[0].scope_spans)}")
     assert len(request.resource_spans[0].scope_spans[0].spans) == 1, (
-        f"Expected 1 span, but got {len(trace_service.request.resource_spans[0].scope_spans[0].spans)}"
-    )
+        f"Expected 1 span, "
+        f"but got {len(request.resource_spans[0].scope_spans[0].spans)}")
 
-    attributes = decode_attributes(trace_service.request.resource_spans[0].
-                                   scope_spans[0].spans[0].attributes)
+    attributes = decode_attributes(
+        request.resource_spans[0].scope_spans[0].spans[0].attributes)
     assert attributes.get(SpanAttributes.LLM_RESPONSE_MODEL) == model
     assert attributes.get(
         SpanAttributes.LLM_REQUEST_ID) == outputs[0].request_id
@@ -157,17 +157,17 @@ def test_traces_with_detailed_steps(trace_service):
 
     request = trace_service.request
     assert len(request.resource_spans) == 1, (
-        f"Expected 1 resource span, but got {len(trace_service.request.resource_spans)}"
-    )
+        f"Expected 1 resource span, "
+        f"but got {len(request.resource_spans)}")
     assert len(request.resource_spans[0].scope_spans) == 1, (
-        f"Expected 1 scope span, but got {len(trace_service.request.resource_spans[0].scope_spans)}"
-    )
+        f"Expected 1 scope span, "
+        f"but got {len(request.resource_spans[0].scope_spans)}")
     assert len(request.resource_spans[0].scope_spans[0].spans) == 1, (
-        f"Expected 1 span, but got {len(trace_service.request.resource_spans[0].scope_spans[0].spans)}"
-    )
+        f"Expected 1 span, "
+        f"but got {len(request.resource_spans[0].scope_spans[0].spans)}")
 
-    attributes = decode_attributes(trace_service.request.resource_spans[0].
-                                   scope_spans[0].spans[0].attributes)
+    attributes = decode_attributes(
+        request.resource_spans[0].scope_spans[0].spans[0].attributes)
     assert attributes.get(SpanAttributes.LLM_RESPONSE_MODEL) == model
     assert attributes.get(
         SpanAttributes.LLM_REQUEST_ID) == outputs[0].request_id

diff --git a/vllm/engine/llm_engine.py b/vllm/engine/llm_engine.py
@@ -1224,7 +1224,7 @@ def update_prefill_num_computed_tokens(
                               skip)
 
             # Tracing
-            self.do_tracing(scheduler_outputs)
+            self.do_tracing(scheduler_outputs, finished_before)
 
         return None
 
@@ -1798,11 +1798,18 @@ def stop_profile(self) -> None:
     def is_tracing_enabled(self) -> bool:
         return self.tracer is not None
 
-    def do_tracing(self, scheduler_outputs: SchedulerOutputs) -> None:
+    def do_tracing(self,
+                   scheduler_outputs: SchedulerOutputs,
+                   finished_before: Optional[List[int]] = None) -> None:
         if self.tracer is None:
             return
 
-        for scheduled_seq_group in scheduler_outputs.scheduled_seq_groups:
+        for idx, scheduled_seq_group in enumerate(
+                scheduler_outputs.scheduled_seq_groups):
+            # Skip double tracing when using async output proc
+            if finished_before and idx in finished_before:
+                continue
+
             seq_group = scheduled_seq_group.seq_group
             if seq_group.is_finished():
                 self.create_trace_span(seq_group)