diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index c1a0eec90..d684bdf45 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -16,7 +16,6 @@ from dataclasses import dataclass from enum import Enum, auto from typing import Any -from sglang.multimodal_gen import envs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import StoreBoolean @@ -208,10 +207,6 @@ class SamplingParams: if env_steps is not None and self.num_inference_steps is not None: self.num_inference_steps = int(env_steps) - # Auto-enable stage logging if dump path is provided - if self.perf_dump_path: - envs.SGLANG_DIFFUSION_STAGE_LOGGING = True - def _validate(self): """ check if the sampling params is correct by itself diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 79d23549e..961d305d9 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -9,6 +9,7 @@ from typing import List import torch from setproctitle import setproctitle +from sglang.multimodal_gen import envs from sglang.multimodal_gen.runtime.distributed import ( get_sp_group, maybe_init_distributed_environment_and_model_parallel, @@ -31,10 +32,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import ( globally_suppress_loggers, init_logger, ) -from sglang.multimodal_gen.runtime.utils.perf_logger import ( - PerformanceLogger, - StageProfiler, -) +from sglang.multimodal_gen.runtime.utils.perf_logger import PerformanceLogger logger = init_logger(__name__) @@ -129,7 +127,8 @@ class GPUWorker: duration_ms = (time.monotonic() - start_time) * 1000 output_batch.timings.total_duration_ms = duration_ms - if StageProfiler.metrics_enabled(): + # TODO: extract to avoid duplication + if req.perf_dump_path is not None or envs.SGLANG_DIFFUSION_STAGE_LOGGING: PerformanceLogger.log_request_summary(timings=output_batch.timings) except Exception as e: logger.error( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py index 38a917eb6..5f8ef94e0 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py @@ -201,7 +201,12 @@ class PipelineStage(ABC): raise # Execute the actual stage logic with unified profiling - with StageProfiler(stage_name, logger=logger, timings=batch.timings): + with StageProfiler( + stage_name, + logger=logger, + timings=batch.timings, + perf_dump_path_provided=batch.perf_dump_path is not None, + ): result = self.forward(batch, server_args) # Post-execution output verification diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 7580f7bf4..d5862a258 100755 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -959,7 +959,10 @@ class DenoisingStage(PipelineStage): with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t_host in enumerate(timesteps_cpu): with StageProfiler( - f"denoising_step_{i}", logger=logger, timings=batch.timings + f"denoising_step_{i}", + logger=logger, + timings=batch.timings, + perf_dump_path_provided=batch.perf_dump_path is not None, ): t_int = int(t_host.item()) t_device = timesteps[i] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py index edf43f4dc..c67e9d017 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py @@ -98,7 +98,10 @@ class DmdDenoisingStage(DenoisingStage): continue with StageProfiler( - f"denoising_step_{i}", logger=logger, timings=batch.timings + f"denoising_step_{i}", + logger=logger, + timings=batch.timings, + perf_dump_path_provided=batch.perf_dump_path is not None, ): # Expand latents for I2V noise_latents = latents.clone() diff --git a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py index 20a1671e1..423722104 100644 --- a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py +++ b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py @@ -131,31 +131,26 @@ class StageProfiler: logger: _SGLDiffusionLogger, timings: Optional["RequestTimings"], simple_log: bool = False, + perf_dump_path_provided: bool = False, ): self.stage_name = stage_name self.timings = timings self.logger = logger self.simple_log = simple_log self.start_time = 0.0 - - self._metrics_enabled = StageProfiler.metrics_enabled() - - @staticmethod - def metrics_enabled(): - # Check env var at runtime to ensure we pick up changes (e.g. from CLI args) - return envs.SGLANG_DIFFUSION_STAGE_LOGGING + self.enabled = perf_dump_path_provided or envs.SGLANG_DIFFUSION_STAGE_LOGGING def __enter__(self): if self.simple_log: self.logger.info(f"[{self.stage_name}] started...") - if (self._metrics_enabled and self.timings) or self.simple_log: + if (self.enabled and self.timings) or self.simple_log: self.start_time = time.perf_counter() return self def __exit__(self, exc_type, exc_val, exc_tb): - if not ((self._metrics_enabled and self.timings) or self.simple_log): + if not ((self.enabled and self.timings) or self.simple_log): return False execution_time_s = time.perf_counter() - self.start_time @@ -175,7 +170,7 @@ class StageProfiler: f"[{self.stage_name}] finished in {execution_time_s:.4f} seconds", ) - if self._metrics_enabled and self.timings: + if self.enabled and self.timings: if "denoising_step_" in self.stage_name: index = int(self.stage_name[len("denoising_step_") :]) self.timings.record_steps(index, execution_time_s)