WIP: initial multimodal-gen support (#12484)

Co-authored-by: yhyang201 <yhyang201@gmail.com>
Co-authored-by: yizhang2077 <1109276519@qq.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>
Co-authored-by: JiLi <leege233@gmail.com>
Co-authored-by: CHEN Xi <78632976+RubiaCx@users.noreply.github.com>
Co-authored-by: laixin <xielx@shanghaitech.edu.cn>
Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>
Co-authored-by: jzhang38 <a1286225768@gmail.com>
Co-authored-by: BrianChen1129 <yongqichcd@gmail.com>
Co-authored-by: Kevin Lin <42618777+kevin314@users.noreply.github.com>
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
Co-authored-by: rlsu9 <r3su@ucsd.edu>
Co-authored-by: Jinzhe Pan <48981407+eigensystem@users.noreply.github.com>
Co-authored-by: foreverpiano <pianoqwz@qq.com>
Co-authored-by: RandNMR73 <notomatthew31@gmail.com>
Co-authored-by: PorridgeSwim <yz3883@columbia.edu>
Co-authored-by: Jiali Chen <90408393+gary-chenjl@users.noreply.github.com>

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

@@ -0,0 +1,291 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+import importlib
+import ipaddress
+import os
+import platform
+import signal
+import socket
+import sys
+import threading
+from functools import lru_cache
+
+import psutil
+import torch
+import zmq
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = None):
+    """Kill the process and all its child processes."""
+    # Remove sigchld handler to avoid spammy logs.
+    if threading.current_thread() is threading.main_thread():
+        signal.signal(signal.SIGCHLD, signal.SIG_DFL)
+
+    if parent_pid is None:
+        parent_pid = os.getpid()
+        include_parent = False
+
+    try:
+        itself = psutil.Process(parent_pid)
+    except psutil.NoSuchProcess:
+        return
+
+    children = itself.children(recursive=True)
+    for child in children:
+        if child.pid == skip_pid:
+            continue
+        try:
+            child.kill()
+        except psutil.NoSuchProcess:
+            pass
+
+    if include_parent:
+        try:
+            if parent_pid == os.getpid():
+                itself.kill()
+                sys.exit(0)
+
+            itself.kill()
+
+            # Sometime processes cannot be killed with SIGKILL (e.g, PID=1 launched by kubernetes),
+            # so we send an additional signal to kill them.
+            itself.send_signal(signal.SIGQUIT)
+        except psutil.NoSuchProcess:
+            pass
+
+
+def add_prefix(name: str, prefix: str) -> str:
+    """Add a weight path prefix to a module name.
+
+    Args:
+        name: base module name.
+        prefix: weight prefix str to added to the front of `name` concatenated with `.`.
+
+    Returns:
+        The string `prefix.name` if prefix is non-empty, otherwise just `name`.
+    """
+    return name if not prefix else f"{prefix}.{name}"
+
+
+def is_valid_ipv6_address(address: str) -> bool:
+    try:
+        ipaddress.IPv6Address(address)
+        return True
+    except ValueError:
+        return False
+
+
+def configure_ipv6(dist_init_addr):
+    addr = dist_init_addr
+    end = addr.find("]")
+    if end == -1:
+        raise ValueError("invalid IPv6 address format: missing ']'")
+
+    host = addr[: end + 1]
+
+    # this only validates the address without brackets: we still need the below checks.
+    # if it's invalid, immediately raise an error so we know it's not formatting issues.
+    if not is_valid_ipv6_address(host[1:end]):
+        raise ValueError(f"invalid IPv6 address: {host}")
+
+    port_str = None
+    if len(addr) > end + 1:
+        if addr[end + 1] == ":":
+            port_str = addr[end + 2 :]
+        else:
+            raise ValueError("received IPv6 address format: expected ':' after ']'")
+
+    if not port_str:
+        raise ValueError(
+            "a port must be specified in IPv6 address (format: [ipv6]:port)"
+        )
+
+    try:
+        port = int(port_str)
+    except ValueError:
+        raise ValueError(f"invalid port in IPv6 address: '{port_str}'")
+    return port, host
+
+
+def is_port_available(port):
+    """Return whether a port is available."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        try:
+            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            s.bind(("", port))
+            s.listen(1)
+            return True
+        except socket.error:
+            return False
+        except OverflowError:
+            return False
+
+
+def get_zmq_socket(
+    context: zmq.Context, socket_type: zmq.SocketType, endpoint: str, bind: bool
+) -> zmq.Socket:
+    mem = psutil.virtual_memory()
+    total_mem = mem.total / 1024**3
+    available_mem = mem.available / 1024**3
+    if total_mem > 32 and available_mem > 16:
+        buf_size = int(0.5 * 1024**3)
+    else:
+        buf_size = -1
+
+    socket = context.socket(socket_type)
+    if endpoint.find("[") != -1:
+        socket.setsockopt(zmq.IPV6, 1)
+
+    def set_send_opt():
+        socket.setsockopt(zmq.SNDHWM, 0)
+        socket.setsockopt(zmq.SNDBUF, buf_size)
+
+    def set_recv_opt():
+        socket.setsockopt(zmq.RCVHWM, 0)
+        socket.setsockopt(zmq.RCVBUF, buf_size)
+
+    if socket_type == zmq.PUSH:
+        set_send_opt()
+    elif socket_type == zmq.PULL:
+        set_recv_opt()
+    elif socket_type == zmq.DEALER:
+        set_send_opt()
+        set_recv_opt()
+    elif socket_type == zmq.REQ:
+        set_send_opt()
+        set_recv_opt()
+    elif socket_type == zmq.REP:
+        set_send_opt()
+        set_recv_opt()
+    else:
+        raise ValueError(f"Unsupported socket type: {socket_type}")
+
+    if bind:
+        socket.bind(endpoint)
+    else:
+        socket.connect(endpoint)
+
+    return socket
+
+
+# https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip
+@lru_cache(maxsize=1)
+def is_hip() -> bool:
+    return torch.version.hip is not None
+
+
+@lru_cache(maxsize=1)
+def is_cuda():
+    return torch.cuda.is_available() and torch.version.cuda
+
+
+@lru_cache(maxsize=1)
+def is_cuda_alike():
+    return is_cuda() or is_hip()
+
+
+@lru_cache(maxsize=1)
+def is_blackwell():
+    if not is_cuda():
+        return False
+    return torch.cuda.get_device_capability()[0] == 10
+
+
+@lru_cache(maxsize=1)
+def is_hpu() -> bool:
+    return hasattr(torch, "hpu") and torch.hpu.is_available()
+
+
+@lru_cache(maxsize=1)
+def is_xpu() -> bool:
+    return hasattr(torch, "xpu") and torch.xpu.is_available()
+
+
+@lru_cache(maxsize=1)
+def is_npu() -> bool:
+    return hasattr(torch, "npu") and torch.npu.is_available()
+
+
+@lru_cache(maxsize=1)
+def is_host_cpu_x86() -> bool:
+    machine = platform.machine().lower()
+    return (
+        machine in ("x86_64", "amd64", "i386", "i686")
+        and hasattr(torch, "cpu")
+        and torch.cpu.is_available()
+    )
+
+
+@lru_cache(maxsize=1)
+def is_cpu() -> bool:
+    return os.getenv("SGLANG_USE_CPU_ENGINE", "0") == "1" and is_host_cpu_x86()
+
+
+# cuda
+
+
+def set_cuda_arch():
+    capability = torch.cuda.get_device_capability()
+    arch = f"{capability[0]}.{capability[1]}"
+    os.environ["TORCH_CUDA_ARCH_LIST"] = f"{arch}{'+PTX' if arch == '9.0' else ''}"
+
+
+def get_bool_env_var(env_var_name: str, default: str | bool = "false") -> bool:
+    raw_value = os.getenv(env_var_name, None)
+    if raw_value is None:
+        raw_value = str(default)
+
+    value_str = str(raw_value).strip().lower()
+    truthy = {"1", "true", "yes", "y", "t", "on"}
+    falsy = {"0", "false", "no", "n", "f", "off", ""}
+
+    if value_str in truthy:
+        return True
+    if value_str in falsy:
+        return False
+
+    default_bool = str(default).strip().lower() in truthy
+    logger.warning(
+        "Unrecognized boolean for %s=%r; falling back to default=%r",
+        env_var_name,
+        raw_value,
+        default_bool,
+    )
+    return default_bool
+
+
+def is_flashinfer_available():
+    """
+    Check whether flashinfer is available.
+    As of Oct. 6, 2024, it is only available on NVIDIA GPUs.
+    """
+    # if not get_bool_env_var("SGLANG_IS_FLASHINFER_AVAILABLE", default="true"):
+    #     return False
+    return importlib.util.find_spec("flashinfer") is not None and is_cuda()
+
+
+# env var managements
+
+_warned_bool_env_var_keys = set()
+
+
+def get_bool_env_var(name: str, default: str = "false") -> bool:
+    # FIXME: move your environment variable to sglang.srt.environ
+    value = os.getenv(name, default)
+    value = value.lower()
+
+    truthy_values = ("true", "1")
+    falsy_values = ("false", "0")
+
+    if (value not in truthy_values) and (value not in falsy_values):
+        if value not in _warned_bool_env_var_keys:
+            logger.warning(
+                f"get_bool_env_var({name}) see non-understandable value={value} and treat as false"
+            )
+        _warned_bool_env_var_keys.add(value)
+
+    return value in truthy_values

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

@@ -0,0 +1,231 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+import pickle
+from typing import Any, List, Optional
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+
+def broadcast_pyobj(
+    data: List[Any],
+    rank: int,
+    dist_group: Optional[torch.distributed.ProcessGroup] = None,
+    src: int = 0,
+    force_cpu_device: bool = True,
+):
+    """Broadcast inputs from src rank to all other ranks with torch.dist backend.
+    The `rank` here refer to the source rank on global process group (regardless
+    of dist_group argument).
+    """
+    device = torch.device(
+        "cuda" if torch.cuda.is_available() and not force_cpu_device else "cpu"
+    )
+
+    if rank == src:
+        if data is None or len(data) == 0:
+            tensor_size = torch.tensor([0], dtype=torch.long, device=device)
+            dist.broadcast(tensor_size, src=src, group=dist_group)
+        else:
+            serialized_data = pickle.dumps(data)
+            size = len(serialized_data)
+
+            tensor_data = torch.ByteTensor(
+                np.frombuffer(serialized_data, dtype=np.uint8)
+            ).to(device)
+            tensor_size = torch.tensor([size], dtype=torch.long, device=device)
+
+            dist.broadcast(tensor_size, src=src, group=dist_group)
+            dist.broadcast(tensor_data, src=src, group=dist_group)
+        return data
+    else:
+        tensor_size = torch.tensor([0], dtype=torch.long, device=device)
+        dist.broadcast(tensor_size, src=src, group=dist_group)
+        size = tensor_size.item()
+
+        if size == 0:
+            return []
+
+        tensor_data = torch.empty(size, dtype=torch.uint8, device=device)
+        dist.broadcast(tensor_data, src=src, group=dist_group)
+
+        serialized_data = bytes(tensor_data.cpu().numpy())
+        data = pickle.loads(serialized_data)
+        return data
+
+
+def generate_masked_orthogonal_rank_groups(
+    world_size: int, parallel_size: list[int], mask: list[bool]
+) -> list[list[int]]:
+    """Generate orthogonal parallel groups based on the parallel size and mask.
+
+    Arguments:
+        world_size (int): world size
+
+        parallel_size (List[int]):
+            The parallel size of each orthogonal parallel type. For example, if
+            tensor_parallel_size = 2, pipeline_model_parallel_group = 3, data_parallel_size = 4,
+            and the parallel mapping order is tp-pp-dp, then the parallel_size = [2, 3, 4].
+
+        mask (List[bool]):
+            The mask controls which parallel methods the generated groups represent. If mask[i] is
+            True, it means the generated group contains the i-th parallelism method. For example,
+            if parallel_size = [tp_size, pp_size, dp_size], and mask = [True, False , True], then
+            the generated group is the `tp-dp` group, if the mask = [False, True, False], then the
+            generated group is the `pp` group.
+
+    Algorithm:
+        For orthogonal parallelism, such as tp/dp/pp/cp, the global_rank and
+
+        If we want to get the `dp_group` (tp_size * pp_size groups of dp_size ranks each.
+        For example,  if the gpu size is 8 and order is 'tp-pp-dp', size is '2-2-2', and the
+        dp_group here is [[0, 4], [1, 5], [2, 6], [3, 7]].)
+        The tp_rank and pp_rank will be combined to form the `dp_group_index`.
+            dp_group_index = tp_rank + pp_rank * tp_size (2)
+
+        So, Given that tp_rank and pp_rank satisfy equation (2), and dp_rank in
+        range(0, dp_size), the ranks in dp_group[dp_group_index] satisfies the
+        equation (1).
+
+        This function solve this math problem.
+
+    For example, if the parallel_size = [tp_size, dp_size, pp_size] = [2, 3, 4],
+    and the mask = [False, True, False]. Then,
+        dp_group_index(0) = tp_rank(0) + pp_rank(0) * 2
+        dp_group_index(1) = tp_rank(1) + pp_rank(0) * 2
+        ...
+        dp_group_index(7) = tp_rank(1) + pp_rank(3) * 2
+
+        dp_group[0] = 0 + range(0, 3) * 2 + 0 = [0, 2, 4]
+        dp_group[1] = 1 + range(0, 3) * 2 + 0 = [1, 3, 5]
+        ...
+        dp_group[7] = 1 + range(0, 3) * 2 + 3 * 2 * 3 = [19, 21, 23]
+    """
+
+    def prefix_product(a: List[int], init=1) -> List[int]:
+        r = [init]
+        for v in a:
+            init = init * v
+            r.append(init)
+        return r
+
+    def inner_product(a: List[int], b: List[int]) -> int:
+        return sum([x * y for x, y in zip(a, b)])
+
+    def decompose(index, shape, stride=None):
+        """
+        This function solve the math problem below:
+            There is an equation:
+                index = sum(idx[i] * stride[i])
+            And given the value of index, stride.
+            Return the idx.
+        This function will used to get the pp/dp/pp_rank
+        from group_index and rank_in_group.
+        """
+        if stride is None:
+            stride = prefix_product(shape)
+        idx = [(index // d) % s for s, d in zip(shape, stride)]
+        # stride is a prefix_product result. And the value of stride[-1]
+        # is not used.
+        assert (
+            sum([x * y for x, y in zip(idx, stride[:-1])]) == index
+        ), "idx {} with shape {} mismatch the return idx {}".format(index, shape, idx)
+        return idx
+
+    masked_shape = [s for s, m in zip(parallel_size, mask) if m]
+    unmasked_shape = [s for s, m in zip(parallel_size, mask) if not m]
+
+    global_stride = prefix_product(parallel_size)
+    masked_stride = [d for d, m in zip(global_stride, mask) if m]
+    unmasked_stride = [d for d, m in zip(global_stride, mask) if not m]
+
+    group_size = prefix_product(masked_shape)[-1]
+    num_of_group = world_size // group_size
+
+    ranks = []
+    for group_index in range(num_of_group):
+        # get indices from unmaksed for group_index.
+        decomposed_group_idx = decompose(group_index, unmasked_shape)
+        rank = []
+        for rank_in_group in range(group_size):
+            # get indices from masked for rank_in_group.
+            decomposed_rank_idx = decompose(rank_in_group, masked_shape)
+            rank.append(
+                inner_product(decomposed_rank_idx, masked_stride)
+                + inner_product(decomposed_group_idx, unmasked_stride)
+            )
+        ranks.append(rank)
+    return ranks
+
+
+class RankGenerator(object):
+    def __init__(
+        self,
+        tp: int,
+        sp: int,
+        pp: int,
+        cfg: int,
+        dp: int,
+        order: str,
+        rank_offset: int = 0,
+    ) -> None:
+        self.tp = tp
+        self.sp = sp
+        self.pp = pp
+        self.cfg = cfg
+        self.dp = dp
+        self.rank_offset = rank_offset
+        self.world_size = tp * sp * pp * cfg * dp
+
+        self.name_to_size = {
+            "tp": self.tp,
+            "sp": self.sp,
+            "pp": self.pp,
+            "cfg": self.cfg,
+            "dp": self.dp,
+        }
+        order = order.lower()
+
+        for name in self.name_to_size.keys():
+            if name not in order and self.name_to_size[name] != 1:
+                raise RuntimeError(
+                    f"The size of ({name}) is ({self.name_to_size[name]}), but you haven't specified the order ({self.order})."
+                )
+            elif name not in order:
+                order = order + "-" + name
+
+        self.order = order
+        self.ordered_size = []
+
+        for token in order.split("-"):
+            self.ordered_size.append(self.name_to_size[token])
+
+    def get_mask(self, order: str, token: str):
+        ordered_token = order.split("-")
+        token = token.split("-")
+        mask = [False] * len(ordered_token)
+        for t in token:
+            mask[ordered_token.index(t)] = True
+        return mask
+
+    def get_ranks(self, token):
+        """Get rank group by input token.
+
+        Arguments:
+            token (str):
+                Specify the ranks type that want to get. If we want
+                to obtain multiple parallel types, we can use a hyphen
+                '-' to separate them. For example, if we want to obtain
+                the TP_DP group, the token should be 'tp-dp'.
+
+        """
+        mask = self.get_mask(self.order, token)
+        ranks = generate_masked_orthogonal_rank_groups(
+            self.world_size, self.ordered_size, mask
+        )
+        if self.rank_offset > 0:
+            for rank_group in ranks:
+                for i in range(len(rank_group)):
+                    rank_group[i] += self.rank_offset
+        return ranks

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

@@ -0,0 +1,384 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/hf_transformers_utils.py
+
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities for Huggingface Transformers."""
+
+import contextlib
+import hashlib
+import json
+import os
+import tempfile
+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
+)
+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
+
+logger = init_logger(__name__)
+_CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = {
+    # ChatGLMConfig.model_type: ChatGLMConfig,
+    # DbrxConfig.model_type: DbrxConfig,
+    # ExaoneConfig.model_type: ExaoneConfig,
+    # Qwen2_5_VLConfig.model_type: Qwen2_5_VLConfig,
+}
+
+for name, cls in _CONFIG_REGISTRY.items():
+    with contextlib.suppress(ValueError):
+        AutoConfig.register(name, cls)
+
+
+def download_from_hf(model_path: str):
+    if os.path.exists(model_path):
+        return model_path
+
+    return snapshot_download(model_path, allow_patterns=["*.json", "*.bin", "*.model"])
+
+
+def get_hf_config(
+    model: str,
+    trust_remote_code: bool,
+    revision: str | None = None,
+    model_override_args: dict | None = None,
+    **kwargs,
+):
+    is_gguf = check_gguf_file(model)
+    if is_gguf:
+        raise NotImplementedError("GGUF models are not supported.")
+
+    config = AutoConfig.from_pretrained(
+        model, trust_remote_code=trust_remote_code, revision=revision, **kwargs
+    )
+    if config.model_type in _CONFIG_REGISTRY:
+        config_class = _CONFIG_REGISTRY[config.model_type]
+        config = config_class.from_pretrained(model, revision=revision)
+        # NOTE(HandH1998): Qwen2VL requires `_name_or_path` attribute in `config`.
+        config._name_or_path = model
+    if model_override_args:
+        config.update(model_override_args)
+
+    # Special architecture mapping check for GGUF models
+    if is_gguf:
+        if config.model_type not in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
+            raise RuntimeError(f"Can't get gguf config for {config.model_type}.")
+        model_type = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[config.model_type]
+        config.update({"architectures": [model_type]})
+
+    return config
+
+
+def get_config(
+    model: str,
+    trust_remote_code: bool,
+    revision: Optional[str] = None,
+    model_override_args: Optional[dict] = None,
+    **kwargs,
+):
+    try:
+        config = AutoConfig.from_pretrained(
+            model, trust_remote_code=trust_remote_code, revision=revision, **kwargs
+        )
+    except ValueError as e:
+        raise e
+
+    return config
+
+
+def load_dict(file_path):
+    if not os.path.exists(file_path):
+        return {}
+    try:
+        # Load the config directly from the file
+        with open(file_path) as f:
+            config_dict: dict[str, Any] = json.load(f)
+        if "_diffusers_version" in config_dict:
+            config_dict.pop("_diffusers_version")
+        # TODO(will): apply any overrides from inference args
+        return config_dict
+    except Exception as e:
+        raise RuntimeError(
+            f"Failed to load diffusers config from {file_path}: {e}"
+        ) from e
+
+
+def get_diffusers_config(
+    model: str,
+) -> dict[str, Any]:
+    """Gets a configuration for the given diffusers model.
+
+    Args:
+        model: The model name or path.
+
+    Returns:
+        The loaded configuration.
+    """
+
+    config_name = "config.json"
+    if "scheduler" in model:
+        config_name = "scheduler_config.json"
+    # Check if the model path exists
+    if os.path.exists(model):
+        config_file = os.path.join(model, config_name)
+        config_dict = load_dict(config_file)
+        generation_config_file = os.path.join(model, "generation_config.json")
+        generation_config_dict = load_dict(generation_config_file)
+        return config_dict | generation_config_dict
+    else:
+        raise RuntimeError(f"Diffusers config file not found at {model}")
+
+
+# Models don't use the same configuration key for determining the maximum
+# context length.  Store them here so we can sanely check them.
+# NOTE: The ordering here is important. Some models have two of these and we
+# have a preference for which value gets used.
+CONTEXT_LENGTH_KEYS = [
+    "max_sequence_length",
+    "seq_length",
+    "max_seq_len",
+    "model_max_length",
+    "max_position_embeddings",
+]
+
+
+def attach_additional_stop_token_ids(tokenizer):
+    # Special handling for stop token <|eom_id|> generated by llama 3 tool use.
+    if "<|eom_id|>" in tokenizer.get_added_vocab():
+        tokenizer.additional_stop_token_ids = set(
+            [tokenizer.get_added_vocab()["<|eom_id|>"]]
+        )
+    else:
+        tokenizer.additional_stop_token_ids = None
+
+
+def check_gguf_file(model: str | os.PathLike) -> bool:
+    """Check if the file is a GGUF model."""
+    model = Path(model)
+    if not model.is_file():
+        return False
+    elif model.suffix == ".gguf":
+        return True
+
+    with open(model, "rb") as f:
+        header = f.read(4)
+    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:
+    """
+    Check if the model path is a Hugging Face Hub model ID and download it if needed.
+    Args:
+        model_name_or_path: Local path or Hugging Face Hub model ID
+        local_dir: Local directory to save the model
+        download: Whether to download the model from Hugging Face Hub
+
+    Returns:
+        Local path to the model
+    """
+
+    local_path = maybe_download_model(model_name_or_path, local_dir, download)
+    weight_name = _best_guess_weight_name(
+        model_name_or_path, file_extension=".safetensors"
+    )
+    return os.path.join(local_path, weight_name)
+
+
+def verify_model_config_and_directory(model_path: str) -> dict[str, Any]:
+    """
+    Verify that the model directory contains a valid diffusers configuration.
+
+    Args:
+        model_path: Path to the model directory
+
+    Returns:
+        The loaded model configuration as a dictionary
+    """
+
+    # Check for model_index.json which is required for diffusers models
+    config_path = os.path.join(model_path, "model_index.json")
+    if not os.path.exists(config_path):
+        raise ValueError(
+            f"Model directory {model_path} does not contain model_index.json. "
+            "Only HuggingFace diffusers format is supported."
+        )
+
+    # Check for transformer and vae directories
+    transformer_dir = os.path.join(model_path, "transformer")
+    vae_dir = os.path.join(model_path, "vae")
+
+    if not os.path.exists(transformer_dir):
+        raise ValueError(
+            f"Model directory {model_path} does not contain a transformer/ directory."
+        )
+
+    if not os.path.exists(vae_dir):
+        raise ValueError(
+            f"Model directory {model_path} does not contain a vae/ directory."
+        )
+
+    # Load the config
+    with open(config_path) as f:
+        config = json.load(f)
+
+    # Verify diffusers version exists
+    if "_diffusers_version" not in config:
+        raise ValueError("model_index.json does not contain _diffusers_version")
+
+    logger.info("Diffusers version: %s", config["_diffusers_version"])
+    return cast(dict[str, Any], config)
+
+
+def maybe_download_model_index(model_name_or_path: str) -> dict[str, Any]:
+    """
+    Download and extract just the model_index.json for a Hugging Face model.
+
+    Args:
+        model_name_or_path: Path or HF Hub model ID
+
+    Returns:
+        The parsed model_index.json as a dictionary
+    """
+    import tempfile
+
+    from huggingface_hub import hf_hub_download
+    from huggingface_hub.errors import EntryNotFoundError
+
+    # If it's a local path, verify it directly
+    if os.path.exists(model_name_or_path):
+        try:
+            return verify_model_config_and_directory(model_name_or_path)
+        except ValueError:
+            # Not a pipeline, maybe a single model.
+            config_path = os.path.join(model_name_or_path, "config.json")
+            if os.path.exists(config_path):
+                with open(config_path) as f:
+                    config = json.load(f)
+                return config
+            raise
+
+    # For remote models, download just the model_index.json
+    try:
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            # Download just the model_index.json file
+            model_index_path = hf_hub_download(
+                repo_id=model_name_or_path,
+                filename="model_index.json",
+                local_dir=tmp_dir,
+            )
+
+            # Load the model_index.json
+            with open(model_index_path) as f:
+                config: dict[str, Any] = json.load(f)
+
+            # Verify it has the required fields
+            if "_class_name" not in config:
+                raise ValueError(
+                    f"model_index.json for {model_name_or_path} does not contain _class_name field"
+                )
+
+            if "_diffusers_version" not in config:
+                raise ValueError(
+                    f"model_index.json for {model_name_or_path} does not contain _diffusers_version field"
+                )
+
+            # Add the pipeline name for downstream use
+            config["pipeline_name"] = config["_class_name"]
+
+            logger.info(
+                "Downloaded model_index.json for %s, pipeline: %s",
+                model_name_or_path,
+                config["_class_name"],
+            )
+            return config
+    except EntryNotFoundError:
+        logger.warning(
+            "model_index.json not found for %s. Assuming it is a single model and downloading it.",
+            model_name_or_path,
+        )
+        local_path = maybe_download_model(model_name_or_path)
+        config_path = os.path.join(local_path, "config.json")
+        if not os.path.exists(config_path):
+            raise ValueError(
+                f"Failed to find config.json for {model_name_or_path} after failing to find model_index.json"
+                f"You might be looking for models ending with '-Diffusers'"
+            )
+        with open(config_path) as f:
+            config = json.load(f)
+        return config
+    except Exception as e:
+        raise ValueError(
+            f"Failed to download or parse model_index.json for {model_name_or_path}: {e}"
+        ) from e
+
+
+def maybe_download_model(
+    model_name_or_path: str, local_dir: str | None = None, download: bool = True
+) -> str:
+    """
+    Check if the model path is a Hugging Face Hub model ID and download it if needed.
+
+    Args:
+        model_name_or_path: Local path or Hugging Face Hub model ID
+        local_dir: Local directory to save the model
+        download: Whether to download the model from Hugging Face Hub
+
+    Returns:
+        Local path to the model
+    """
+
+    # If the path exists locally, return it
+    if os.path.exists(model_name_or_path):
+        logger.info("Model already exists locally")
+        return model_name_or_path
+
+    # Otherwise, assume it's a HF Hub model ID and try to download it
+    try:
+        logger.info(
+            "Downloading model snapshot from HF Hub for %s...", model_name_or_path
+        )
+        with get_lock(model_name_or_path):
+            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)
+        return str(local_path)
+    except Exception as e:
+        raise ValueError(
+            f"Could not find model at {model_name_or_path} and failed to download from HF Hub: {e}"
+        ) from e

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

@@ -0,0 +1,401 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# 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 datetime
+import json
+import logging
+import os
+import sys
+import warnings
+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
+
+import sglang.multimodal_gen.envs as envs
+
+SGL_DIFFUSION_CONFIGURE_LOGGING = envs.SGL_DIFFUSION_CONFIGURE_LOGGING
+SGL_DIFFUSION_LOGGING_CONFIG_PATH = envs.SGL_DIFFUSION_LOGGING_CONFIG_PATH
+SGL_DIFFUSION_LOGGING_LEVEL = envs.SGL_DIFFUSION_LOGGING_LEVEL
+SGL_DIFFUSION_LOGGING_PREFIX = envs.SGL_DIFFUSION_LOGGING_PREFIX
+
+RED = "\033[91m"
+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"
+)
+
+# _FORMAT = "[%(asctime)s] %(message)s"
+_DATE_FORMAT = "%m-%d %H:%M:%S"
+
+DEFAULT_LOGGING_CONFIG = {
+    "formatters": {
+        "sgl_diffusion": {
+            "class": "sglang.multimodal_gen.runtime.utils.logging_utils.ColoredFormatter",
+            "datefmt": _DATE_FORMAT,
+            "format": _FORMAT,
+        },
+    },
+    "handlers": {
+        "sgl_diffusion": {
+            "class": "logging.StreamHandler",
+            "formatter": "sgl_diffusion",
+            "level": SGL_DIFFUSION_LOGGING_LEVEL,
+            "stream": "ext://sys.stdout",
+        },
+    },
+    "loggers": {
+        "sgl_diffusion": {
+            "handlers": ["sgl_diffusion"],
+            "level": "WARNING",
+            "propagate": False,
+        },
+    },
+    "root": {
+        "handlers": ["sgl_diffusion"],
+        "level": "DEBUG",
+    },
+    "version": 1,
+    "disable_existing_loggers": False,
+}
+
+
+class NewLineFormatter(logging.Formatter):
+    """Adds logging prefix to newlines to align multi-line messages."""
+
+    def __init__(self, fmt, datefmt=None, style="%"):
+        logging.Formatter.__init__(self, fmt, datefmt, style)
+
+    def format(self, record):
+        msg = logging.Formatter.format(self, record)
+        if record.message != "":
+            parts = msg.split(record.message)
+            msg = msg.replace("\n", "\r\n" + parts[0])
+        return msg
+
+
+class ColoredFormatter(NewLineFormatter):
+    """A logging formatter that adds color to log levels."""
+
+    LEVEL_COLORS = {
+        logging.ERROR: RED,
+        logging.WARNING: YELLOW,
+    }
+
+    def format(self, record: logging.LogRecord) -> str:
+        """Adds color to the log level name."""
+        original_levelname = record.levelname
+        color = self.LEVEL_COLORS.get(record.levelno)
+        if color:
+            record.levelname = f"{color}{original_levelname}{RESET}"
+
+        formatted_message = super().format(record)
+
+        if color:
+            record.levelname = original_levelname
+
+        return formatted_message
+
+
+class SortedHelpFormatter(argparse.HelpFormatter):
+    """SortedHelpFormatter that sorts arguments by their option strings."""
+
+    def add_arguments(self, actions):
+        actions = sorted(actions, key=lambda x: x.option_strings)
+        super().add_arguments(actions)
+
+
+@lru_cache
+def _print_info_once(logger: Logger, msg: str) -> None:
+    # Set the stacklevel to 2 to print the original caller's line info
+    logger.info(msg, stacklevel=2)
+
+
+@lru_cache
+def _print_warning_once(logger: Logger, msg: str) -> None:
+    # Set the stacklevel to 2 to print the original caller's line info
+    logger.warning(msg, stacklevel=2)
+
+
+# TODO(will): add env variable to control this process-aware logging behavior
+def _info(
+    logger: Logger,
+    msg: object,
+    *args: Any,
+    main_process_only: bool = True,
+    local_main_process_only: bool = True,
+    **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
+
+    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)
+
+    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)
+
+
+class _SGLDiffusionLogger(Logger):
+    """
+    Note:
+        This class is just to provide type information.
+        We actually patch the methods directly on the :class:`logging.Logger`
+        instance to avoid conflicting with other libraries such as
+        `intel_extension_for_pytorch.utils._logger`.
+    """
+
+    def info_once(self, msg: str) -> None:
+        """
+        As :meth:`info`, but subsequent calls with the same message
+        are silently dropped.
+        """
+        _print_info_once(self, msg)
+
+    def warning_once(self, msg: str) -> None:
+        """
+        As :meth:`warning`, but subsequent calls with the same message
+        are silently dropped.
+        """
+        _print_warning_once(self, msg)
+
+    def info(  # type: ignore[override]
+        self,
+        msg: object,
+        *args: Any,
+        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,
+        )
+
+
+def _configure_sgl_diffusion_root_logger() -> None:
+    logging_config = dict[str, Any]()
+
+    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 init_logger(name: str) -> _SGLDiffusionLogger:
+    """The main purpose of this function is to ensure that loggers are
+    retrieved in such a way that we can be sure the root sgl_diffusion logger has
+    already been configured."""
+
+    logger = logging.getLogger(name)
+
+    methods_to_patch = {
+        "info_once": _print_info_once,
+        "warning_once": _print_warning_once,
+        "info": _info,
+    }
+
+    for method_name, method in methods_to_patch.items():
+        setattr(
+            logger, method_name, MethodType(method, logger)
+        )  # type: ignore[arg-type]
+
+    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__)
+
+
+def _trace_calls(log_path, root_dir, frame, event, arg=None):
+    if event in ["call", "return"]:
+        # Extract the filename, line number, function name, and the code object
+        filename = frame.f_code.co_filename
+        lineno = frame.f_lineno
+        func_name = frame.f_code.co_name
+        if not filename.startswith(root_dir):
+            # only log the functions in the sgl_diffusion root_dir
+            return
+        # Log every function call or return
+        try:
+            last_frame = frame.f_back
+            if last_frame is not None:
+                last_filename = last_frame.f_code.co_filename
+                last_lineno = last_frame.f_lineno
+                last_func_name = last_frame.f_code.co_name
+            else:
+                # initial frame
+                last_filename = ""
+                last_lineno = 0
+                last_func_name = ""
+            with open(log_path, "a") as f:
+                ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")
+                if event == "call":
+                    f.write(
+                        f"{ts} Call to"
+                        f" {func_name} in {filename}:{lineno}"
+                        f" from {last_func_name} in {last_filename}:"
+                        f"{last_lineno}\n"
+                    )
+                else:
+                    f.write(
+                        f"{ts} Return from"
+                        f" {func_name} in {filename}:{lineno}"
+                        f" to {last_func_name} in {last_filename}:"
+                        f"{last_lineno}\n"
+                    )
+        except NameError:
+            # modules are deleted during shutdown
+            pass
+    return partial(_trace_calls, log_path, root_dir)
+
+
+def enable_trace_function_call(log_file_path: str, root_dir: str | None = None):
+    """
+    Enable tracing of every function call in code under `root_dir`.
+    This is useful for debugging hangs or crashes.
+    `log_file_path` is the path to the log file.
+    `root_dir` is the root directory of the code to trace. If None, it is the
+    sgl_diffusion root directory.
+
+    Note that this call is thread-level, any threads calling this function
+    will have the trace enabled. Other threads will not be affected.
+    """
+    logger.warning(
+        "SGL_DIFFUSION_TRACE_FUNCTION is enabled. It will record every"
+        " function executed by Python. This will slow down the code. It "
+        "is suggested to be used for debugging hang or crashes only."
+    )
+    logger.info("Trace frame log is saved to %s", log_file_path)
+    if root_dir is None:
+        # by default, this is the sgl_diffusion root directory
+        root_dir = os.path.dirname(os.path.dirname(__file__))
+    sys.settrace(partial(_trace_calls, log_file_path, root_dir))
+
+
+def set_uvicorn_logging_configs():
+    from uvicorn.config import LOGGING_CONFIG
+
+    LOGGING_CONFIG["formatters"]["default"][
+        "fmt"
+    ] = "[%(asctime)s] %(levelprefix)s %(message)s"
+    LOGGING_CONFIG["formatters"]["default"]["datefmt"] = "%Y-%m-%d %H:%M:%S"
+    LOGGING_CONFIG["formatters"]["access"][
+        "fmt"
+    ] = '[%(asctime)s] %(levelprefix)s %(client_addr)s - "%(request_line)s" %(status_code)s'
+    LOGGING_CONFIG["formatters"]["access"]["datefmt"] = "%Y-%m-%d %H:%M:%S"
+
+
+def configure_logger(server_args, prefix: str = ""):
+    log_format = f"[%(asctime)s{prefix}] %(message)s"
+    datefmt = "%m-%d %H:%M:%S"
+    logging.basicConfig(
+        level=getattr(logging, server_args.log_level.upper()),
+        format=log_format,
+        datefmt=datefmt,
+        force=True,
+    )
+
+    set_uvicorn_logging_configs()
+
+
+def suppress_other_loggers():
+    warnings.filterwarnings(
+        "ignore", category=UserWarning, message="The given NumPy array is not writable"
+    )

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

@@ -0,0 +1,76 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+import json
+import logging
+import os
+import subprocess
+import time
+from datetime import datetime
+
+from dateutil.tz import UTC
+
+project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))
+LOG_DIR = os.path.join(project_root, "logs")
+
+# Configure a specific logger for performance metrics
+perf_logger = logging.getLogger("performance")
+perf_logger.setLevel(logging.INFO)
+perf_logger.propagate = False  # Prevent perf logs from going to the main logger
+
+# Ensure the logs directory exists
+if not os.path.exists(LOG_DIR):
+    os.makedirs(LOG_DIR)
+
+# Set up a file handler for the performance logger
+handler = logging.FileHandler(os.path.join(LOG_DIR, "performance.log"))
+handler.setFormatter(logging.Formatter("%(message)s"))
+perf_logger.addHandler(handler)
+
+
+def get_git_commit_hash() -> str:
+    """Get the current git commit hash."""
+    try:
+        commit_hash = (
+            subprocess.check_output(["git", "rev-parse", "HEAD"])
+            .strip()
+            .decode("utf-8")
+        )
+        return commit_hash
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        return "N/A"
+
+
+class PerformanceLogger:
+    """
+    A utility class for logging performance metrics.
+    """
+
+    def __init__(self, request_id: str):
+        self.request_id = request_id
+        self.start_time = time.monotonic()
+        self.step_timings = []
+        self.commit_hash = get_git_commit_hash()
+
+    def record_step_start(self):
+        """Records the start time of a step."""
+        self.step_start_time = time.monotonic()
+
+    def record_step_end(self, step_name: str, step_index: int | None = None):
+        """Records the end time of a step and calculates the duration."""
+        duration = time.monotonic() - self.step_start_time
+        self.step_timings.append(
+            {"name": step_name, "index": step_index, "duration_ms": duration * 1000}
+        )
+
+    def log_total_duration(self, tag: str):
+        """Logs the total duration of the operation and all recorded steps."""
+        total_duration = time.monotonic() - self.start_time
+        log_entry = {
+            "timestamp": datetime.now(UTC).isoformat(),
+            "request_id": self.request_id,
+            "commit_hash": self.commit_hash,
+            "tag": tag,
+            "total_duration_ms": total_duration * 1000,
+            "steps": self.step_timings,
+        }
+        perf_logger.info(json.dumps(log_entry))