tracing

Author: Your Name

vllm/beam/beam.py

@@ -5,6 +5,7 @@
 from vllm.beam.penalty import PenaltyComputer
 import torch
 from vllm.beam.ranking import RankingComputer
+from vllm.beam.tracing import trace_async_method
 from vllm.entrypoints.openai.protocol import CompletionResponse, ErrorResponse, CompletionResponseChoice
 from vllm.logger import init_logger
 
@@ -16,6 +17,7 @@ def __init__(self, classi_idx):
         self.penalty_computer = PenaltyComputer(classi_idx)
         self.ranking_computer = RankingComputer(classi_idx)
 
+    @trace_async_method(span_name='pick_best_beam')
     async def pick_best_beam(self, responses: list[
         Union[AsyncGenerator[str, None], CompletionResponseChoice, ErrorResponse]]) -> Union[
         AsyncGenerator[str, None], CompletionResponseChoice, ErrorResponse]:

vllm/beam/filtering.py

@@ -6,6 +6,7 @@
 import torch
 from starlette.datastructures import MutableHeaders
 
+from vllm.beam.tracing import trace_async_method
 from vllm.entrypoints.openai.protocol import CompletionRequest, CompletionResponse, \
     ErrorResponse, CompletionResponseChoice
 from vllm.logger import init_logger
@@ -32,6 +33,7 @@ def __init__(self, classi_idx, classifier_names):
         self.classi_idx = classi_idx
         self.classifier_names = classifier_names
 
+    @trace_async_method(span_name='get_n_valid_beams')
     async def get_n_valid_beams(self, create_completion: Callable,
                                 request: CompletionRequest,
                                 chunk_num: int,

vllm/beam/tracing.py

@@ -0,0 +1,86 @@
+from functools import wraps
+from typing import Optional, Union, AsyncGenerator
+import time
+
+from vllm.entrypoints.openai.protocol import CompletionRequest, ErrorResponse
+from vllm.tracing import extract_trace_context, SpanAttributes, init_tracer
+from vllm.v1.request import Request
+from opentelemetry import trace
+
+tracer = init_tracer(
+                "vllm.entrypoints.openai.serving_completion",
+                "http://localhost:4317")
+
+def trace_streaming_completion(tracer_attr='tracer'):
+    """
+    Decorator specifically for tracing streaming completion functions.
+    Handles both the initial processing and the async generator.
+    """
+
+    def decorator(func):
+        async def wrapper(self, request: CompletionRequest, raw_request: Request | None = None):
+            ctx = extract_trace_context(dict(raw_request.headers)) if raw_request else None
+            parent_span = tracer.start_span("chunkwise_beam_completion", context=ctx)
+
+            # keep the span current until we’re done
+            scope = trace.use_span(parent_span, end_on_exit=False)
+
+            try:
+                parent_span.set_attribute(SpanAttributes.GEN_AI_REQUEST_MAX_TOKENS, request.max_tokens)
+                parent_span.set_attribute(SpanAttributes.GEN_AI_REQUEST_N, request.n)
+                if hasattr(request, "request_id"):
+                    parent_span.set_attribute(SpanAttributes.GEN_AI_REQUEST_ID, request.request_id)
+
+                gen = await func(self, request, raw_request)
+                if isinstance(gen, ErrorResponse):
+                    parent_span.end()
+                    scope.__exit__(None, None, None)
+                    return gen
+
+                async def traced_generator():
+                    with trace.use_span(parent_span, end_on_exit=False):
+                        with tracer.start_as_current_span("chunk_generation"):
+                            async for item in gen:
+                                yield item
+
+
+                    # now it’s safe to close the parent
+                    parent_span.end()
+                    scope.__exit__(None, None, None)
+
+                return traced_generator()
+
+            except Exception as e:
+                parent_span.record_exception(e)
+                parent_span.end()
+                scope.__exit__(type(e), e, e.__traceback__)
+                raise
+
+        return wrapper
+
+    return decorator
+
+
+def trace_async_method(span_name: Optional[str] = None, tracer_attr='tracer'):
+    """
+    Simple decorator for tracing regular async methods.
+    """
+
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(self, *args, **kwargs):
+            name = span_name or func.__name__
+
+            with tracer.start_as_current_span(name) as span:
+                start_time = time.time()
+                try:
+                    result = await func(self, *args, **kwargs)
+                    span.set_attribute("execution_time_ms", (time.time() - start_time) * 1000)
+                    return result
+                except Exception as e:
+                    span.record_exception(e)
+                    raise
+
+        return wrapper
+
+    return decorator

vllm/entrypoints/openai/serving_completion.py

@@ -15,6 +15,7 @@
 from vllm.beam.filtering import _CHUNK_SIZE, BeamValidator
 from vllm.beam.metrics import report_metrics
 from vllm.beam.penalty import MEOW_CLASSI_IDX, PenaltyComputer
+from vllm.beam.tracing import trace_streaming_completion, trace_async_method
 from vllm.config import ModelConfig
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.logger import RequestLogger
@@ -71,10 +72,8 @@ def __init__(
 
         self.beam_scorer = BeamScorer(classi_idx=MEOW_CLASSI_IDX)
         self.beam_validator = BeamValidator(classi_idx=MEOW_CLASSI_IDX, classifier_names=MEOW_CLASSI_IDX.keys())
-        self.tracer = init_tracer(
-                "vllm.entrypoints.openai.serving_completion",
-                "http://localhost:4317")
 
+    @trace_streaming_completion()
     async def create_completion_with_chunkwise_beam(
         self,
         request: CompletionRequest,
@@ -83,7 +82,7 @@ async def create_completion_with_chunkwise_beam(
         """
     Chunkwise beam search hack
     """
-        
+        @trace_async_method(span_name='_process_prefix')
         async def _process_prefix(request: CompletionRequest):
             og_max_tokens = request.max_tokens
             og_n = request.n