[diffusion] log: unify generation performance logging (#14117)

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,25 +289,14 @@ 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()
+                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])
-                gen_time = time.perf_counter() - start_time
                     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",
@@ -318,7 +309,9 @@ class DiffGenerator:
                             sample,
                             fps=req.fps,
                             save_output=req.save_output,
-                        save_file_path=req.output_file_path(num_outputs, output_idx),
+                            save_file_path=req.output_file_path(
+                                num_outputs, output_idx
+                            ),
                             data_type=req.data_type,
                         )
 
@@ -327,7 +320,7 @@ class DiffGenerator:
                             "frames": frames,
                             "prompts": req.prompt,
                             "size": (req.height, req.width, req.num_frames),
-                        "generation_time": gen_time,
+                            "generation_time": timer.duration,
                             "timings": (
                                 output_batch.timings.to_dict()
                                 if output_batch.timings
@@ -339,21 +332,11 @@ class DiffGenerator:
                             "prompt_index": output_idx,
                         }
                         results.append(result_item)
-            except Exception as e:
+            except Exception:
-                logger.error(
-                    "Failed to generate output for prompt %d: %s",
-                    request_idx + 1,
-                    e,
-                    exc_info=True,
-                )
                 continue
 
         total_gen_time = time.perf_counter() - total_start_time
-        logger.info(
+        log_batch_completion(logger, len(results), total_gen_time)
-            "Completed batch processing. Generated %d outputs in %.2f seconds.",
-            len(results),
-            total_gen_time,
-        )
 
         if len(results) == 0:
             return None

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,7 +103,10 @@ async def process_generation_batch(
     scheduler_client,
     batch,
 ):
+    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.")
 
@@ -110,4 +118,8 @@ async def process_generation_batch(
             batch.save_output,
             save_file_path,
         )
+
+    total_time = time.perf_counter() - total_start_time
+    log_batch_completion(logger, 1, total_time)
+
     return save_file_path

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,
+    )