[Refactor] Deduplicate trace header handling in BaseServing (#60828)

Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>

tests/entrypoints/generate/structured_decisions/test_strategies.py

@@ -222,11 +222,12 @@ async def test_request_metadata_reaches_every_question(
     )
     response = await serving.create_decision(request, raw_request)
     assert isinstance(response, StructuredDecisionResponse)
+    engine.is_tracing_enabled.assert_not_awaited()
     assert engine.generate.call_count == 2
     for call in engine.generate.call_args_list:
         assert call.args[0].get("cache_salt") == cache_salt
         assert call.kwargs["trace_headers"] == (
-            TRACE_HEADERS if tracing_enabled and request_headers else None
+            TRACE_HEADERS if request_headers else None
         )
 
 

vllm/entrypoints/generate/base/serving.py

@@ -2,14 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import json
 import time
-from collections.abc import Awaitable, Mapping
+from collections.abc import Awaitable
 from dataclasses import dataclass, field
 from http import HTTPStatus
 from typing import ClassVar, Generic, TypeVar
 
 from fastapi import Request
 from pydantic import ConfigDict
-from starlette.datastructures import Headers
 
 from vllm import RequestOutput
 from vllm.engine.protocol import EngineClient
@@ -33,11 +32,6 @@ from vllm.logprobs import Logprob, PromptLogprobs
 from vllm.lora.request import LoRARequest
 from vllm.tokenizers import TokenizerLike
 from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens
-from vllm.tracing import (
-    contains_trace_headers,
-    extract_trace_headers,
-    log_tracing_disabled_warning,
-)
 from vllm.v1.metrics.stats import RequestStateStats
 
 logger = init_logger(__name__)
@@ -222,20 +216,6 @@ class GenerateBaseServing(BaseServing, BeamSearchOnlineMixin):
             status_code=e.status_code,
         )
 
-    async def _get_trace_headers(
-        self,
-        headers: Headers,
-    ) -> Mapping[str, str] | None:
-        is_tracing_enabled = await self.engine_client.is_tracing_enabled()
-
-        if is_tracing_enabled:
-            return extract_trace_headers(headers)
-
-        if contains_trace_headers(headers):
-            log_tracing_disabled_warning()
-
-        return None
-
     @staticmethod
     def _get_data_parallel_rank(raw_request: Request | None) -> int | None:
         """Pulls the data parallel rank from a header, if provided."""

vllm/entrypoints/generate/generative_scoring/serving.py

@@ -11,7 +11,6 @@ logits (task="generate").
 import asyncio
 import math
 import time
-from collections.abc import Mapping
 from typing import Any, Literal
 
 from fastapi import Request
@@ -32,11 +31,6 @@ from vllm.inputs import EngineInput, tokens_input
 from vllm.logger import init_logger
 from vllm.sampling_params import SamplingParams
 from vllm.tokenizers import TokenizerLike
-from vllm.tracing import (
-    contains_trace_headers,
-    extract_trace_headers,
-    log_tracing_disabled_warning,
-)
 from vllm.utils import random_uuid
 
 logger = init_logger(__name__)
@@ -462,17 +456,3 @@ class ServingGenerativeScoring(BaseServing):
                 token_id: math.exp(logprob)
                 for token_id, logprob in label_logprobs.items()
             }
-
-    async def _get_trace_headers(
-        self,
-        headers: Mapping[str, str],
-    ) -> Mapping[str, str] | None:
-        """Extract trace headers from request headers."""
-        if not contains_trace_headers(headers):
-            return None
-
-        if not await self.engine_client.is_tracing_enabled():
-            log_tracing_disabled_warning()
-            return None
-
-        return extract_trace_headers(headers)

vllm/entrypoints/generate/structured_decisions/serving.py

@@ -2,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import asyncio
 import json
-from collections.abc import Mapping
 from typing import Any
 
 from fastapi import Request
@@ -12,11 +11,6 @@ from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.entrypoints.serve.engine.serving import BaseServing
 from vllm.entrypoints.serve.utils.request_logger import RequestLogger
 from vllm.logger import init_logger
-from vllm.tracing import (
-    contains_trace_headers,
-    extract_trace_headers,
-    log_tracing_disabled_warning,
-)
 
 from .protocol import (
     DecisionUsage,
@@ -74,19 +68,10 @@ class ServingStructuredDecisions(BaseServing):
         super().__init__(
             models=models, model_config=model_config, request_logger=request_logger
         )
+        self.engine_client = strategy.context.engine_client
         self.strategy = strategy
         self.limits = strategy.limits()
 
-    async def _get_trace_headers(
-        self, headers: Mapping[str, str]
-    ) -> Mapping[str, str] | None:
-        if not contains_trace_headers(headers):
-            return None
-        if not await self.strategy.context.engine_client.is_tracing_enabled():
-            log_tracing_disabled_warning()
-            return None
-        return extract_trace_headers(headers)
-
     async def create_decision(
         self,
         request: StructuredDecisionRequest,

vllm/entrypoints/pooling/base/serving.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import asyncio
 from abc import ABC, abstractmethod
-from collections.abc import AsyncGenerator, Mapping
+from collections.abc import AsyncGenerator
 from concurrent.futures import Executor
 from http import HTTPStatus
 from typing import ClassVar
@@ -10,7 +10,6 @@ from typing import ClassVar
 import torch
 from fastapi import Request
 from fastapi.responses import Response
-from starlette.datastructures import Headers
 
 from vllm import PoolingRequestOutput, envs
 from vllm.config import VllmConfig
@@ -24,11 +23,6 @@ from vllm.entrypoints.serve.utils.request_logger import RequestLogger
 from vllm.exceptions import GenerationError, RetryableRequestError
 from vllm.lora.request import LoRARequest
 from vllm.renderers.base import BaseRenderer
-from vllm.tracing import (
-    contains_trace_headers,
-    extract_trace_headers,
-    log_tracing_disabled_warning,
-)
 from vllm.utils.async_utils import make_async, merge_async_iterators
 
 from ...serve.engine.protocol import ErrorResponse
@@ -262,20 +256,6 @@ class PoolingBaseServing(ABC, BaseServing):
 
         return None
 
-    async def _get_trace_headers(
-        self,
-        headers: Headers,
-    ) -> Mapping[str, str] | None:
-        is_tracing_enabled = await self.engine_client.is_tracing_enabled()
-
-        if is_tracing_enabled:
-            return extract_trace_headers(headers)
-
-        if contains_trace_headers(headers):
-            log_tracing_disabled_warning()
-
-        return None
-
 
 class PoolingServing(PoolingBaseServing, ABC):
     def __init__(self, *args, **kwargs):

vllm/entrypoints/serve/engine/serving.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Mapping
 from http import HTTPStatus
 
 from fastapi import Request
@@ -23,6 +24,7 @@ from vllm.renderers.inputs.preprocess import (
     extract_prompt_len,
 )
 from vllm.sampling_params import BeamSearchParams
+from vllm.tracing import contains_trace_headers, extract_trace_headers
 from vllm.utils import random_uuid
 
 
@@ -37,6 +39,13 @@ class BaseServing:
         self.model_config = model_config
         self.request_logger = request_logger
 
+    async def _get_trace_headers(
+        self, headers: Mapping[str, str]
+    ) -> Mapping[str, str] | None:
+        if not contains_trace_headers(headers):
+            return None
+        return extract_trace_headers(headers)
+
     async def _check_model(
         self,
         request: AnyRequest | AnyPoolingRequest,

vllm/v1/engine/async_llm.py

@@ -44,7 +44,7 @@ from vllm.renderers.inputs.preprocess import extract_prompt_components
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.tasks import SupportedTask
 from vllm.tokenizers import TokenizerLike
-from vllm.tracing import init_tracer
+from vllm.tracing import init_tracer, log_tracing_disabled_warning
 from vllm.transformers_utils.config import maybe_register_config_serialize_by_value
 from vllm.usage.usage_lib import UsageContext
 from vllm.utils.async_utils import cancel_task_threadsafe
@@ -403,6 +403,10 @@ class AsyncLLM(EngineClient):
         if self.errored:
             raise EngineDeadError()
 
+        if trace_headers and self.observability_config.otlp_traces_endpoint is None:
+            log_tracing_disabled_warning()
+            trace_headers = None
+
         is_pooling = isinstance(params, PoolingParams)
 
         if (