[diffusion] CI: support returning request id from endpoint (#15844)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py

@@ -21,6 +21,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
 from sglang.multimodal_gen.runtime.entrypoints.openai.stores import IMAGE_STORE
 from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
     _parse_size,
+    add_common_data_to_response,
     merge_image_input_list,
     process_generation_batch,
     save_image_to_path,
@@ -145,10 +146,11 @@ async def generations(
                     b64_json=b64,
                     revised_prompt=request.prompt,
                 )
-            ]
+            ],
         }
-        if result.peak_memory_mb and result.peak_memory_mb > 0:
-            response_kwargs["peak_memory_mb"] = result.peak_memory_mb
+        response_kwargs = add_common_data_to_response(
+            response_kwargs, request_id=request_id, result=result
+        )
         return ImageResponse(**response_kwargs)
     else:
         # Return error, not supported
@@ -244,17 +246,19 @@ async def edits(
         with open(save_file_path, "rb") as f:
             b64 = base64.b64encode(f.read()).decode("utf-8")
         response_kwargs = {
-            "data": [ImageResponseData(b64_json=b64, revised_prompt=prompt)]
+            "data": [ImageResponseData(b64_json=b64, revised_prompt=prompt)],
         }
-        if result.peak_memory_mb and result.peak_memory_mb > 0:
-            response_kwargs["peak_memory_mb"] = result.peak_memory_mb
-        return ImageResponse(**response_kwargs)
     else:
         url = f"/v1/images/{request_id}/content"
-        response_kwargs = {"data": [ImageResponseData(url=url, revised_prompt=prompt)]}
-        if result.peak_memory_mb and result.peak_memory_mb > 0:
-            response_kwargs["peak_memory_mb"] = result.peak_memory_mb
-        return ImageResponse(**response_kwargs)
+        response_kwargs = {
+            "data": [ImageResponseData(url=url, revised_prompt=prompt)],
+        }
+
+    response_kwargs = add_common_data_to_response(
+        response_kwargs, request_id=request_id, result=result
+    )
+
+    return ImageResponse(**response_kwargs)
 
 
 @router.get("/{image_id}/content")

python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py

@@ -12,6 +12,7 @@ class ImageResponseData(BaseModel):
 
 
 class ImageResponse(BaseModel):
+    id: str
     created: int = Field(default_factory=lambda: int(time.time()))
     data: List[ImageResponseData]
     peak_memory_mb: Optional[float] = None

python/sglang/multimodal_gen/runtime/entrypoints/openai/stores.py

@@ -42,5 +42,6 @@ class AsyncDictStore:
 
 
 # Global stores shared by OpenAI entrypoints
+# [request_id, dict]
 VIDEO_STORE = AsyncDictStore()
 IMAGE_STORE = AsyncDictStore()

python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py

@@ -10,6 +10,7 @@ import httpx
 from fastapi import UploadFile
 
 from sglang.multimodal_gen.runtime.entrypoints.utils import post_process_sample
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
 from sglang.multimodal_gen.runtime.scheduler_client import AsyncSchedulerClient
 from sglang.multimodal_gen.runtime.utils.logging_utils import (
     init_logger,
@@ -172,7 +173,7 @@ async def _save_base64_image_to_path(base64_data: str, target_path: str) -> str:
 async def process_generation_batch(
     scheduler_client: AsyncSchedulerClient,
     batch,
-):
+) -> tuple[str, OutputBatch]:
     total_start_time = time.perf_counter()
     with log_generation_timer(logger, batch.prompt):
         result = await scheduler_client.forward([batch])
@@ -227,3 +228,14 @@ def merge_image_input_list(*inputs: Union[List, Any, None]) -> List:
             else:
                 result.append(input_item)
     return result
+
+
+def add_common_data_to_response(
+    response: dict, request_id: str, result: OutputBatch
+) -> dict:
+    if result.peak_memory_mb and result.peak_memory_mb > 0:
+        response["peak_memory_mb"] = result.peak_memory_mb
+
+    response["id"] = request_id
+
+    return response

python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py

@@ -30,6 +30,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
 from sglang.multimodal_gen.runtime.entrypoints.openai.stores import VIDEO_STORE
 from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
     _parse_size,
+    add_common_data_to_response,
     merge_image_input_list,
     process_generation_batch,
     save_image_to_path,
@@ -124,8 +125,9 @@ async def _dispatch_job_async(job_id: str, batch: Req) -> None:
             "progress": 100,
             "completed_at": int(time.time()),
         }
-        if result.peak_memory_mb and result.peak_memory_mb > 0:
-            update_fields["peak_memory_mb"] = result.peak_memory_mb
+        update_fields = add_common_data_to_response(
+            update_fields, request_id=job_id, result=result
+        )
         await VIDEO_STORE.update_fields(job_id, update_fields)
     except Exception as e:
         logger.error(f"{e}")
@@ -135,6 +137,7 @@ async def _dispatch_job_async(job_id: str, batch: Req) -> None:
 
 
 # TODO: support image to video generation
+# TODO: this is currently not used
 @router.post("", response_model=VideoResponse)
 async def create_video(
     request: Request,

python/sglang/multimodal_gen/test/server/test_server_common.py

@@ -38,7 +38,6 @@ from sglang.multimodal_gen.test.server.testcase_configs import (
 from sglang.multimodal_gen.test.test_utils import (
     get_dynamic_server_port,
     is_image_url,
-    read_perf_logs,
     wait_for_req_perf_record,
 )
 
@@ -191,15 +190,13 @@ Consider updating perf_baselines.json with the snippets below:
     ) -> RequestPerfRecord:
         """Run generation and collect performance records."""
         log_path = ctx.perf_log_path
-        prev_len = len(read_perf_logs(log_path))
         log_wait_timeout = 30
 
         client = self._client(ctx)
         rid = generate_fn(case_id, client)
 
-        req_perf_record, _ = wait_for_req_perf_record(
+        req_perf_record = wait_for_req_perf_record(
             rid,
-            prev_len,
             log_path,
             timeout=log_wait_timeout,
         )

python/sglang/multimodal_gen/test/server/test_server_utils.py

@@ -733,6 +733,8 @@ def get_generate_fn(
         """
         Create a video job via /v1/videos, poll until completion,
         then download the binary content and validate it.
+
+        Returns request-id
         """
 
         create_kwargs: dict[str, Any] = {
@@ -842,6 +844,8 @@ def get_generate_fn(
         result = response.parse()
         validate_image(result.data[0].b64_json)
 
+        rid = result.id
+
         img_data = base64.b64decode(result.data[0].b64_json)
         # Infer expected format from request parameters
         expected_ext = get_expected_image_format(req_output_format, req_background)
@@ -869,7 +873,7 @@ def get_generate_fn(
         )
         os.remove(tmp_path)
 
-        return str(result.created)
+        return rid
 
     def generate_image_edit(case_id, client) -> str:
         """TI2I: Text + Image ? Image edit."""
@@ -910,12 +914,12 @@ def get_generate_fn(
             for img in images:
                 img.close()
 
-        rid = response.headers.get("x-request-id", "")
-
         result = response.parse()
         validate_image(result.data[0].b64_json)
 
         img_data = base64.b64decode(result.data[0].b64_json)
+        rid = result.id
+
         # Infer expected format from request parameters
         expected_ext = get_expected_image_format(req_output_format, req_background)
         expected_filename = f"{rid}.{expected_ext}"
@@ -976,8 +980,9 @@ def get_generate_fn(
             extra_body={"url": image_urls},
         )
 
-        rid = response.headers.get("x-request-id", "")
         result = response.parse()
+        rid = result.id
+
         validate_image(result.data[0].b64_json)
 
         # Save and upload result for verification

python/sglang/multimodal_gen/test/test_utils.py

@@ -198,10 +198,9 @@ def read_perf_logs(log_path: Path) -> list[RequestPerfRecord]:
 
 def wait_for_req_perf_record(
     request_id: str,
-    prev_len: int,
     log_path: Path,
     timeout: float = 30.0,
-) -> tuple[RequestPerfRecord | None, int]:
+) -> RequestPerfRecord | None:
     """
     the stage metrics of this request should be in the performance_log file with {request-id}
     """
@@ -209,16 +208,14 @@ def wait_for_req_perf_record(
     deadline = time.time() + timeout
     while time.time() < deadline:
         records = read_perf_logs(log_path)
-        if len(records) >= prev_len + 1:
-            # FIXME: unable to get rid from openai apis, this is a hack. we should compare rid
-            # potential error when there are multiple servers
-            return records[-1], len(records)
+        for record in records:
+            if record.request_id == request_id:
+                return record
 
         time.sleep(0.5)
 
     if os.environ.get("SGLANG_GEN_BASELINE", "0") == "1":
-        records = read_perf_logs(log_path)
-        return None, len(records)
+        return None
 
     logger.error(f"record: {records}")
     raise AssertionError(f"Timeout waiting for stage metrics for request {request_id} ")