[diffusion] UX: suppress excessive loggers (#14900)

python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py

@@ -23,7 +23,10 @@ from sglang.multimodal_gen.runtime.distributed.device_communicators.base_device_
 from sglang.multimodal_gen.runtime.distributed.device_communicators.cpu_communicator import (
     CpuCommunicator,
 )
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.logging_utils import (
+    init_logger,
+    suppress_stdout,
+)
 
 try:
     import torch_musa  # noqa: F401
@@ -172,7 +175,8 @@ class GroupCoordinator:
             )
             # a group with `gloo` backend, to allow direct coordination between
             # processes through the CPU.
-            cpu_group = torch.distributed.new_group(ranks, backend="gloo")
+            with suppress_stdout():
+                cpu_group = torch.distributed.new_group(ranks, backend="gloo")
             if self.rank in ranks:
                 self.ranks = ranks
                 self.world_size = len(ranks)
@@ -803,7 +807,8 @@ class PipelineGroupCoordinator(GroupCoordinator):
                 )
                 # a group with `gloo` backend, to allow direct coordination between
                 # processes through the CPU.
-                cpu_group = torch.distributed.new_group(ranks, backend="gloo")
+                with suppress_stdout():
+                    cpu_group = torch.distributed.new_group(ranks, backend="gloo")
                 if self.rank in ranks:
                     self.ranks = ranks
                     self.world_size = len(ranks)
@@ -826,8 +831,9 @@ class PipelineGroupCoordinator(GroupCoordinator):
                 )
                 # a group with `gloo` backend, to allow direct coordination between
                 # processes through the CPU.
-                cpu_group_0_1 = torch.distributed.new_group(ranks, backend="gloo")
+                with suppress_stdout():
-                cpu_group_1_0 = torch.distributed.new_group(ranks, backend="gloo")
+                    cpu_group_0_1 = torch.distributed.new_group(ranks, backend="gloo")
+                    cpu_group_1_0 = torch.distributed.new_group(ranks, backend="gloo")
                 if self.rank in ranks:
                     self.ranks = ranks
                     self.world_size = len(ranks)

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

@@ -19,15 +19,12 @@
 """Utilities for Huggingface Transformers."""
 
 import contextlib
-import hashlib
 import json
 import os
-import tempfile
 from functools import reduce
 from pathlib import Path
 from typing import Any, Optional, cast
 
-import filelock
 from diffusers.loaders.lora_base import (
     _best_guess_weight_name,  # watch out for potetential removal from diffusers
 )
@@ -35,6 +32,7 @@ from huggingface_hub import snapshot_download
 from transformers import AutoConfig, PretrainedConfig
 from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
 
+from sglang.multimodal_gen.runtime.loader.weight_utils import get_lock
 from sglang.multimodal_gen.runtime.utils.logging_utils import (
     init_logger,
     suppress_other_loggers,
@@ -200,18 +198,6 @@ def check_gguf_file(model: str | os.PathLike) -> bool:
     return header == b"GGUF"
 
 
-def get_lock(model_name_or_path: str):
-    lock_dir = tempfile.gettempdir()
-    os.makedirs(os.path.dirname(lock_dir), exist_ok=True)
-    model_name = model_name_or_path.replace("/", "-")
-    hash_name = hashlib.sha256(model_name.encode()).hexdigest()
-    # add hash to avoid conflict with old users' lock files
-    lock_file_name = hash_name + model_name + ".lock"
-    # mode 0o666 is required for the filelock to be shared across users
-    lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666)
-    return lock
-
-
 def maybe_download_lora(
     model_name_or_path: str, local_dir: str | None = None, download: bool = True
 ) -> str:
@@ -391,15 +377,15 @@ def maybe_download_model(
             "Downloading model snapshot from HF Hub for %s...", model_name_or_path
         )
         with (
+            suppress_other_loggers(not_suppress_on_main_rank=False),
             get_lock(model_name_or_path).acquire(poll_interval=2),
-            suppress_other_loggers(not_suppress_on_main_rank=True),
         ):
             local_path = snapshot_download(
                 repo_id=model_name_or_path,
                 ignore_patterns=["*.onnx", "*.msgpack"],
                 local_dir=local_dir,
             )
-        logger.info("Downloaded model to %s", local_path)
+            logger.info("Downloaded model to %s", local_path)
         return str(local_path)
     except Exception as e:
         raise ValueError(

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

@@ -4,6 +4,7 @@
 # adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/logger.py
 """Logging configuration for sglang.multimodal_gen."""
 import argparse
+import contextlib
 import datetime
 import logging
 import os
@@ -382,13 +383,13 @@ def configure_logger(server_args, prefix: str = ""):
     set_uvicorn_logging_configs()
 
 
-def suppress_loggers(loggers_to_suppress: list[str]):
+def suppress_loggers(loggers_to_suppress: list[str], level: int = logging.WARNING):
     original_levels = {}
 
     for logger_name in loggers_to_suppress:
         logger = logging.getLogger(logger_name)
         original_levels[logger_name] = logger.level
-        logger.setLevel(logging.WARNING)
+        logger.setLevel(level)
 
     return original_levels
 
@@ -409,11 +410,21 @@ def suppress_other_loggers(not_suppress_on_main_rank: bool = False):
 
     should_suppress = True
     if not_suppress_on_main_rank:
-        if get_is_main_process() == 0:
+        if get_is_main_process():
             should_suppress = False
 
-    loggers_to_suppress = ["urllib3", "imageio", "imageio_ffmpeg", "PIL", "PIL_Image"]
+    loggers_to_suppress = [
+        "urllib3",
+        "imageio",
+        "imageio_ffmpeg",
+        "PIL",
+        "PIL_Image",
+    ]
+    filelock_loggers = [
+        "filelock",
+    ]
     original_levels = suppress_loggers(loggers_to_suppress)
+    original_levels.update(suppress_loggers(filelock_loggers, level=logging.ERROR))
 
     try:
         yield
@@ -423,6 +434,35 @@ def suppress_other_loggers(not_suppress_on_main_rank: bool = False):
                 logging.getLogger(logger_name).setLevel(level)
 
 
+# source: https://github.com/vllm-project/vllm/blob/a11f4a81e027efd9ef783b943489c222950ac989/vllm/utils/system_utils.py#L60
+@contextlib.contextmanager
+def suppress_stdout():
+    """
+    Suppress stdout from C libraries at the file descriptor level.
+
+    Only suppresses stdout, not stderr, to preserve error messages.
+    Example:
+        with suppress_stdout():
+            # C library calls that would normally print to stdout
+            torch.distributed.new_group(ranks, backend="gloo")
+    """
+    # Don't suppress if logging level is DEBUG
+
+    stdout_fd = sys.stdout.fileno()
+    stdout_dup = os.dup(stdout_fd)
+    devnull_fd = os.open(os.devnull, os.O_WRONLY)
+
+    try:
+        sys.stdout.flush()
+        os.dup2(devnull_fd, stdout_fd)
+        yield
+    finally:
+        sys.stdout.flush()
+        os.dup2(stdout_dup, stdout_fd)
+        os.close(stdout_dup)
+        os.close(devnull_fd)
+
+
 class GenerationTimer:
     def __init__(self):
         self.start_time = 0.0