[backend:feature] Propagate W3C trace context through Python UDF server and AI clients MCDB-103303 #closes

pyproject.toml

@@ -51,6 +51,7 @@ test = [
     "fastapi",
     "ipython",
     "jupysql",
+    "opentelemetry-api",
     "pandas",
     "parameterized",
     "polars",

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(

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

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,
                 },
             )
 

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:

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(

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()