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>
This commit is contained in:
Mick
291
python/sglang/multimodal_gen/runtime/utils/common.py
Normal file
291
python/sglang/multimodal_gen/runtime/utils/common.py
Normal file
@@ -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
|
||||
231
python/sglang/multimodal_gen/runtime/utils/distributed.py
Normal file
231
python/sglang/multimodal_gen/runtime/utils/distributed.py
Normal file
@@ -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
|
||||
384
python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py
Normal file
384
python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py
Normal file
@@ -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
|
||||
401
python/sglang/multimodal_gen/runtime/utils/logging_utils.py
Normal file
401
python/sglang/multimodal_gen/runtime/utils/logging_utils.py
Normal file
@@ -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"
|
||||
)
|
||||
@@ -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))
|
||||
Reference in New Issue
Block a user