diff --git a/python/sglang/multimodal_gen/runtime/models/vision_utils.py b/python/sglang/multimodal_gen/runtime/models/vision_utils.py index ac22579a0..c50eefc26 100644 --- a/python/sglang/multimodal_gen/runtime/models/vision_utils.py +++ b/python/sglang/multimodal_gen/runtime/models/vision_utils.py @@ -15,6 +15,8 @@ import requests import torch from packaging import version +from sglang.multimodal_gen.runtime.utils.logging_utils import suppress_other_loggers + if version.parse(version.parse(PIL.__version__).base_version) >= version.parse("9.1.0"): PIL_INTERPOLATION = { "linear": PIL.Image.Resampling.BILINEAR, @@ -108,7 +110,8 @@ def load_image( """ if isinstance(image, str): if image.startswith("http://") or image.startswith("https://"): - image = PIL.Image.open(requests.get(image, stream=True).raw) + with suppress_other_loggers(not_suppress_on_main_rank=True): + image = PIL.Image.open(requests.get(image, stream=True).raw) elif os.path.isfile(image): image = PIL.Image.open(image) else: diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising_dmd.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising_dmd.py index 890428dcb..47516fa70 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising_dmd.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising_dmd.py @@ -9,7 +9,6 @@ from sglang.multimodal_gen.runtime.distributed import ( get_local_torch_device, get_sp_parallel_rank, get_sp_world_size, - logger, sequence_model_parallel_all_gather, ) from sglang.multimodal_gen.runtime.layers.attention.backends.sliding_tile_attn import ( @@ -268,7 +267,7 @@ class DmdDenoisingStage(DenoisingStage): denoising_loop_end_time = time.time() if len(timesteps) > 0: - logger.info( + self.log_info( "Average time per step: %.4f seconds", (denoising_loop_end_time - denoising_loop_start_time) / len(timesteps), ) diff --git a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py index d553069aa..56877d159 100644 --- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py @@ -34,7 +34,10 @@ 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.utils.logging_utils import init_logger +from sglang.multimodal_gen.runtime.utils.logging_utils import ( + init_logger, + suppress_other_loggers, +) logger = init_logger(__name__) _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = { @@ -370,7 +373,9 @@ def maybe_download_model( logger.info( "Downloading model snapshot from HF Hub for %s...", model_name_or_path ) - with get_lock(model_name_or_path).acquire(poll_interval=2): + with 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"], diff --git a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py index be99b3ac2..669bb4b46 100644 --- a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py @@ -5,15 +5,13 @@ """Logging configuration for sglang.multimodal_gen.""" import argparse import datetime -import json import logging import os import sys import warnings +from contextlib import contextmanager from functools import lru_cache, partial from logging import Logger -from logging.config import dictConfig -from os import path from types import MethodType from typing import Any, cast @@ -29,9 +27,6 @@ GREEN = "\033[92m" YELLOW = "\033[93m" RESET = "\033[0;0m" -_warned_local_main_process = False -_warned_main_process = False - _FORMAT = ( f"{SGL_DIFFUSION_LOGGING_PREFIX}%(levelname)s %(asctime)s " "[%(filename)s: %(lineno)d] %(message)s" @@ -129,69 +124,42 @@ def _print_warning_once(logger: Logger, msg: str) -> None: logger.warning(msg, stacklevel=2) -# TODO(will): add env variable to control this process-aware logging behavior -def _info( - logger: Logger, +def _get_rank_info(): + """Get rank and local rank from environment variables.""" + try: + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + except (KeyError, ValueError): + rank = 0 + local_rank = 0 + return rank, local_rank + + +def _log_process_aware( + level: int, + logger_self: Logger, msg: object, *args: Any, - main_process_only: bool = True, - local_main_process_only: bool = True, + main_process_only: bool, + local_main_process_only: bool, **kwargs: Any, ) -> None: - """Process-aware INFO level logging function. - - This function controls logging behavior based on the process rank, allowing for - selective logging from specific processes in a distributed environment. - - Args: - logger: The logger instance to use for logging - msg: The message format string to log - *args: Format string arguments - main_process_only: If True, only log if this is the global main process (RANK=0) - local_main_process_only: If True, only log if this is the local main process (LOCAL_RANK=0) - **kwargs: Additional keyword arguments to pass to the logger.log method - - stacklevel: Defaults to 2 to show the original caller's location - - Note: - - When both main_process_only and local_main_process_only are True, - the message will be logged only if both conditions are met - - When both are False, the message will be logged from all processes - - By default, only logs from processes with LOCAL_RANK=0 - """ - try: - local_rank = int(os.environ["LOCAL_RANK"]) - rank = int(os.environ["RANK"]) - except Exception: - local_rank = 0 - rank = 0 - + """Helper function to log a message if the process rank matches the criteria.""" + rank, local_rank = _get_rank_info() is_main_process = rank == 0 is_local_main_process = local_rank == 0 - if (main_process_only and is_main_process) or ( - local_main_process_only and is_local_main_process - ): - logger.log(logging.INFO, msg, *args, stacklevel=2, **kwargs) + 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) + ) - global _warned_local_main_process, _warned_main_process - - if not _warned_local_main_process and local_main_process_only: - # logger.warning( - # "%sBy default, logger.info(..) will only log from the local main process. Set logger.info(..., is_local_main_process=False) to log from all processes.%s", - # GREEN, - # RESET, - # ) - _warned_local_main_process = True - if not _warned_main_process and main_process_only and is_main_process: - # logger.warning( - # "%sis_main_process_only is set to True, logging only from the main (RANK==0) process.%s", - # GREEN, - # RESET, - # ) - _warned_main_process = True - - if not main_process_only and not local_main_process_only: - logger.log(logging.INFO, msg, *args, stacklevel=2, **kwargs) + if should_log: + # stacklevel=3 to show the original caller's location, + # as this function is called by the patched methods. + logger_self.log(level, msg, *args, stacklevel=3, **kwargs) class _SGLDiffusionLogger(Logger): @@ -224,54 +192,34 @@ class _SGLDiffusionLogger(Logger): main_process_only: bool = True, local_main_process_only: bool = True, **kwargs: Any, - ) -> None: - _info( - self, - msg, - *args, - main_process_only=main_process_only, - local_main_process_only=local_main_process_only, - **kwargs, - ) + ) -> None: ... + def debug( # type: ignore[override] + self, + msg: object, + *args: Any, + main_process_only: bool = True, + local_main_process_only: bool = True, + **kwargs: Any, + ) -> None: ... -def _configure_sgl_diffusion_root_logger() -> None: - logging_config = dict[str, Any]() + def warning( # type: ignore[override] + self, + msg: object, + *args: Any, + main_process_only: bool = False, + local_main_process_only: bool = True, + **kwargs: Any, + ) -> None: ... - if not SGL_DIFFUSION_CONFIGURE_LOGGING and SGL_DIFFUSION_LOGGING_CONFIG_PATH: - raise RuntimeError( - "SGL_DIFFUSION_CONFIGURE_LOGGING evaluated to false, but " - "SGL_DIFFUSION_LOGGING_CONFIG_PATH was given. SGL_DIFFUSION_LOGGING_CONFIG_PATH " - "implies SGL_DIFFUSION_CONFIGURE_LOGGING. Please enable " - "SGL_DIFFUSION_CONFIGURE_LOGGING or unset SGL_DIFFUSION_LOGGING_CONFIG_PATH." - ) - - if SGL_DIFFUSION_CONFIGURE_LOGGING: - logging_config = DEFAULT_LOGGING_CONFIG - - if SGL_DIFFUSION_LOGGING_CONFIG_PATH: - if not path.exists(SGL_DIFFUSION_LOGGING_CONFIG_PATH): - raise RuntimeError( - "Could not load logging config. File does not exist: %s", - SGL_DIFFUSION_LOGGING_CONFIG_PATH, - ) - with open(SGL_DIFFUSION_LOGGING_CONFIG_PATH, encoding="utf-8") as file: - custom_config = json.loads(file.read()) - - if not isinstance(custom_config, dict): - raise ValueError( - "Invalid logging config. Expected Dict, got %s.", - type(custom_config).__name__, - ) - logging_config = custom_config - - for formatter in logging_config.get("formatters", {}).values(): - # This provides backwards compatibility after #10134. - if formatter.get("class") == "sglang.multimodal_gen.logging.NewLineFormatter": - formatter["class"] = "sglang.multimodal_gen.logging_utils.NewLineFormatter" - - if logging_config: - dictConfig(logging_config) + def error( # type: ignore[override] + self, + msg: object, + *args: Any, + main_process_only: bool = False, + local_main_process_only: bool = True, + **kwargs: Any, + ) -> None: ... def init_logger(name: str) -> _SGLDiffusionLogger: @@ -281,25 +229,59 @@ def init_logger(name: str) -> _SGLDiffusionLogger: logger = logging.getLogger(name) - methods_to_patch = { - "info_once": _print_info_once, - "warning_once": _print_warning_once, - "info": _info, - } + # Patch instance methods + setattr(logger, "info_once", MethodType(_print_info_once, logger)) + setattr(logger, "warning_once", MethodType(_print_warning_once, logger)) - for method_name, method in methods_to_patch.items(): - setattr( - logger, method_name, MethodType(method, logger) - ) # type: ignore[arg-type] + def _create_patched_method( + level: int, + main_process_only_default: bool, + local_main_process_only_default: bool, + ): + def _method( + self: Logger, + msg: object, + *args: Any, + main_process_only: bool = main_process_only_default, + local_main_process_only: bool = local_main_process_only_default, + **kwargs: Any, + ) -> None: + _log_process_aware( + level, + self, + msg, + *args, + main_process_only=main_process_only, + local_main_process_only=local_main_process_only, + **kwargs, + ) + + return _method + + setattr( + logger, + "info", + MethodType(_create_patched_method(logging.INFO, True, True), logger), + ) + setattr( + logger, + "debug", + MethodType(_create_patched_method(logging.DEBUG, True, True), logger), + ) + setattr( + logger, + "warning", + MethodType(_create_patched_method(logging.WARNING, False, True), logger), + ) + setattr( + logger, + "error", + MethodType(_create_patched_method(logging.ERROR, False, True), logger), + ) return cast(_SGLDiffusionLogger, logger) -# The root logger is initialized when the module is imported. -# This is thread-safe as the module is only imported once, -# guaranteed by the Python GIL. -# _configure_sgl_diffusion_root_logger() - logger = init_logger(__name__) @@ -395,7 +377,38 @@ def configure_logger(server_args, prefix: str = ""): set_uvicorn_logging_configs() -def suppress_other_loggers(): +@contextmanager +def suppress_other_loggers(not_suppress_on_main_rank: bool = False): + """ + A context manager to temporarily suppress specified loggers. + + Args: + not_suppress_on_main_rank (bool): If True, loggers will not be + suppressed on the main process (rank 0). + """ + # This is a global setting that we want to apply to all ranks warnings.filterwarnings( "ignore", category=UserWarning, message="The given NumPy array is not writable" ) + + should_suppress = True + if not_suppress_on_main_rank: + rank, _ = _get_rank_info() + if rank == 0: + should_suppress = False + + loggers_to_suppress = ["urllib3"] + original_levels = {} + + if should_suppress: + for logger_name in loggers_to_suppress: + logger = logging.getLogger(logger_name) + original_levels[logger_name] = logger.level + logger.setLevel(logging.WARNING) + + try: + yield + finally: + if should_suppress: + for logger_name, level in original_levels.items(): + logging.getLogger(logger_name).setLevel(level)