Sglang Tracing: optimize trace_event_batch() (#13036)

Signed-off-by: Feng Su <sufeng@linux.alibaba.com>

python/sglang/srt/disaggregation/decode.py

@@ -879,8 +879,7 @@ class SchedulerDisaggregationDecodeMixin:
             ret = self.prepare_mlp_sync_batch(ret)
 
         if ret:
-            attrs = {"bid": hex(id(ret)), "batch_size": ret.batch_size()}
+            trace_event_batch("schedule", ret.reqs)
-            trace_event_batch("schedule", ret.reqs, attrs=attrs)
         return ret
 
     def get_new_prebuilt_batch(self: Scheduler) -> Optional[ScheduleBatch]:

python/sglang/srt/disaggregation/prefill.py

@@ -321,8 +321,7 @@ class SchedulerDisaggregationPrefillMixin:
             batch = self.prepare_mlp_sync_batch(batch)
 
         if batch:
-            attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()}
+            trace_event_batch("schedule", batch.reqs)
-            trace_event_batch("schedule", batch.reqs, attrs=attrs)
 
         return batch
 

python/sglang/srt/managers/scheduler.py

@@ -1702,8 +1702,7 @@ class Scheduler(
             ret = self.prepare_mlp_sync_batch(ret)
 
         if ret:
-            attrs = {"bid": hex(id(ret)), "batch_size": ret.batch_size()}
+            trace_event_batch("schedule", ret.reqs)
-            trace_event_batch("schedule", ret.reqs, attrs=attrs)
 
         return ret
 

python/sglang/srt/tracing/trace.py

@@ -678,10 +678,14 @@ def trace_event_batch(
     name: str,
     reqs: List[Req],
     ts: Optional[int] = None,
-    attrs: Dict[str, Any] = None,
+    attrs: Dict[str, Any] = {},
 ):
     if not tracing_enabled:
         return
 
+    bid = uuid.uuid4().hex[:8]
+    _attrs = {"bid": bid, "batch_size": len(reqs)}
+    _attrs.update(attrs)
+
     for req in reqs:
-        trace_event(name, req.rid, ts=ts, attrs=attrs)
+        trace_event(name, req.rid, ts=ts, attrs=_attrs)