From 2c5679f314760fe32bb9ea76a383c805ccb81f21 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 24 Dec 2025 23:26:09 +0800 Subject: [PATCH] [diffusion] refactor: unify the profiling api for all executors (#15718) --- .../pipelines_core/composed_pipeline_base.py | 6 ++--- .../executors/parallel_executor.py | 17 +++--------- .../executors/pipeline_executor.py | 27 ++++++++++++++----- .../pipelines_core/executors/sync_executor.py | 9 +++---- .../multimodal_gen/runtime/utils/profiler.py | 17 +++++------- 5 files changed, 38 insertions(+), 38 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index 4375138a7..39094e118 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py index 52d85d32b..72622e254 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py +++ b/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: - rank = get_classifier_free_guidance_rank() - - 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) - + ) -> OutputBatch: + batch = self._execute(stages, batch, server_args) return batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py index 772981380..8e87cc181 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py +++ b/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(export_trace=should_export, dump_rank=dump_rank) + profiler.stop(dump_rank=dump_rank) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py index 47491a111..b7eb3ca7b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/sync_executor.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/utils/profiler.py b/python/sglang/multimodal_gen/runtime/utils/profiler.py index 08842dd6b..037133e8a 100644 --- a/python/sglang/multimodal_gen/runtime/utils/profiler.py +++ b/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): - 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 + def _export_trace(self): 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)