[diffusion] fix: fix suppressing error log on non-main ranks (#17712)

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -111,6 +111,41 @@ class GPUWorker:
             f"Worker {self.rank}: Initialized device, model, and distributed environment."
         )
 
+    def do_mem_analysis(self, output_batch: OutputBatch):
+        peak_memory_bytes = torch.cuda.max_memory_allocated()
+        output_batch.peak_memory_mb = peak_memory_bytes / (1024**2)
+        peak_memory_gb = peak_memory_bytes / (1024**3)
+        remaining_gpu_mem_gb = (
+            current_platform.get_device_total_memory() / (1024**3) - peak_memory_gb
+        )
+        can_stay_resident = self.get_can_stay_resident_components(remaining_gpu_mem_gb)
+        suggested_args = set()
+        component_to_arg = {
+            "vae": "--vae-cpu-offload",
+            "text_encoder": "--text-encoder-cpu-offload",
+            "text_encoder_2": "--text-encoder-cpu-offload",
+            "image_encoder": "--image-encoder-cpu-offload",
+        }
+
+        for component in can_stay_resident:
+            if component == "transformer":
+                if self.server_args.dit_layerwise_offload:
+                    suggested_args.add("--dit-layerwise-offload")
+                elif self.server_args.dit_cpu_offload:
+                    suggested_args.add("--dit-cpu-offload")
+            elif component in component_to_arg:
+                suggested_args.add(component_to_arg[component])
+
+        suggested_args_str = (
+            ", ".join(sorted(suggested_args)) if suggested_args else "None"
+        )
+        logger.info(
+            f"Peak GPU memory: {peak_memory_gb:.2f} GB, "
+            f"Remaining GPU memory at peak: {remaining_gpu_mem_gb:.2f} GB. "
+            f"Components that could stay resident (based on the last request workload): {can_stay_resident}. "
+            f"Related offload server args to disable: {suggested_args_str}"
+        )
+
     def execute_forward(self, batch: List[Req]) -> OutputBatch:
         """
         Execute a forward pass.
@@ -138,23 +173,8 @@ class GPUWorker:
             else:
                 output_batch = result
 
-            if self.rank == 0:
+            if self.rank == 0 and not req.suppress_logs:
-                peak_memory_bytes = torch.cuda.max_memory_allocated()
+                self.do_mem_analysis(output_batch)
-                output_batch.peak_memory_mb = peak_memory_bytes / (1024**2)
-                peak_memory_gb = peak_memory_bytes / (1024**3)
-                remaining_gpu_mem_gb = (
-                    current_platform.get_device_total_memory() / (1024**3)
-                    - peak_memory_gb
-                )
-                can_stay_resident = self.get_can_stay_resident_components(
-                    remaining_gpu_mem_gb
-                )
-                if not req.suppress_logs:
-                    logger.info(
-                        f"Peak GPU memory: {peak_memory_gb:.2f} GB, "
-                        f"Remaining GPU memory at peak: {remaining_gpu_mem_gb:.2f} GB. "
-                        f"Components that can stay resident: {can_stay_resident}"
-                    )
 
             duration_ms = (time.monotonic() - start_time) * 1000
             output_batch.timings.total_duration_ms = duration_ms

python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py

@@ -138,7 +138,7 @@ class InputValidationStage(PipelineStage):
 
             scale = max(ow / iw, oh / ih)
             img = img.resize((round(iw * scale), round(ih * scale)), Image.LANCZOS)
-            logger.info("resized img height: %s, img width: %s", img.height, img.width)
+            logger.debug("resized img height: %s, img width: %s", img.height, img.width)
 
             # center-crop
             x1 = (img.width - ow) // 2

python/sglang/multimodal_gen/runtime/server_args.py

@@ -12,7 +12,6 @@ import os
 import random
 import sys
 import tempfile
-from contextlib import contextmanager
 from dataclasses import field
 from enum import Enum
 from typing import Any, Optional
@@ -1090,8 +1089,6 @@ class PortArgs:
         )
 
 
-# TODO: not sure what _current_server_args is for, using a _global_server_args instead
-_current_server_args = None
 _global_server_args = None
 
 
@@ -1103,29 +1100,9 @@ def prepare_server_args(argv: list[str]) -> ServerArgs:
     ServerArgs.add_cli_args(parser)
     raw_args = parser.parse_args(argv)
     server_args = ServerArgs.from_cli_args(raw_args)
-    global _current_server_args
-    _current_server_args = server_args
     return server_args
 
 
-@contextmanager
-def set_current_server_args(server_args: ServerArgs):
-    """
-    Temporarily set the current sgl_diffusion config.
-    Used during model initialization.
-    We save the current sgl_diffusion config in a global variable,
-    so that all modules can access it, e.g. custom ops
-    can access the sgl_diffusion config to determine how to dispatch.
-    """
-    global _current_server_args
-    old_server_args = _current_server_args
-    try:
-        _current_server_args = server_args
-        yield
-    finally:
-        _current_server_args = old_server_args
-
-
 def set_global_server_args(server_args: ServerArgs):
     """
     Set the global sgl_diffusion config for each process
@@ -1134,17 +1111,7 @@ def set_global_server_args(server_args: ServerArgs):
     _global_server_args = server_args
 
 
-def get_current_server_args() -> ServerArgs | None:
+def get_global_server_args() -> ServerArgs:
-    if _current_server_args is None:
-        # in ci, usually when we test custom ops/modules directly,
-        # we don't set the sgl_diffusion config. In that case, we set a default
-        # config.
-        # TODO(will): may need to handle this for CI.
-        raise ValueError("Current sgl_diffusion args is not set.")
-    return _current_server_args
-
-
-def get_global_server_args() -> ServerArgs | None:
     if _global_server_args is None:
         # in ci, usually when we test custom ops/modules directly,
         # we don't set the sgl_diffusion config. In that case, we set a default

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

@@ -142,6 +142,7 @@ def get_is_local_main_process():
 
 
 def _log_process_aware(
+    server_log_level: int,
     level: int,
     logger_self: Logger,
     msg: object,
@@ -153,12 +154,12 @@ def _log_process_aware(
     """Helper function to log a message if the process rank matches the criteria."""
     is_main_process = get_is_main_process()
     is_local_main_process = get_is_local_main_process()
-
     should_log = (
         not main_process_only
         and not local_main_process_only
         or (main_process_only and is_main_process)
         or (local_main_process_only and is_local_main_process)
+        or server_log_level <= logging.DEBUG
     )
 
     if should_log:
@@ -234,6 +235,8 @@ def init_logger(name: str) -> _SGLDiffusionLogger:
 
     logger = logging.getLogger(name)
 
+    server_log_level = logger.getEffectiveLevel()
+
     # Patch instance methods
     setattr(logger, "info_once", MethodType(_print_info_once, logger))
     setattr(logger, "warning_once", MethodType(_print_warning_once, logger))
@@ -252,6 +255,7 @@ def init_logger(name: str) -> _SGLDiffusionLogger:
             **kwargs: Any,
         ) -> None:
             _log_process_aware(
+                server_log_level,
                 level,
                 self,
                 msg,
@@ -281,7 +285,7 @@ def init_logger(name: str) -> _SGLDiffusionLogger:
     setattr(
         logger,
         "error",
-        MethodType(_create_patched_method(logging.ERROR, False, True), logger),
+        MethodType(_create_patched_method(logging.ERROR, False, False), logger),
     )
 
     return cast(_SGLDiffusionLogger, logger)