[Diffusion] Support peak memory record in offline generate and serving (#15610)

python/sglang/multimodal_gen/benchmarks/bench_serving.py

@@ -63,6 +63,7 @@ class RequestFuncOutput:
     error: str = ""
     start_time: float = 0.0
     response_body: Dict[str, Any] = field(default_factory=dict)
+    peak_memory_mb: float = 0.0
 
 
 class BaseDataset(ABC):
@@ -371,6 +372,8 @@ async def async_request_image_sglang(
                     resp_json = await response.json()
                     output.response_body = resp_json
                     output.success = True
+                    if "peak_memory_mb" in resp_json:
+                        output.peak_memory_mb = resp_json["peak_memory_mb"]
                 else:
                     output.error = f"HTTP {response.status}: {await response.text()}"
                     output.success = False
@@ -398,6 +401,8 @@ async def async_request_image_sglang(
                     resp_json = await response.json()
                     output.response_body = resp_json
                     output.success = True
+                    if "peak_memory_mb" in resp_json:
+                        output.peak_memory_mb = resp_json["peak_memory_mb"]
                 else:
                     output.error = f"HTTP {response.status}: {await response.text()}"
                     output.success = False
@@ -406,6 +411,7 @@ async def async_request_image_sglang(
             output.success = False
 
     output.latency = time.perf_counter() - output.start_time
+
     if pbar:
         pbar.update(1)
     return output
@@ -537,6 +543,8 @@ async def async_request_video_sglang(
                     if status == "completed":
                         output.success = True
                         output.response_body = status_data
+                        if "peak_memory_mb" in status_data:
+                            output.peak_memory_mb = status_data["peak_memory_mb"]
                         break
                     elif status == "failed":
                         output.success = False
@@ -557,6 +565,7 @@ async def async_request_video_sglang(
             break
 
     output.latency = time.perf_counter() - output.start_time
+
     if pbar:
         pbar.update(1)
     return output
@@ -568,6 +577,7 @@ def calculate_metrics(outputs: List[RequestFuncOutput], total_duration: float):
 
     num_success = len(success_outputs)
     latencies = [o.latency for o in success_outputs]
+    peak_memories = [o.peak_memory_mb for o in success_outputs if o.peak_memory_mb > 0]
 
     metrics = {
         "duration": total_duration,
@@ -578,6 +588,9 @@ def calculate_metrics(outputs: List[RequestFuncOutput], total_duration: float):
         "latency_median": np.median(latencies) if latencies else 0,
         "latency_p99": np.percentile(latencies, 99) if latencies else 0,
         "latency_p50": np.percentile(latencies, 50) if latencies else 0,
+        "peak_memory_mb_max": max(peak_memories) if peak_memories else 0,
+        "peak_memory_mb_mean": np.mean(peak_memories) if peak_memories else 0,
+        "peak_memory_mb_median": np.median(peak_memories) if peak_memories else 0,
     }
 
     return metrics
@@ -719,6 +732,24 @@ async def benchmark(args):
     print("{:<40} {:<15.4f}".format("Latency Median (s):", metrics["latency_median"]))
     print("{:<40} {:<15.4f}".format("Latency P99 (s):", metrics["latency_p99"]))
 
+    if metrics["peak_memory_mb_max"] > 0:
+        print(f"{'-' * 50}")
+        print(
+            "{:<40} {:<15.2f}".format(
+                "Peak Memory Max (MB):", metrics["peak_memory_mb_max"]
+            )
+        )
+        print(
+            "{:<40} {:<15.2f}".format(
+                "Peak Memory Mean (MB):", metrics["peak_memory_mb_mean"]
+            )
+        )
+        print(
+            "{:<40} {:<15.2f}".format(
+                "Peak Memory Median (MB):", metrics["peak_memory_mb_median"]
+            )
+        )
+
     print("\n" + "=" * 60)
 
     if args.output_file: