[diffusion] logging: improve request and component load logs (#19253)
This commit is contained in:
@@ -19,6 +19,7 @@ from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
|||||||
from sglang.multimodal_gen.runtime.loader.utils import (
|
from sglang.multimodal_gen.runtime.loader.utils import (
|
||||||
_normalize_component_type,
|
_normalize_component_type,
|
||||||
component_name_to_loader_cls,
|
component_name_to_loader_cls,
|
||||||
|
get_memory_usage_of_component,
|
||||||
)
|
)
|
||||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||||
@@ -124,14 +125,35 @@ class ComponentLoader(ABC):
|
|||||||
if isinstance(component, nn.Module):
|
if isinstance(component, nn.Module):
|
||||||
component = component.eval()
|
component = component.eval()
|
||||||
current_gpu_mem = current_platform.get_available_gpu_memory()
|
current_gpu_mem = current_platform.get_available_gpu_memory()
|
||||||
|
model_size = get_memory_usage_of_component(component)
|
||||||
consumed = gpu_mem_before_loading - current_gpu_mem
|
consumed = gpu_mem_before_loading - current_gpu_mem
|
||||||
logger.info(
|
|
||||||
f"Loaded %s: %s ({source} version). consumed: %.2f GB, avail mem: %.2f GB",
|
# detect component device
|
||||||
component_name,
|
try:
|
||||||
component.__class__.__name__,
|
component_device = str(next(component.parameters()).device)
|
||||||
consumed,
|
is_on_gpu = "cuda" in component_device
|
||||||
current_gpu_mem,
|
except (StopIteration, AttributeError):
|
||||||
)
|
is_on_gpu = False
|
||||||
|
component_device = "unknown"
|
||||||
|
|
||||||
|
if is_on_gpu:
|
||||||
|
logger.info(
|
||||||
|
f"Loaded %s: %s ({source} version). model size: %.2f GB, consumed GPU: %.2f GB, avail GPU mem: %.2f GB",
|
||||||
|
component_name,
|
||||||
|
component.__class__.__name__,
|
||||||
|
model_size,
|
||||||
|
consumed,
|
||||||
|
current_gpu_mem,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
logger.info(
|
||||||
|
f"Loaded %s: %s ({source} version). model size: %.2f GB, device: %s, avail GPU mem: %.2f GB",
|
||||||
|
component_name,
|
||||||
|
component.__class__.__name__,
|
||||||
|
model_size,
|
||||||
|
component_device,
|
||||||
|
current_gpu_mem,
|
||||||
|
)
|
||||||
return component, consumed
|
return component, consumed
|
||||||
|
|
||||||
def load_native(
|
def load_native(
|
||||||
|
|||||||
@@ -179,5 +179,25 @@ def _list_safetensors_files(model_path: str) -> list[str]:
|
|||||||
|
|
||||||
BYTES_PER_GB = 1024**3
|
BYTES_PER_GB = 1024**3
|
||||||
|
|
||||||
|
|
||||||
|
def get_memory_usage_of_component(module) -> float | None:
|
||||||
|
"""
|
||||||
|
returned value is in GB, rounded to 2 decimal digits
|
||||||
|
"""
|
||||||
|
if not isinstance(module, nn.Module):
|
||||||
|
return None
|
||||||
|
if hasattr(module, "get_memory_footprint"):
|
||||||
|
usage = module.get_memory_footprint() / BYTES_PER_GB
|
||||||
|
else:
|
||||||
|
# manually
|
||||||
|
param_size = sum(p.numel() * p.element_size() for p in module.parameters())
|
||||||
|
buffer_size = sum(b.numel() * b.element_size() for b in module.buffers())
|
||||||
|
|
||||||
|
total_size_bytes = param_size + buffer_size
|
||||||
|
usage = total_size_bytes / (1024**3)
|
||||||
|
|
||||||
|
return round(usage, 2)
|
||||||
|
|
||||||
|
|
||||||
# component name -> ComponentLoader class
|
# component name -> ComponentLoader class
|
||||||
component_name_to_loader_cls: Dict[str, Type[Any]] = {}
|
component_name_to_loader_cls: Dict[str, Type[Any]] = {}
|
||||||
|
|||||||
@@ -24,7 +24,10 @@ from sglang.multimodal_gen.configs.sample.teacache import (
|
|||||||
TeaCacheParams,
|
TeaCacheParams,
|
||||||
WanTeaCacheParams,
|
WanTeaCacheParams,
|
||||||
)
|
)
|
||||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
from sglang.multimodal_gen.runtime.server_args import (
|
||||||
|
ServerArgs,
|
||||||
|
_sanitize_for_logging,
|
||||||
|
)
|
||||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||||
from sglang.multimodal_gen.runtime.utils.perf_logger import RequestMetrics
|
from sglang.multimodal_gen.runtime.utils.perf_logger import RequestMetrics
|
||||||
from sglang.multimodal_gen.utils import align_to
|
from sglang.multimodal_gen.utils import align_to
|
||||||
@@ -291,14 +294,20 @@ class Req:
|
|||||||
else:
|
else:
|
||||||
target_width = -1
|
target_width = -1
|
||||||
|
|
||||||
# Log sampling parameters
|
# sanitize prompts for info-level logging
|
||||||
debug_str = f"""Sampling params:
|
sanitized_prompt = _sanitize_for_logging(self.prompt, key_hint="prompt")
|
||||||
|
sanitized_neg_prompt = _sanitize_for_logging(
|
||||||
|
self.negative_prompt, key_hint="negative_prompt"
|
||||||
|
)
|
||||||
|
|
||||||
|
# log non-sensitive parameters at info level
|
||||||
|
info_str = f"""Sampling params:
|
||||||
width: {target_width}
|
width: {target_width}
|
||||||
height: {target_height}
|
height: {target_height}
|
||||||
num_frames: {self.num_frames}
|
num_frames: {self.num_frames}
|
||||||
fps: {self.fps}
|
fps: {self.fps}
|
||||||
prompt: {self.prompt}
|
prompt: {sanitized_prompt}
|
||||||
neg_prompt: {self.negative_prompt}
|
neg_prompt: {sanitized_neg_prompt}
|
||||||
seed: {self.seed}
|
seed: {self.seed}
|
||||||
infer_steps: {self.num_inference_steps}
|
infer_steps: {self.num_inference_steps}
|
||||||
num_outputs_per_prompt: {self.num_outputs_per_prompt}
|
num_outputs_per_prompt: {self.num_outputs_per_prompt}
|
||||||
@@ -310,7 +319,15 @@ class Req:
|
|||||||
save_output: {self.save_output}
|
save_output: {self.save_output}
|
||||||
output_file_path: {self.output_file_path()}
|
output_file_path: {self.output_file_path()}
|
||||||
""" # type: ignore[attr-defined]
|
""" # type: ignore[attr-defined]
|
||||||
|
|
||||||
|
# log full prompts at debug level only (for debugging purposes)
|
||||||
|
debug_str = f"""Full prompts:
|
||||||
|
prompt: {self.prompt}
|
||||||
|
neg_prompt: {self.negative_prompt}
|
||||||
|
"""
|
||||||
|
|
||||||
if not self.suppress_logs:
|
if not self.suppress_logs:
|
||||||
|
logger.info(info_str)
|
||||||
logger.debug(debug_str)
|
logger.debug(debug_str)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -63,10 +63,15 @@ def _sanitize_for_logging(obj: Any, key_hint: str | None = None) -> Any:
|
|||||||
- Render torch.Tensor as a compact summary; if key name is 'scaling_factor', include stats.
|
- Render torch.Tensor as a compact summary; if key name is 'scaling_factor', include stats.
|
||||||
- Dataclasses are expanded to dicts and sanitized recursively.
|
- Dataclasses are expanded to dicts and sanitized recursively.
|
||||||
- Callables/functions are rendered as their qualified name.
|
- Callables/functions are rendered as their qualified name.
|
||||||
|
- Redact sensitive fields like 'prompt' and 'negative_prompt' (only show length).
|
||||||
- Fallback to str(...) for unknown types.
|
- Fallback to str(...) for unknown types.
|
||||||
"""
|
"""
|
||||||
# Handle simple types quickly
|
# Handle simple types quickly
|
||||||
if obj is None or isinstance(obj, (str, int, float, bool)):
|
if obj is None or isinstance(obj, (str, int, float, bool)):
|
||||||
|
# redact sensitive prompt fields
|
||||||
|
if key_hint in ("prompt", "negative_prompt"):
|
||||||
|
if isinstance(obj, str):
|
||||||
|
return f"<redacted, len={len(obj)}>"
|
||||||
return obj
|
return obj
|
||||||
|
|
||||||
# Enum -> value for readability
|
# Enum -> value for readability
|
||||||
@@ -134,7 +139,7 @@ def _sanitize_for_logging(obj: Any, key_hint: str | None = None) -> Any:
|
|||||||
|
|
||||||
# Sequences/Sets -> list
|
# Sequences/Sets -> list
|
||||||
if isinstance(obj, (list, tuple, set)):
|
if isinstance(obj, (list, tuple, set)):
|
||||||
return [_sanitize_for_logging(x) for x in obj]
|
return [_sanitize_for_logging(x, key_hint=key_hint) for x in obj]
|
||||||
|
|
||||||
# Functions / Callables -> qualified name
|
# Functions / Callables -> qualified name
|
||||||
try:
|
try:
|
||||||
|
|||||||
Reference in New Issue
Block a user