From 6c4bf8a0be7f6bf0090ecd637ef65c5ae61a7d07 Mon Sep 17 00:00:00 2001
From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
Date: Wed, 17 Dec 2025 20:44:55 +0800
Subject: [PATCH] [diffusion] profiling: enhance trace export with gzip and
 integrity check (#15326)

Co-authored-by: Mick <mickjagger19@icloud.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
---
 .../multimodal_gen/runtime/utils/profiler.py  | 29 ++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/python/sglang/multimodal_gen/runtime/utils/profiler.py b/python/sglang/multimodal_gen/runtime/utils/profiler.py
index 18f8ad553..08842dd6b 100644
--- a/python/sglang/multimodal_gen/runtime/utils/profiler.py
+++ b/python/sglang/multimodal_gen/runtime/utils/profiler.py
@@ -1,3 +1,4 @@
+import gzip
 import os
 
 import torch
@@ -118,6 +119,12 @@ class SGLDiffusionProfiler:
         if dump_rank is None:
             dump_rank = self.rank
 
+        current_rank = (
+            torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+        )
+        if current_rank != dump_rank:
+            return
+
         try:
             os.makedirs(self.log_dir, exist_ok=True)
             sanitized_profile_mode_id = self.profile_mode_id.replace(" ", "_")
@@ -128,6 +135,26 @@ class SGLDiffusionProfiler:
                 )
             )
             self.profiler.export_chrome_trace(trace_path)
-            logger.info(f"Saved profiler traces to: {CYAN}{trace_path}{RESET}")
+
+            if self._check_trace_integrity(trace_path):
+                logger.info(f"Saved profiler traces to: {CYAN}{trace_path}{RESET}")
+            else:
+                logger.warning(f"Trace file may be corrupted: {trace_path}")
         except Exception as e:
             logger.error(f"Failed to save trace: {e}")
+
+    def _check_trace_integrity(self, trace_path: str) -> bool:
+        try:
+            if not os.path.exists(trace_path) or os.path.getsize(trace_path) == 0:
+                return False
+
+            with gzip.open(trace_path, "rb") as f:
+                content = f.read()
+                if content.count(b"\x1f\x8b") > 1:
+                    logger.warning("Multiple gzip headers detected")
+                    return False
+
+            return True
+        except Exception as e:
+            logger.warning(f"Trace file integrity check failed: {e}")
+            return False