| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import threading |
| 4 | from collections.abc import Callable, Generator, Iterable |
| 5 | from concurrent import futures |
| 6 | from typing import Any, Literal |
| 7 | |
| 8 | import grpc |
| 9 | import pytest |
| 10 | from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( |
| 11 | ExportTraceServiceRequest, |
| 12 | ExportTraceServiceResponse, |
| 13 | ) |
| 14 | from opentelemetry.proto.collector.trace.v1.trace_service_pb2_grpc import ( |
| 15 | TraceServiceServicer, |
| 16 | add_TraceServiceServicer_to_server, |
| 17 | ) |
| 18 | from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue |
| 19 | |
| 20 | FAKE_TRACE_SERVER_ADDRESS = "localhost:4317" |
| 21 | |
| 22 | FieldName = Literal[ |
| 23 | "bool_value", "string_value", "int_value", "double_value", "array_value" |
| 24 | ] |
| 25 | |
| 26 | |
| 27 | def decode_value(value: AnyValue): |
| 28 | """Decode an OpenTelemetry AnyValue protobuf message to a Python value.""" |
| 29 | field_decoders: dict[FieldName, Callable] = { |
| 30 | "bool_value": (lambda v: v.bool_value), |
| 31 | "string_value": (lambda v: v.string_value), |
| 32 | "int_value": (lambda v: v.int_value), |
| 33 | "double_value": (lambda v: v.double_value), |
| 34 | "array_value": ( |
| 35 | lambda v: [decode_value(item) for item in v.array_value.values] |
| 36 | ), |
| 37 | } |
| 38 | for field, decoder in field_decoders.items(): |
| 39 | if value.HasField(field): |
| 40 | return decoder(value) |
| 41 | raise ValueError(f"Couldn't decode value: {value}") |
| 42 | |
| 43 | |
| 44 | def decode_attributes(attributes: Iterable[KeyValue]) -> dict[str, Any]: |
| 45 | """Decode OpenTelemetry KeyValue attributes to a Python dictionary.""" |
| 46 | return {kv.key: decode_value(kv.value) for kv in attributes} |
| 47 | |
| 48 | |
| 49 | class FakeTraceService(TraceServiceServicer): |
| 50 | """A fake gRPC trace service for testing OpenTelemetry trace exports.""" |
| 51 | |
| 52 | def __init__(self): |
| 53 | self.requests: list[ExportTraceServiceRequest] = [] |
| 54 | self.evt = threading.Event() |
| 55 | self._lock = threading.Lock() |
| 56 | |
| 57 | def Export(self, request, context): |
| 58 | with self._lock: |
| 59 | self.requests.append(request) |
| 60 | self.evt.set() |
| 61 | return ExportTraceServiceResponse() |
| 62 | |
| 63 | @property |
| 64 | def request(self) -> ExportTraceServiceRequest | None: |
| 65 | """Returns the first request received (for backward compatibility).""" |
| 66 | with self._lock: |
| 67 | return self.requests[0] if self.requests else None |
| 68 | |
| 69 | def get_all_spans(self) -> list[dict]: |
| 70 | """Returns all spans from all received requests as decoded dicts.""" |
| 71 | spans = [] |
| 72 | with self._lock: |
| 73 | for request in self.requests: |
| 74 | for resource_span in request.resource_spans: |
| 75 | for scope_span in resource_span.scope_spans: |
| 76 | for span in scope_span.spans: |
| 77 | spans.append( |
| 78 | { |
| 79 | "name": span.name, |
| 80 | "attributes": decode_attributes(span.attributes), |
| 81 | "trace_id": span.trace_id.hex(), |
| 82 | "span_id": span.span_id.hex(), |
| 83 | "parent_span_id": span.parent_span_id.hex() |
| 84 | if span.parent_span_id |
| 85 | else None, |
| 86 | "start_time_unix_nano": span.start_time_unix_nano, |
| 87 | "end_time_unix_nano": span.end_time_unix_nano, |
| 88 | } |
| 89 | ) |
| 90 | return spans |
| 91 | |
| 92 | def wait_for_spans(self, count: int = 1, timeout: float = 10) -> bool: |
| 93 | """Wait until at least `count` spans have been received.""" |
| 94 | import time |
| 95 | |
| 96 | deadline = time.time() + timeout |
| 97 | while time.time() < deadline: |
| 98 | if len(self.get_all_spans()) >= count: |
| 99 | return True |
| 100 | time.sleep(0.1) |
| 101 | return False |
| 102 | |
| 103 | def clear(self): |
| 104 | """Clear all received requests.""" |
| 105 | with self._lock: |
| 106 | self.requests.clear() |
| 107 | self.evt.clear() |
| 108 | |
| 109 | |
| 110 | def _wait_for_server_ready(address: str, timeout: float = 5.0) -> bool: |
| 111 | """Wait for the gRPC server to be ready to accept connections.""" |
| 112 | import socket |
| 113 | import time |
| 114 | |
| 115 | host, port = address.rsplit(":", 1) |
| 116 | deadline = time.monotonic() + timeout |
| 117 | while time.monotonic() < deadline: |
| 118 | try: |
| 119 | with socket.create_connection((host, int(port)), timeout=0.5): |
| 120 | return True |
| 121 | except (OSError, ConnectionRefusedError): |
| 122 | time.sleep(0.1) |
| 123 | return False |
| 124 | |
| 125 | |
| 126 | @pytest.fixture |
| 127 | def trace_service() -> Generator[FakeTraceService, None, None]: |
| 128 | """Fixture to set up a fake gRPC trace service.""" |
| 129 | server = grpc.server(futures.ThreadPoolExecutor(max_workers=2)) |
| 130 | service = FakeTraceService() |
| 131 | add_TraceServiceServicer_to_server(service, server) |
| 132 | server.add_insecure_port(FAKE_TRACE_SERVER_ADDRESS) |
| 133 | server.start() |
| 134 | |
| 135 | # Wait for the server to be ready to accept connections |
| 136 | if not _wait_for_server_ready(FAKE_TRACE_SERVER_ADDRESS): |
| 137 | server.stop(grace=None) |
| 138 | raise RuntimeError( |
| 139 | f"Fake trace server failed to start on {FAKE_TRACE_SERVER_ADDRESS}" |
| 140 | ) |
| 141 | |
| 142 | yield service |
| 143 | |
| 144 | server.stop(grace=None) |