From 0a362d653fb878f557e16f29f237cdb960864ae9 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 29 Nov 2025 12:21:59 +0800 Subject: [PATCH] [diffusion] log: unify generation performance logging (#14117) --- .../entrypoints/diffusion_generator.py | 105 ++++++++---------- .../runtime/entrypoints/openai/utils.py | 36 ++++-- .../runtime/utils/logging_utils.py | 60 ++++++++++ 3 files changed, 128 insertions(+), 73 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 20ed203a7..228e3c87c 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -34,6 +34,8 @@ from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs from sglang.multimodal_gen.runtime.sync_scheduler_client import sync_scheduler_client from sglang.multimodal_gen.runtime.utils.logging_utils import ( init_logger, + log_batch_completion, + log_generation_timer, suppress_loggers, suppress_other_loggers, ) @@ -287,73 +289,54 @@ class DiffGenerator: # 2. send requests to scheduler, one at a time # TODO: send batch when supported for request_idx, req in enumerate(requests): - logger.info( - "Processing prompt %d/%d: %s", - request_idx + 1, - len(requests), - req.prompt[:100], - ) try: - start_time = time.perf_counter() - output_batch = self._send_to_scheduler_and_wait_for_response([req]) - gen_time = time.perf_counter() - start_time - if output_batch.error: - raise Exception(f"{output_batch.error}") + with log_generation_timer( + logger, req.prompt, request_idx + 1, len(requests) + ) as timer: + output_batch = self._send_to_scheduler_and_wait_for_response([req]) + if output_batch.error: + raise Exception(f"{output_batch.error}") - # FIXME: in generate mode, an internal assertion error won't raise an error - logger.info( - "Pixel data generated successfully in %.2f seconds", - gen_time, - ) + if output_batch.output is None: + logger.error( + "Received empty output from scheduler for prompt %d", + request_idx + 1, + ) + continue + for output_idx, sample in enumerate(output_batch.output): + num_outputs = len(output_batch.output) + frames = self.post_process_sample( + sample, + fps=req.fps, + save_output=req.save_output, + save_file_path=req.output_file_path( + num_outputs, output_idx + ), + data_type=req.data_type, + ) - if output_batch.output is None: - logger.error( - "Received empty output from scheduler for prompt %d", - request_idx + 1, - ) - continue - for output_idx, sample in enumerate(output_batch.output): - num_outputs = len(output_batch.output) - frames = self.post_process_sample( - sample, - fps=req.fps, - save_output=req.save_output, - save_file_path=req.output_file_path(num_outputs, output_idx), - data_type=req.data_type, - ) - - result_item: dict[str, Any] = { - "samples": sample, - "frames": frames, - "prompts": req.prompt, - "size": (req.height, req.width, req.num_frames), - "generation_time": gen_time, - "timings": ( - output_batch.timings.to_dict() - if output_batch.timings - else {} - ), - "trajectory": output_batch.trajectory_latents, - "trajectory_timesteps": output_batch.trajectory_timesteps, - "trajectory_decoded": output_batch.trajectory_decoded, - "prompt_index": output_idx, - } - results.append(result_item) - except Exception as e: - logger.error( - "Failed to generate output for prompt %d: %s", - request_idx + 1, - e, - exc_info=True, - ) + result_item: dict[str, Any] = { + "samples": sample, + "frames": frames, + "prompts": req.prompt, + "size": (req.height, req.width, req.num_frames), + "generation_time": timer.duration, + "timings": ( + output_batch.timings.to_dict() + if output_batch.timings + else {} + ), + "trajectory": output_batch.trajectory_latents, + "trajectory_timesteps": output_batch.trajectory_timesteps, + "trajectory_decoded": output_batch.trajectory_decoded, + "prompt_index": output_idx, + } + results.append(result_item) + except Exception: continue total_gen_time = time.perf_counter() - total_start_time - logger.info( - "Completed batch processing. Generated %d outputs in %.2f seconds.", - len(results), - total_gen_time, - ) + log_batch_completion(logger, len(results), total_gen_time) if len(results) == 0: return None diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py index 89c08c222..f26a027ed 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py @@ -1,6 +1,7 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import dataclasses import os +import time from typing import Optional import imageio @@ -11,7 +12,11 @@ from einops import rearrange from fastapi import UploadFile from sglang.multimodal_gen.configs.sample.base import DataType -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.runtime.utils.logging_utils import ( + init_logger, + log_batch_completion, + log_generation_timer, +) logger = init_logger(__name__) @@ -98,16 +103,23 @@ async def process_generation_batch( scheduler_client, batch, ): - result = await scheduler_client.forward([batch]) - if result.output is None: - raise RuntimeError("Model generation returned no output.") + total_start_time = time.perf_counter() + with log_generation_timer(logger, batch.prompt): + result = await scheduler_client.forward([batch]) + + if result.output is None: + raise RuntimeError("Model generation returned no output.") + + save_file_path = str(os.path.join(batch.output_path, batch.output_file_name)) + post_process_sample( + result.output[0], + batch.data_type, + batch.fps, + batch.save_output, + save_file_path, + ) + + total_time = time.perf_counter() - total_start_time + log_batch_completion(logger, 1, total_time) - save_file_path = str(os.path.join(batch.output_path, batch.output_file_name)) - post_process_sample( - result.output[0], - batch.data_type, - batch.fps, - batch.save_output, - save_file_path, - ) return save_file_path diff --git a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py index b6f54c3b6..86d2e7b8d 100644 --- a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py @@ -8,6 +8,7 @@ import datetime import logging import os import sys +import time import warnings from contextlib import contextmanager from functools import lru_cache, partial @@ -420,3 +421,62 @@ def suppress_other_loggers(not_suppress_on_main_rank: bool = False): if should_suppress: for logger_name, level in original_levels.items(): logging.getLogger(logger_name).setLevel(level) + + +class GenerationTimer: + def __init__(self): + self.start_time = 0.0 + self.end_time = 0.0 + self.duration = 0.0 + + +@contextmanager +def log_generation_timer( + logger: logging.Logger, + prompt: str, + request_idx: int | None = None, + total_requests: int | None = None, +): + if request_idx is not None and total_requests is not None: + logger.info( + "Processing prompt %d/%d: %s", + request_idx, + total_requests, + prompt[:100], + ) + else: + max_len = 100 + suffix = "..." if len(prompt) > max_len else "" + logger.info(f"Processing prompt: {prompt[:100]}{suffix}") + + timer = GenerationTimer() + timer.start_time = time.perf_counter() + try: + yield timer + timer.end_time = time.perf_counter() + timer.duration = timer.end_time - timer.start_time + logger.info("Pixel data generated successfully in %.2f seconds", timer.duration) + except Exception as e: + if request_idx is not None: + logger.error( + "Failed to generate output for prompt %d: %s", + request_idx, + e, + exc_info=True, + ) + else: + logger.error( + f"Failed to generate output for prompt: {e}", + exc_info=True, + ) + raise + + +def log_batch_completion( + logger: logging.Logger, num_outputs: int, total_time: float +) -> None: + logger.info( + "Completed batch processing. Generated %d outputs in %.2f seconds.", + num_outputs, + total_time, + )