[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,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

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

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