Add support for capturing generated audio duration in TTS responses

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
 	}

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()

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
 

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
 

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)
+""",
+)

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),

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) {