diff --git a/pyproject.toml b/pyproject.toml index c0459f2ee..4146a85ae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,6 +51,7 @@ test = [ "fastapi", "ipython", "jupysql", + "opentelemetry-api", "pandas", "parameterized", "polars", diff --git a/singlestoredb/ai/chat.py b/singlestoredb/ai/chat.py index 83fe52442..d842c7d67 100644 --- a/singlestoredb/ai/chat.py +++ b/singlestoredb/ai/chat.py @@ -30,28 +30,9 @@ from botocore import UNSIGNED from botocore.config import Config - -def _inject_otel_headers(headers: Any) -> None: - try: - from opentelemetry.propagate import inject as otel_inject - except ImportError: - return - try: - otel_inject(headers) - except Exception: - return - - -def _httpx_inject_otel(request: httpx.Request) -> None: - _inject_otel_headers(request.headers) - - -def _attach_otel_request_hook( - client: Union[httpx.Client, httpx.AsyncClient], -) -> None: - hooks = client.event_hooks.setdefault('request', []) - if _httpx_inject_otel not in hooks: - hooks.append(_httpx_inject_otel) +from singlestoredb.utils.otel import attach_otel_request_hook as _attach_otel_request_hook +from singlestoredb.utils.otel import httpx_inject_otel as _httpx_inject_otel # noqa: F401 +from singlestoredb.utils.otel import inject_otel_headers as _inject_otel_headers def SingleStoreChatFactory( diff --git a/singlestoredb/ai/embeddings.py b/singlestoredb/ai/embeddings.py index a4a2ef6c9..be61816d2 100644 --- a/singlestoredb/ai/embeddings.py +++ b/singlestoredb/ai/embeddings.py @@ -43,6 +43,9 @@ from botocore import UNSIGNED from botocore.config import Config +from singlestoredb.utils.otel import attach_otel_request_hook +from singlestoredb.utils.otel import inject_otel_headers + _DEFAULT_TOKEN_ID_PLATFORMS = frozenset({'Nova', 'NovaMultiTenant'}) # Throwaway text used to compare a bare encode against a wrapped one. @@ -475,6 +478,7 @@ def _inject_headers(request: Any, **_ignored: Any) -> None: obo_val = obo_token_getter() if obo_val: request.headers['X-S2-OBO'] = obo_val + inject_otel_headers(request.headers) request.headers.pop('X-Amz-Date', None) request.headers.pop('X-Amz-Security-Token', None) @@ -507,8 +511,16 @@ def _inject_headers(request: Any, **_ignored: Any) -> None: api_key=token, model=model_name, ) - if http_client is not None: - openai_kwargs['http_client'] = http_client + http_async_client = kwargs.pop('http_async_client', None) + default_timeout = httpx.Timeout(timeout=600.0, connect=5.0) + if http_client is None: + http_client = httpx.Client(timeout=default_timeout) + if http_async_client is None: + http_async_client = httpx.AsyncClient(timeout=http_client.timeout) + attach_otel_request_hook(http_client) + attach_otel_request_hook(http_async_client) + openai_kwargs['http_client'] = http_client + openai_kwargs['http_async_client'] = http_async_client if info.hosting_platform == 'Azure': # Real OpenAI models: tiktoken is the right tokenizer, and the model name passed diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py index 8cbd15982..94d95e31e 100755 --- a/singlestoredb/functions/ext/asgi.py +++ b/singlestoredb/functions/ext/asgi.py @@ -119,6 +119,46 @@ async def to_thread( return await loop.run_in_executor(None, func_call) +def _extract_trace_context( + headers: Dict[bytes, bytes], +) -> Tuple[Optional[Any], Optional[str]]: + """ + Return the caller's W3C trace context and its hex trace ID. + + Both are ``None`` when ``opentelemetry-api`` is not installed or the + request has no valid ``traceparent`` header. + + """ + try: + from opentelemetry import propagate + from opentelemetry import trace as otel_trace + except ImportError: + return None, None + carrier = { + k.decode('latin-1'): v.decode('latin-1') for k, v in headers.items() + } + ctx = propagate.extract(carrier) + span_context = otel_trace.get_current_span(ctx).get_span_context() + if not span_context.is_valid: + return None, None + return ctx, otel_trace.format_trace_id(span_context.trace_id) + + +def _attach_trace_context(ctx: Optional[Any]) -> Optional[Any]: + """Make ``ctx`` the current OpenTelemetry context; return the detach token.""" + if ctx is None: + return None + from opentelemetry import context as otel_context + return otel_context.attach(ctx) + + +def _detach_trace_context(token: Optional[Any]) -> None: + if token is None: + return + from opentelemetry import context as otel_context + otel_context.detach(token) + + async def _poll_cancel(cancel_event: threading.Event) -> None: """ Return once ``cancel_event`` is set, polling it on the running loop. @@ -1339,6 +1379,8 @@ async def __call__( method = scope['method'] path = tuple(x for x in scope['path'].split('/') if x) headers = dict(scope['headers']) + trace_context, trace_id = _extract_trace_context(headers) + trace_extra = {'trace_id': trace_id} if trace_id else {} content_type = headers.get( b'content-type', @@ -1368,6 +1410,7 @@ async def __call__( 'function_name': func_name.decode('utf-8'), 'content_type': content_type.decode('utf-8'), 'accepts': accepts.decode('utf-8'), + **trace_extra, }, ) @@ -1387,6 +1430,9 @@ async def __call__( input_handler = self.handlers[(content_type, data_version, args_data_format)] output_handler = self.handlers[(accepts, data_version, returns_data_format)] + # The UDF tasks created below copy the current context, so outbound + # calls made by the UDF propagate the caller's traceparent. + trace_token = _attach_trace_context(trace_context) try: all_tasks = [] result = [] @@ -1470,6 +1516,7 @@ async def __call__( 'request_id': request_id, 'function_name': func_name.decode('utf-8'), 'timeout': func_info['timeout'], + **trace_extra, }, ) body = ( @@ -1486,6 +1533,7 @@ async def __call__( 'app_name': self.name, 'request_id': request_id, 'function_name': func_name.decode('utf-8'), + **trace_extra, }, ) body = b'CancelledError: Function call was cancelled' @@ -1499,6 +1547,7 @@ async def __call__( 'request_id': request_id, 'function_name': func_name.decode('utf-8'), 'exception_type': type(e).__name__, + **trace_extra, }, ) msg = traceback.format_exc().strip().split(' File ')[-1] @@ -1512,6 +1561,7 @@ async def __call__( finally: cancel_event.set() await cancel_all_tasks(all_tasks) + _detach_trace_context(trace_token) # Handle api reflection elif method == 'GET' and path == self.show_create_function_path: @@ -1567,6 +1617,7 @@ async def __call__( 'request_id': request_id, 'function_name': timer.metadata.get('function', ''), 'metrics': metrics, + **trace_extra, }, ) diff --git a/singlestoredb/functions/ext/utils.py b/singlestoredb/functions/ext/utils.py index 4eda318b3..e86d9ae59 100644 --- a/singlestoredb/functions/ext/utils.py +++ b/singlestoredb/functions/ext/utils.py @@ -51,7 +51,7 @@ def format(self, record: logging.LogRecord) -> str: # Add extra fields if present allowed_fields = [ - 'app_name', 'request_id', 'function_name', + 'app_name', 'request_id', 'trace_id', 'function_name', 'content_type', 'accepts', 'metrics', ] for field in allowed_fields: diff --git a/singlestoredb/tests/test_ai_chat.py b/singlestoredb/tests/test_ai_chat.py index 6658ce89a..ed77bd071 100644 --- a/singlestoredb/tests/test_ai_chat.py +++ b/singlestoredb/tests/test_ai_chat.py @@ -48,6 +48,39 @@ def test_httpx_hook_appends_once(self): finally: client.close() + def test_async_client_request_injects_trace_context(self): + seen = {} + + def handler(request): + seen['traceparent'] = request.headers.get('traceparent') + return httpx.Response(200) + + def fake_inject(target): + target['traceparent'] = 'trace-context' + + fake_propagate = ModuleType('opentelemetry.propagate') + fake_propagate.inject = fake_inject + fake_otel = ModuleType('opentelemetry') + fake_otel.__path__ = [] + + async def send(): + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler), + ) as client: + chat_mod._attach_otel_request_hook(client) + await client.get('http://localhost/') + + with patch.dict( + 'sys.modules', + { + 'opentelemetry': fake_otel, + 'opentelemetry.propagate': fake_propagate, + }, + ): + import asyncio + asyncio.run(send()) + self.assertEqual(seen['traceparent'], 'trace-context') + @patch('singlestoredb.ai.chat.ChatOpenAI') def test_openai_default_timeout_is_600(self, mock_chat_openai): chat_mod.SingleStoreChatFactory( diff --git a/singlestoredb/tests/test_embeddings.py b/singlestoredb/tests/test_embeddings.py index 00269b86b..38c642b62 100644 --- a/singlestoredb/tests/test_embeddings.py +++ b/singlestoredb/tests/test_embeddings.py @@ -7,7 +7,7 @@ import os import sys import types -import unittest +import unittest.mock # Stand-ins for a model's BOS/EOS, outside the fake tokenizers' ord()-derived range. # Not the real Qwen3 EOS (151643): affixes come from the tokenizer, so no test should @@ -23,6 +23,8 @@ 'botocore.config', 'boto3', 'transformers', + # Imports httpx, so it must be reloaded against the fakes and dropped afterwards. + 'singlestoredb.utils.otel', ) @@ -75,24 +77,37 @@ def __init__(self, **kwargs): class MockClient: + + def __init__(self, timeout=None, **kwargs): + self.timeout = timeout + self.event_hooks = {} + + +class MockAsyncClient(MockClient): pass class MockTimeout: - def __init__(self, connect=None, read=None): - self.connect = connect - self.read = read + def __init__(self, timeout=None, connect=None, read=None): + self.connect = connect if connect is not None else timeout + self.read = read if read is not None else timeout class MockBoto3(types.ModuleType): + # Bedrock before-send handlers registered by the most recent client, by event. + handlers = {} + def client(self, *args, **kwargs): + MockBoto3.handlers = {} + + def register_first(event, handler): + MockBoto3.handlers[event] = handler + return types.SimpleNamespace( _endpoint=types.SimpleNamespace( - _event_emitter=types.SimpleNamespace( - register_first=lambda *args, **kwargs: None, - ), + _event_emitter=types.SimpleNamespace(register_first=register_first), ), ) @@ -150,9 +165,12 @@ def setUpClass(cls): sys.modules.pop('singlestoredb.ai.embeddings', None) sys.modules.pop('_test_embeddings_module', None) + sys.modules.pop('singlestoredb.utils.otel', None) httpx = types.ModuleType('httpx') httpx.Client = MockClient + httpx.AsyncClient = MockAsyncClient + httpx.Request = object httpx.Timeout = MockTimeout sys.modules['httpx'] = httpx @@ -440,6 +458,67 @@ def test_max_input_tokens_override_beats_registry_and_info(self): 'Qwen/Qwen3-Embedding-0.6B', 'Nova', info=info, max_input_tokens=256, ).max_input_tokens == 256 + def test_openai_clients_inject_trace_context(self): + otel = sys.modules['singlestoredb.utils.otel'] + embedding = self.embeddings.SingleStoreEmbeddingsFactory( + model_name='text-embedding-3-small', + api_key='token', + base_url='http://localhost:8000', + hosting_platform='Azure', + ) + + http_client = embedding.kwargs['http_client'] + http_async_client = embedding.kwargs['http_async_client'] + assert isinstance(http_async_client, MockAsyncClient) + assert http_client.event_hooks['request'] == [otel.httpx_inject_otel] + assert http_async_client.event_hooks['request'] == [ + otel.httpx_inject_otel_async, + ] + assert http_client.timeout.read == 600.0 + assert http_client.timeout.connect == 5.0 + assert http_async_client.timeout is http_client.timeout + + def test_openai_custom_client_gets_trace_hook(self): + otel = sys.modules['singlestoredb.utils.otel'] + custom = MockClient(timeout=MockTimeout(42.0)) + embedding = self.embeddings.SingleStoreEmbeddingsFactory( + model_name='text-embedding-3-small', + api_key='token', + base_url='http://localhost:8000', + hosting_platform='Azure', + http_client=custom, + ) + + assert embedding.kwargs['http_client'] is custom + assert custom.event_hooks['request'] == [otel.httpx_inject_otel] + assert embedding.kwargs['http_async_client'].timeout.read == 42.0 + + def test_bedrock_request_injects_trace_context(self): + self.embeddings.SingleStoreEmbeddingsFactory( + model_name='amazon.titan-embed-text-v2:0', + api_key='token', + base_url='http://localhost:8000', + hosting_platform='Amazon', + ) + handler = MockBoto3.handlers['before-send.bedrock-runtime.InvokeModel'] + + def fake_inject(carrier): + carrier['traceparent'] = 'trace-context' + + fake_otel = types.ModuleType('opentelemetry') + fake_otel.__path__ = [] + fake_propagate = types.ModuleType('opentelemetry.propagate') + fake_propagate.inject = fake_inject + request = types.SimpleNamespace(headers={}) + with unittest.mock.patch.dict( + sys.modules, + {'opentelemetry': fake_otel, 'opentelemetry.propagate': fake_propagate}, + ): + handler(request) + + assert request.headers['traceparent'] == 'trace-context' + assert request.headers['Authorization'] == 'Bearer token' + if __name__ == '__main__': unittest.main() diff --git a/singlestoredb/tests/test_udf_trace_context.py b/singlestoredb/tests/test_udf_trace_context.py new file mode 100644 index 000000000..ac8bf6881 --- /dev/null +++ b/singlestoredb/tests/test_udf_trace_context.py @@ -0,0 +1,99 @@ +"""Tests for W3C trace-context propagation through the UDF application.""" +import asyncio +import unittest +from typing import Any +from typing import Dict +from typing import List +from typing import Tuple + +from ..functions import udf +from ..functions.ext.asgi import Application +from .test_udf_event_loop import _make_invoke_args + +try: + from opentelemetry import propagate +except ImportError: + propagate = None # type: ignore[assignment] + + +TRACE_ID = '4bf92f3577b34da6a3ce929d0e0e4736' +TRACEPARENT = f'00-{TRACE_ID}-00f067aa0ba902b7-01' + +# Outbound traceparent each UDF would send, keyed by the UDF's input tag. +_outbound: Dict[str, str] = {} + + +def _record_outbound(tag: str) -> None: + carrier: Dict[str, str] = {} + propagate.inject(carrier) + _outbound[tag] = carrier.get('traceparent', '') + + +@udf +async def _async_trace_udf(tag: str) -> int: + _record_outbound(tag) + return len(tag) + + +@udf +def _sync_trace_udf(tag: str) -> int: + _record_outbound(tag) + return len(tag) + + +@unittest.skipIf(propagate is None, 'opentelemetry-api is not installed') +class TestUDFTraceContext(unittest.TestCase): + + def setUp(self) -> None: + _outbound.clear() + self.app = Application( + functions=[_async_trace_udf, _sync_trace_udf], + app_name='trace_test', + disable_metrics=True, + ) + + def _invoke( + self, + name: str, + tag: str, + extra_headers: List[Tuple[bytes, bytes]], + ) -> List[Any]: + scope, receive, send, sent = _make_invoke_args(name, [(tag,)]) + scope['headers'] = list(scope['headers']) + extra_headers + with self.assertLogs( + 'singlestoredb.functions.ext.asgi.trace_test', level='INFO', + ) as logs: + asyncio.run(self.app(scope, receive, send)) + statuses = [m for m in sent if m.get('type') == 'http.response.start'] + self.assertTrue(statuses and statuses[0]['status'] == 200, sent) + return logs.records + + def test_async_udf_propagates_incoming_traceparent(self) -> None: + records = self._invoke( + '_async_trace_udf', 'alpha', [(b'traceparent', TRACEPARENT.encode())], + ) + self.assertEqual(_outbound['alpha'], TRACEPARENT) + self.assertEqual([getattr(r, 'trace_id') for r in records], [TRACE_ID]) + + def test_sync_udf_propagates_incoming_traceparent(self) -> None: + records = self._invoke( + '_sync_trace_udf', 'beta', [(b'traceparent', TRACEPARENT.encode())], + ) + self.assertEqual(_outbound['beta'], TRACEPARENT) + self.assertEqual([getattr(r, 'trace_id') for r in records], [TRACE_ID]) + + def test_request_without_traceparent_propagates_nothing(self) -> None: + records = self._invoke('_async_trace_udf', 'gamma', []) + self.assertEqual(_outbound['gamma'], '') + self.assertFalse(any(hasattr(r, 'trace_id') for r in records)) + + def test_invalid_traceparent_is_ignored(self) -> None: + records = self._invoke( + '_sync_trace_udf', 'delta', [(b'traceparent', b'00-not-a-trace-01')], + ) + self.assertEqual(_outbound['delta'], '') + self.assertFalse(any(hasattr(r, 'trace_id') for r in records)) + + +if __name__ == '__main__': + unittest.main() diff --git a/singlestoredb/utils/otel.py b/singlestoredb/utils/otel.py new file mode 100644 index 000000000..f1dbd6fe2 --- /dev/null +++ b/singlestoredb/utils/otel.py @@ -0,0 +1,40 @@ +"""OpenTelemetry trace-context propagation for outbound HTTP requests.""" +from typing import Any +from typing import Union + +import httpx + + +def inject_otel_headers(headers: Any) -> None: + """Add the current W3C trace context to ``headers``, if OpenTelemetry is installed.""" + try: + from opentelemetry.propagate import inject as otel_inject + except ImportError: + return + try: + otel_inject(headers) + except Exception: + return + + +def httpx_inject_otel(request: httpx.Request) -> None: + inject_otel_headers(request.headers) + + +async def httpx_inject_otel_async(request: httpx.Request) -> None: + inject_otel_headers(request.headers) + + +def attach_otel_request_hook( + client: Union[httpx.Client, httpx.AsyncClient], +) -> None: + """Inject the trace context into every request sent by ``client``.""" + # httpx.AsyncClient awaits its event hooks, so it needs a coroutine hook. + hook: Any = ( + httpx_inject_otel_async + if isinstance(client, httpx.AsyncClient) + else httpx_inject_otel + ) + hooks = client.event_hooks.setdefault('request', []) + if hook not in hooks: + hooks.append(hook)