From fa5cfd48d561989d2f5e7fd42788e2f6cfb479c6 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Sat, 10 Oct 2026 13:27:58 +0800 Subject: [PATCH] [Refactor] Deduplicate trace header handling in BaseServing (#60828) Signed-off-by: chaunceyjiang --- .../structured_decisions/test_strategies.py | 3 ++- vllm/entrypoints/generate/base/serving.py | 22 +------------------ .../generate/generative_scoring/serving.py | 20 ----------------- .../generate/structured_decisions/serving.py | 17 +------------- vllm/entrypoints/pooling/base/serving.py | 22 +------------------ vllm/entrypoints/serve/engine/serving.py | 9 ++++++++ vllm/v1/engine/async_llm.py | 6 ++++- 7 files changed, 19 insertions(+), 80 deletions(-) diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index a63e258752..41b9a8f012 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/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 ) diff --git a/vllm/entrypoints/generate/base/serving.py b/vllm/entrypoints/generate/base/serving.py index d1f3e94b00..a0eda6d05f 100644 --- a/vllm/entrypoints/generate/base/serving.py +++ b/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.""" diff --git a/vllm/entrypoints/generate/generative_scoring/serving.py b/vllm/entrypoints/generate/generative_scoring/serving.py index 74dbfd0759..92c0847716 100644 --- a/vllm/entrypoints/generate/generative_scoring/serving.py +++ b/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) diff --git a/vllm/entrypoints/generate/structured_decisions/serving.py b/vllm/entrypoints/generate/structured_decisions/serving.py index e6e990bd46..2633647a0b 100644 --- a/vllm/entrypoints/generate/structured_decisions/serving.py +++ b/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, diff --git a/vllm/entrypoints/pooling/base/serving.py b/vllm/entrypoints/pooling/base/serving.py index 92cdfaec25..2fb0b8e1bc 100644 --- a/vllm/entrypoints/pooling/base/serving.py +++ b/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): diff --git a/vllm/entrypoints/serve/engine/serving.py b/vllm/entrypoints/serve/engine/serving.py index 8b42449552..15f6adc377 100644 --- a/vllm/entrypoints/serve/engine/serving.py +++ b/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, diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 7e50aa0451..bdab43e925 100644 --- a/vllm/v1/engine/async_llm.py +++ b/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 (