[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/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} ")