diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 9edb6138cc06..84fed73f2b9f 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -1939,7 +1939,7 @@ class ImageEditRequest(OpenAIBaseModel): description= "Optional edit mask. Currently accepted for compatibility but unsupported.", ) - response_format: Literal["url", "b64_json"] = "url" + response_format: Literal["url", "b64_json", "path"] = "url" output_format: Literal["png", "webp", "jpeg"] = Field( default="png", validation_alias=AliasChoices("output_format", "format"), @@ -2169,9 +2169,9 @@ class VideoJob(OpenAIBaseModel): description= "Server-side paths for n>1 (internal; excluded from the wire).") # exclude=True internal timings, never on the wire (status/list - # model_dump() stays status-only). ``request_started`` is a - # ``perf_counter()`` stamped at the POST handler; the background task - # computes ``total`` from it and stores the header timings + # model_dump() stays status-only). ``request_started`` carries the + # steady-clock ``server_arrival_time``; the background task computes + # ``total`` from it and stores the header timings # (``generation``/``denoise``/``total``) in ``timing_metrics`` so # ``/content`` emits the same Server-Timing header as the sync route. request_started: Optional[float] = Field(default=None, exclude=True) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 001f2e478b1f..dd5669fac0e9 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -3056,6 +3056,7 @@ async def openai_image_generation(self, request: ImageGenerationRequest, with ``request.format`` extended to accept tensor payloads (``"safetensors"``/``"pt"``) alongside the PNG/WebP/JPEG encoders. """ + request_received = raw_request.state.server_arrival_time try: image_id = f"image_{uuid.uuid4().hex}" @@ -3168,8 +3169,9 @@ async def openai_image_generation(self, request: ImageGenerationRequest, logger.info(f"Image {image_id} generated and encoded: " f"latency={latency:.3f}s generation={generation:.3f}s " f"denoise={denoise:.3f}s") + total = get_steady_clock_now_in_seconds() - request_received headers = build_visual_gen_timing_headers( - build_visual_gen_server_timings(metrics)) + build_visual_gen_server_timings(metrics, total=total)) return JSONResponse(content=response.model_dump(), headers=headers) @@ -3247,13 +3249,15 @@ def _reject_disabled_path( ) return None - def _image_object(self, request: ImageGenerationRequest, + def _image_object(self, request: Union[ImageGenerationRequest, + ImageEditRequest], raw_request: Request, image_id: str, i: int, path: Path) -> ImageObject: """Build the per-item ``ImageObject`` for the ``path``/``url`` transports. - ``b64_json`` is handled separately. Shared by the tensor and encoder - branches so they cannot drift when a transport changes. + ``b64_json`` is handled separately. Shared by the generation + (tensor + encoder) and edit routes so they cannot drift when a + transport changes. """ if request.response_format == "path": return ImageObject(path=str(path), revised_prompt=request.prompt) @@ -3316,6 +3320,7 @@ async def _parse_image_edit_request( async def openai_image_edit(self, raw_request: Request) -> Response: """OpenAI-compatible image editing endpoint.""" + request_received = raw_request.state.server_arrival_time if not self._supports_image_edit(): return self._create_not_supported_error( "Image editing is not supported by the loaded visual generation model." @@ -3327,6 +3332,9 @@ async def openai_image_edit(self, raw_request: Request) -> Response: try: request = await self._parse_image_edit_request(raw_request) + path_error = self._reject_disabled_path(request.response_format) + if path_error is not None: + return path_error params = parse_visual_gen_params( request, image_id, @@ -3382,11 +3390,8 @@ async def openai_image_edit(self, raw_request: Request) -> Response: path = self.media_storage_path / f"{image_id}_{i}{ext}" path.write_bytes(image_to_bytes(image, format=pil_format)) data.append( - ImageObject( - url=self._build_image_content_url( - raw_request, image_id, i), - revised_prompt=request.prompt, - )) + self._image_object(request, raw_request, image_id, i, + path)) response = ImageGenerationResponse( created=int(time.time()), @@ -3402,8 +3407,9 @@ async def openai_image_edit(self, raw_request: Request) -> Response: logger.info(f"Image {image_id} edited and encoded: " f"latency={latency:.3f}s generation={generation:.3f}s " f"denoise={denoise:.3f}s") + total = get_steady_clock_now_in_seconds() - request_received headers = build_visual_gen_timing_headers( - build_visual_gen_server_timings(metrics)) + build_visual_gen_server_timings(metrics, total=total)) return JSONResponse(content=response.model_dump(), headers=headers) diff --git a/tensorrt_llm/serve/openai_video_routes.py b/tensorrt_llm/serve/openai_video_routes.py index cf69905ad505..8fbaca541071 100644 --- a/tensorrt_llm/serve/openai_video_routes.py +++ b/tensorrt_llm/serve/openai_video_routes.py @@ -26,6 +26,7 @@ from fastapi.responses import FileResponse, JSONResponse, Response from pydantic import ValidationError +from tensorrt_llm._utils import get_steady_clock_now_in_seconds from tensorrt_llm.logger import logger from tensorrt_llm.media.encoding import resolve_video_format from tensorrt_llm.media.tensor_payload import is_tensor_format @@ -142,9 +143,7 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: - JSON: Send VideoGenerationRequest as application/json - Multipart: Send form fields + optional input_reference file """ - # Stamp request arrival for the Server-Timing ``total`` (full server - # time, incl. request parsing) before any work. - request_received = time.perf_counter() + request_received = raw_request.state.server_arrival_time try: # Client-side ValueErrors from content-type parsing, request # translation, encoder-format preflight, parameter validation, @@ -213,7 +212,7 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: f"Video {video_id} serialized as tensor: latency={latency:.3f}s " f"generation={getattr(output.metrics, 'generation', 0.0):.3f}s" ) - total = time.perf_counter() - request_received + total = get_steady_clock_now_in_seconds() - request_received headers = build_visual_gen_timing_headers( build_visual_gen_server_timings(output.metrics, total=total) ) @@ -252,7 +251,7 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: f"latency={latency:.3f}s generation={generation:.3f}s " f"denoise={denoise:.3f}s" ) - total = time.perf_counter() - request_received + total = get_steady_clock_now_in_seconds() - request_received headers = build_visual_gen_timing_headers( build_visual_gen_server_timings(metrics, total=total) ) @@ -376,10 +375,7 @@ async def openai_video_generation_async( - JSON: Send VideoGenerationRequest as application/json - Multipart: Send form fields + optional input_reference file """ - # Stamp request arrival for the Server-Timing ``total`` (full server - # time, POST arrival -> job completed); the background task reads it - # back off the job to compute ``total``. - request_received = time.perf_counter() + request_received = raw_request.state.server_arrival_time try: # Parse request based on content-type request = await self._parse_video_generation_request(raw_request) @@ -542,7 +538,7 @@ async def _generate_video_background( # from the status wire). ``total`` spans POST arrival -> # completion. total = ( - time.perf_counter() - job.request_started + get_steady_clock_now_in_seconds() - job.request_started if job.request_started is not None else None ) diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py index 1a35d3161abc..27d658371451 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -628,6 +628,39 @@ def test_basic_image_generation_b64(self, image_client): assert len(decoded) > 0 assert img_obj["revised_prompt"] == "A cat sitting on a mat" + def test_image_generation_server_timing_has_total(self, image_client): + """The Server-Timing header carries generation, denoise, and ``total`` + (full server time; real wall-clock, so only checked > 0).""" + resp = image_client.post( + "/v1/images/generations", + json={"prompt": "timing", "response_format": "b64_json", "size": "64x64"}, + ) + assert resp.status_code == 200 + assert _server_timing_ms(resp.headers, "generation") == 1250.0 + assert _server_timing_ms(resp.headers, "denoise") == 750.0 + assert _server_timing_ms(resp.headers, "total") > 0 + + def test_image_generation_total_anchored_to_server_arrival(self, image_client, monkeypatch): + """``total`` measures from the middleware's arrival stamp, not from a + handler-local clock — this route is handed an already-parsed + ``ImageGenerationRequest``, so stamping in the handler would silently + exclude body parsing. + + Backdating only the middleware's clock leaves the handler's end + reading on the real one, so ``total`` must absorb the full offset. + """ + import tensorrt_llm.serve.responses_utils as _ru + + real = _ru.get_steady_clock_now_in_seconds + monkeypatch.setattr(_ru, "get_steady_clock_now_in_seconds", lambda: real() - 5.0) + + resp = image_client.post( + "/v1/images/generations", + json={"prompt": "timing", "response_format": "b64_json", "size": "64x64"}, + ) + assert resp.status_code == 200 + assert _server_timing_ms(resp.headers, "total") >= 5000.0 + def test_image_generation_with_optional_params(self, image_client): resp = image_client.post( "/v1/images/generations", @@ -1060,6 +1093,65 @@ def test_image_edit_default_url_returns_fetchable_output(self, tmp_path, monkeyp assert content.content.startswith(b"\x89PNG\r\n\x1a\n") assert content.headers["content-type"] == "image/png" + def test_image_edit_server_timing_has_total(self, tmp_path, monkeypatch): + """The edit route reports ``total`` too (real wall-clock, so > 0).""" + client, _ = self._client(tmp_path, monkeypatch) + + resp = client.post( + "/v1/images/edits", + json={ + "prompt": "timing", + "image": _b64_white_png_1x1(), + "response_format": "b64_json", + }, + ) + + assert resp.status_code == 200 + assert _server_timing_ms(resp.headers, "generation") == 1250.0 + assert _server_timing_ms(resp.headers, "denoise") == 750.0 + assert _server_timing_ms(resp.headers, "total") > 0 + + def test_image_edit_response_format_path_returns_on_disk_path(self, tmp_path, monkeypatch): + """``response_format='path'`` returns the server-side output path + instead of a fetchable URL, matching ``/v1/images/generations``.""" + client, _ = self._client(tmp_path, monkeypatch) + + resp = client.post( + "/v1/images/edits", + json={ + "prompt": "split layers", + "image": _b64_white_png_1x1(), + "response_format": "path", + }, + ) + + assert resp.status_code == 200 + item = resp.json()["data"][0] + assert item["url"] is None + assert item["path"].startswith(str(tmp_path)) + assert os.path.exists(item["path"]) + + def test_image_edit_response_format_path_rejected_when_disabled(self, tmp_path, monkeypatch): + """``TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1`` gates the edit route too, + so the path transport cannot leak server-side paths on shared + deployments.""" + client, _ = self._client(tmp_path, monkeypatch) + monkeypatch.setenv("TRTLLM_DISALLOW_LOCAL_MEDIA_PATH", "1") + + resp = client.post( + "/v1/images/edits", + json={ + "prompt": "split layers", + "image": _b64_white_png_1x1(), + "response_format": "path", + }, + ) + + assert resp.status_code == 400 + body = resp.json() + _assert_llm_envelope(body, code=400) + assert "path" in body["message"] and "disabled" in body["message"] + def test_image_edit_rejects_json_array_body(self, tmp_path, monkeypatch): """Non-object JSON bodies are client errors, not server errors.""" client, gen = self._client(tmp_path, monkeypatch) @@ -2940,6 +3032,36 @@ async def test_async_content_server_timing_has_total(self, async_video_client): assert _server_timing_ms(content.headers, "denoise") == 750.0 assert _server_timing_ms(content.headers, "total") > 0 + @pytest.mark.asyncio + async def test_async_total_anchored_to_server_arrival(self, async_video_client, monkeypatch): + """The async path round-trips the arrival stamp through + ``VideoJob.request_started`` and closes ``total`` out in a background + task, so the stamp and the end reading must stay on one clock. + """ + import tensorrt_llm.serve.responses_utils as _ru + + real = _ru.get_steady_clock_now_in_seconds + monkeypatch.setattr(_ru, "get_steady_clock_now_in_seconds", lambda: real() - 5.0) + + resp = await async_video_client.post( + "/v1/videos", + json={ + "prompt": "timing", + "size": "32x32", + "seconds": 1.0, + "fps": 8, + "format": "avi", + }, + headers={"content-type": "application/json"}, + ) + assert resp.status_code == 202 + video_id = resp.json()["id"] + assert await _adrive_job_to_completion(async_video_client, video_id) == "completed" + + content = await async_video_client.get(f"/v1/videos/{video_id}/content") + assert content.status_code == 200 + assert _server_timing_ms(content.headers, "total") >= 5000.0 + @pytest.mark.asyncio async def test_async_file_still_returns_file_response(self, async_video_client): """Default and explicit ``response_format='file'`` keep the