[diffusion] refactor: unify the profiling api for all executors (#15718)

python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py

@@ -22,7 +22,7 @@ from sglang.multimodal_gen.runtime.loader.component_loader import (
 from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import (
     PipelineExecutor,
 )
-from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages import PipelineStage
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
@@ -345,7 +345,7 @@ class ComposedPipelineBase(ABC):
         self,
         batch: Req,
         server_args: ServerArgs,
-    ) -> Req:
+    ) -> OutputBatch:
         """
         Generate a video or image using the pipeline.
 
@@ -372,4 +372,4 @@ class ComposedPipelineBase(ABC):
             main_process_only=True,
         )
 
-        return self.executor.execute(self.stages, batch, server_args)
+        return self.executor.execute_with_profiling(self.stages, batch, server_args)

python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py

@@ -8,13 +8,13 @@ from sglang.multimodal_gen.runtime.distributed import get_sp_group
 from sglang.multimodal_gen.runtime.distributed.parallel_state import (
     get_cfg_group,
     get_classifier_free_guidance_rank,
-    get_world_rank,
 )
 from sglang.multimodal_gen.runtime.pipelines_core import Req
 from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import (
     PipelineExecutor,
     Timer,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
     PipelineStage,
     StageParallelismType,
@@ -57,7 +57,7 @@ class ParallelExecutor(PipelineExecutor):
         stages: List[PipelineStage],
         batch: Req,
         server_args: ServerArgs,
-    ) -> Req:
+    ) -> OutputBatch:
         """
         Execute all pipeline stages respecting their declared parallelism type.
         """
@@ -95,15 +95,6 @@ class ParallelExecutor(PipelineExecutor):
         stages: List[PipelineStage],
         batch: Req,
         server_args: ServerArgs,
-    ) -> Req:
+    ) -> OutputBatch:
-        rank = get_classifier_free_guidance_rank()
+        batch = self._execute(stages, batch, server_args)
-
-        if batch.profile and batch.profile_all_stages:
-            world_rank = get_world_rank()
-        else:
-            world_rank = 0
-
-        with self.profile_execution(batch, check_rank=rank, dump_rank=world_rank):
-            batch = self._execute(stages, batch, server_args)
-
         return batch

python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py

@@ -9,7 +9,8 @@ import contextlib
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, List
 
-from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.distributed import get_world_rank
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
@@ -43,13 +44,25 @@ class PipelineExecutor(ABC):
     def __init__(self, server_args):
         self.server_args = server_args
 
+    def execute_with_profiling(
+        self,
+        stages: List["PipelineStage"],
+        batch: Req,
+        server_args: ServerArgs,
+    ) -> OutputBatch:
+
+        with self.profile_execution(batch, dump_rank=0):
+            batch = self.execute(stages, batch, server_args)
+
+        return batch
+
     @abstractmethod
     def execute(
         self,
         stages: List["PipelineStage"],
         batch: Req,
         server_args: ServerArgs,
-    ) -> Req:
+    ) -> OutputBatch:
         """
         Execute the pipeline stages.
 
@@ -64,20 +77,23 @@ class PipelineExecutor(ABC):
         raise NotImplementedError
 
     @contextlib.contextmanager
-    def profile_execution(self, batch: Req, check_rank: int = 0, dump_rank: int = 0):
+    def profile_execution(self, batch: Req, dump_rank: int = 0):
         """
         Context manager for profiling execution.
         """
         do_profile = batch.profile
 
         if not do_profile:
+            # fast forward
             yield
             return
 
         request_id = batch.request_id
+        rank = get_world_rank()
+
         profiler = SGLDiffusionProfiler(
             request_id=request_id,
-            rank=check_rank,
+            rank=rank,
             full_profile=batch.profile_all_stages,
             num_steps=batch.num_profiled_timesteps,
             num_inference_steps=batch.num_inference_steps,
@@ -85,5 +101,4 @@ class PipelineExecutor(ABC):
         try:
             yield
         finally:
-            should_export = check_rank == 0
+            profiler.stop(dump_rank=dump_rank)
-            profiler.stop(export_trace=should_export, dump_rank=dump_rank)

python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py

@@ -11,7 +11,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor im
     SGLDiffusionProfiler,
     Timer,
 )
-from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages import PipelineStage
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 
@@ -26,7 +26,7 @@ class SyncExecutor(PipelineExecutor):
         stages: List[PipelineStage],
         batch: Req,
         server_args: ServerArgs,
-    ) -> Req:
+    ) -> OutputBatch:
         """
         Execute all pipeline stages sequentially.
         """
@@ -44,12 +44,11 @@ class SyncExecutor(PipelineExecutor):
         stages: List[PipelineStage],
         batch: Req,
         server_args: ServerArgs,
-    ) -> Req:
+    ) -> OutputBatch:
         """
         Execute the pipeline stages sequentially.
         """
 
-        with self.profile_execution(batch, check_rank=0, dump_rank=0):
+        batch = self.run_profile_all_stages(stages, batch, server_args)
-            batch = self.run_profile_all_stages(stages, batch, server_args)
 
         return batch

python/sglang/multimodal_gen/runtime/utils/profiler.py

@@ -111,19 +111,14 @@ class SGLDiffusionProfiler:
         self.profiler.stop()
 
         if export_trace:
-            self._export_trace(dump_rank)
+            if dump_rank is not None and dump_rank != self.rank:
+                pass
+            else:
+                self._export_trace()
 
         SGLDiffusionProfiler._instance = None
 
-    def _export_trace(self, dump_rank: int | None = None):
+    def _export_trace(self):
-        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)
@@ -131,7 +126,7 @@ class SGLDiffusionProfiler:
             trace_path = os.path.abspath(
                 os.path.join(
                     self.log_dir,
-                    f"{self.request_id}-{sanitized_profile_mode_id}-global-rank{dump_rank}.trace.json.gz",
+                    f"{self.request_id}-{sanitized_profile_mode_id}-global-rank{self.rank}.trace.json.gz",
                 )
             )
             self.profiler.export_chrome_trace(trace_path)