From 0be2c55132e03c0af3c4109d7ceecb031f2754e7 Mon Sep 17 00:00:00 2001 From: ganisback Date: Sat, 1 Aug 2026 18:02:01 +0800 Subject: [PATCH] Add support for capturing generated audio duration in TTS responses --- .../handler/response_writer_wrapper_speech.go | 12 +- .../response_writer_wrapper_speech_test.go | 18 ++ docker/inference/Dockerfile.vllm | 2 + docker/inference/Dockerfile.vllm-amd | 2 + docker/inference/vllm/patch_tts_duration.py | 184 ++++++++++++++++++ runner/component/service.go | 9 +- runner/component/service_test.go | 3 + 7 files changed, 222 insertions(+), 8 deletions(-) create mode 100644 docker/inference/vllm/patch_tts_duration.py diff --git a/aigateway/handler/response_writer_wrapper_speech.go b/aigateway/handler/response_writer_wrapper_speech.go index 8a45b5409..9e9543ccf 100644 --- a/aigateway/handler/response_writer_wrapper_speech.go +++ b/aigateway/handler/response_writer_wrapper_speech.go @@ -25,8 +25,8 @@ type speechResponseWriter interface { // NewResponseWriterWrapperSpeech wraps the response of an OpenAI-compatible // text-to-speech request (POST /v1/audio/speech). Binary audio responses and // raw audio streams pass through unchanged; SSE responses (stream=true or -// stream_format="sse") are decoded to capture token usage from the terminal -// speech.audio.done event. +// stream_format="sse") are decoded to capture token usage and generated +// duration from the terminal speech.audio.done event. func NewResponseWriterWrapperSpeech(internalWriter http.ResponseWriter, tokenCounter *token.AudioUsageCounter) speechResponseWriter { return &speechAudioResponseWriter{ internalWriter: internalWriter, @@ -89,8 +89,9 @@ func (rw *speechAudioResponseWriter) captureUsage(payload []byte) { return } var event struct { - Type string `json:"type"` - Usage struct { + Type string `json:"type"` + Duration float64 `json:"duration"` + Usage struct { InputTokens int64 `json:"input_tokens"` OutputTokens int64 `json:"output_tokens"` TotalTokens int64 `json:"total_tokens"` @@ -102,6 +103,9 @@ func (rw *speechAudioResponseWriter) captureUsage(payload []byte) { if event.Type != "speech.audio.done" { return } + if event.Duration > 0 { + rw.tokenCounter.Duration(event.Duration) + } if event.Usage.TotalTokens == 0 && event.Usage.InputTokens == 0 && event.Usage.OutputTokens == 0 { return } diff --git a/aigateway/handler/response_writer_wrapper_speech_test.go b/aigateway/handler/response_writer_wrapper_speech_test.go index b85d35709..75d3fe00e 100644 --- a/aigateway/handler/response_writer_wrapper_speech_test.go +++ b/aigateway/handler/response_writer_wrapper_speech_test.go @@ -26,6 +26,24 @@ func TestSpeechAudioResponseWriterCapturesDurationHeader(t *testing.T) { require.Equal(t, 2.5, usage.Duration) } +func TestSpeechAudioResponseWriterCapturesDurationFromSSEDone(t *testing.T) { + counter := token.NewAudioUsageCounter(nil) + recorder := httptest.NewRecorder() + writer := NewResponseWriterWrapperSpeech(recorder, counter) + + writer.Header().Set("Content-Type", "text/event-stream") + writer.WriteHeader(http.StatusOK) + _, err := writer.Write([]byte( + "event: speech.audio.done\n" + + `data: {"type":"speech.audio.done","duration":1.23}` + "\n\n", + )) + require.NoError(t, err) + + usage, err := counter.Usage(context.Background()) + require.NoError(t, err) + require.Equal(t, 1.23, usage.Duration) +} + func TestSpeechBatchResponseWriterSumsDuration(t *testing.T) { counter := token.NewAudioUsageCounter(nil) recorder := httptest.NewRecorder() diff --git a/docker/inference/Dockerfile.vllm b/docker/inference/Dockerfile.vllm index ba1aa7043..099ba95d7 100644 --- a/docker/inference/Dockerfile.vllm +++ b/docker/inference/Dockerfile.vllm @@ -13,6 +13,8 @@ COPY ./supervisord.conf /etc/supervisor/conf.d/supervisord.conf COPY ./vllm/ /etc/csghub/ # Add Audio-Duration-Seconds to successful /v1/audio/transcriptions responses. RUN python3 /etc/csghub/patch_audio_duration_header.py +# Add generated TTS duration to the terminal speech.audio.done SSE event. +RUN python3 /etc/csghub/patch_tts_duration.py RUN chmod +x /etc/csghub/*.sh RUN chmod +x /vllm-workspace/examples/ray_serving/*.sh diff --git a/docker/inference/Dockerfile.vllm-amd b/docker/inference/Dockerfile.vllm-amd index f077ffab4..ab18d2e9b 100644 --- a/docker/inference/Dockerfile.vllm-amd +++ b/docker/inference/Dockerfile.vllm-amd @@ -12,6 +12,8 @@ COPY ./supervisord.conf /etc/supervisor/conf.d/supervisord.conf COPY ./vllm/ /etc/csghub/ # Add Audio-Duration-Seconds to successful /v1/audio/transcriptions responses. RUN python3 /etc/csghub/patch_audio_duration_header.py +# Add generated TTS duration to speech responses. +RUN python3 /etc/csghub/patch_tts_duration.py RUN chmod +x /etc/csghub/*.sh RUN chmod +x /app/vllm/examples/ray_serving/*.sh diff --git a/docker/inference/vllm/patch_tts_duration.py b/docker/inference/vllm/patch_tts_duration.py new file mode 100644 index 000000000..dbf0f4cc6 --- /dev/null +++ b/docker/inference/vllm/patch_tts_duration.py @@ -0,0 +1,184 @@ +"""Patch vLLM-Omni v0.24.0 to expose generated TTS duration. + +The patch is intentionally version-locked. It must fail during the image +build when vLLM-Omni changes, rather than silently producing an image without +the duration field. +""" + +from pathlib import Path + +import vllm_omni + + +def replace_once(path: Path, old: str, new: str) -> None: + content = path.read_text() + if content.count(old) != 1: + raise RuntimeError(f"unexpected vLLM-Omni v0.24.0 source in {path}") + path.write_text(content.replace(old, new)) + + +serving_path = ( + Path(vllm_omni.__file__).resolve().parent + / "entrypoints/openai/serving_speech.py" +) + +replace_once( + serving_path, + """ first_audio_chunk_s: float | None = None + stream_start_s = request_start_s if request_start_s is not None else time.perf_counter() + artifact_ready = False +""", + """ first_audio_chunk_s: float | None = None + total_audio_samples = 0 + stream_start_s = request_start_s if request_start_s is not None else time.perf_counter() + artifact_ready = False +""", +) + +replace_once( + serving_path, + """ if chunk_np.ndim > 1: + chunk_np = chunk_np.squeeze() + # For WAV format, emit header before first audio chunk +""", + """ if chunk_np.ndim > 1: + chunk_np = chunk_np.squeeze() + chunk_array = np.asarray(chunk_np) + chunk_channels = _infer_audio_num_channels(chunk_array) + total_audio_samples += int(chunk_array.size // max(chunk_channels, 1)) + # For WAV format, emit header before first audio chunk +""", +) + +replace_once( + serving_path, + """ else: + logger.info( + "[SpeechE2E] request_id=%s stream=true status=ok total_ms=%.2f first_chunk_ms=NA", + request_id, + total_ms, + ) + except asyncio.CancelledError: +""", + """ else: + logger.info( + "[SpeechE2E] request_id=%s stream=true status=ok total_ms=%.2f first_chunk_ms=NA", + request_id, + total_ms, + ) + if raw_request is not None and total_audio_samples > 0: + raw_request.state.audio_duration_s = total_audio_samples / sample_rate_val + except asyncio.CancelledError: +""", +) + +replace_once( + serving_path, + """ The terminal ``speech.audio.done`` event carries a ``usage`` object + (``input_tokens``/``output_tokens``/``total_tokens`` + a per-modality + ``input_token_details`` breakdown), matching OpenAI's documented + ``speech.audio.done`` schema. ``output_tokens`` is accumulated from the + stage-0 deltas as they stream (see ``SpeechOutputTokenCounter``); + ``input_tokens`` is computed from the request text + reference audio. +""", + """ The terminal ``speech.audio.done`` event carries a ``usage`` object + (``input_tokens``/``output_tokens``/``total_tokens`` + a per-modality + ``input_token_details`` breakdown), matching OpenAI's documented + ``speech.audio.done`` schema. ``output_tokens`` is accumulated from the + stage-0 deltas as they stream (see ``SpeechOutputTokenCounter``); + ``input_tokens`` is computed from the request text + reference audio. + ``duration`` is the generated audio duration in seconds. +""", +) + +replace_once( + serving_path, + """ done_payload: dict[str, Any] = {"type": "speech.audio.done"} + if request is not None: + # Streaming path: output_tokens = sum of stage-0 deltas. + usage = self._build_speech_usage(request, tts_params or {}, usage_acc.total()) + done_payload["usage"] = usage.model_dump() +""", + """ done_payload: dict[str, Any] = {"type": "speech.audio.done"} + duration_s = getattr(raw_request.state, "audio_duration_s", None) if raw_request is not None else None + if duration_s is not None: + done_payload["duration"] = round(duration_s, 3) + if request is not None: + # Streaming path: output_tokens = sum of stage-0 deltas. + usage = self._build_speech_usage(request, tts_params or {}, usage_acc.total()) + done_payload["usage"] = usage.model_dump() +""", +) + +replace_once( + serving_path, + """ usage_out: list[SpeechTokenUsage] | None = None, + ) -> tuple[bytes | str, str]: +""", + """ raw_request: Request | None = None, + usage_out: list[SpeechTokenUsage] | None = None, + ) -> tuple[bytes | str, str]: +""", +) + +replace_once( + serving_path, + """ if hasattr(audio_tensor, "float"): + audio_tensor = audio_tensor.float().detach().cpu().numpy() + + if audio_tensor.ndim > 1: + audio_tensor = audio_tensor.squeeze() + + audio_obj = CreateAudio( + audio_tensor=audio_tensor, + sample_rate=sample_rate, + response_format=request.response_format or "wav", + speed=request.speed or 1.0, + base64_encode=base64_encode, + ) +""", + """ if hasattr(audio_tensor, "float"): + audio_tensor = audio_tensor.float().detach().cpu().numpy() + + if audio_tensor.ndim > 1: + audio_tensor = audio_tensor.squeeze() + + audio_duration_s = None + if sample_rate > 0: + audio_array = np.asarray(audio_tensor) + audio_channels = _infer_audio_num_channels(audio_array) + audio_duration_s = audio_array.size / max(audio_channels, 1) / sample_rate + if raw_request is not None and audio_duration_s > 0: + raw_request.state.audio_duration_s = audio_duration_s + + audio_obj = CreateAudio( + audio_tensor=audio_tensor, + sample_rate=sample_rate, + response_format=request.response_format or "wav", + speed=request.speed or 1.0, + base64_encode=base64_encode, + ) +""", +) + +replace_once( + serving_path, + """ audio_bytes, media_type = await self._generate_audio_bytes(request, request_id=request_id) +""", + """ audio_bytes, media_type = await self._generate_audio_bytes( + request, + request_id=request_id, + raw_request=raw_request, + ) +""", +) + +replace_once( + serving_path, + """ return Response(content=audio_bytes, media_type=media_type) +""", + """ duration_s = getattr(raw_request.state, "audio_duration_s", None) if raw_request is not None else None + headers = {"Audio-Duration-Seconds": f"{duration_s:.3f}"} if duration_s is not None else None + return Response(content=audio_bytes, media_type=media_type, headers=headers) +""", +) diff --git a/runner/component/service.go b/runner/component/service.go index 1c292fd12..0aa55e82f 100644 --- a/runner/component/service.go +++ b/runner/component/service.go @@ -282,10 +282,11 @@ func (s *serviceComponentImpl) generateService(ctx context.Context, cluster *clu PodSpec: corev1.PodSpec{ NodeSelector: nodeSelector, Containers: []corev1.Container{{ - Image: containerImg, - Ports: exposePorts, - Resources: resources, - Env: environments, + Image: containerImg, + ImagePullPolicy: corev1.PullAlways, + Ports: exposePorts, + Resources: resources, + Env: environments, ReadinessProbe: &corev1.Probe{ InitialDelaySeconds: int32(initialDelaySeconds), PeriodSeconds: int32(periodSeconds), diff --git a/runner/component/service_test.go b/runner/component/service_test.go index 8ab69587c..0600dfe45 100644 --- a/runner/component/service_test.go +++ b/runner/component/service_test.go @@ -78,6 +78,9 @@ func TestServiceComponent_RunService(t *testing.T) { kss.EXPECT().Add(mock.Anything, mock.Anything).Return(nil) err := sc.RunService(ctx, req) require.Nil(t, err) + service, err := expectCluster.KnativeClient.ServingV1().Services(sc.k8sNameSpace).Get(ctx, req.SvcName, metav1.GetOptions{}) + require.NoError(t, err) + require.Equal(t, corev1.PullAlways, service.Spec.Template.Spec.Containers[0].ImagePullPolicy) } func TestServiceComponent_StopService(t *testing.T) {