Skip to content

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

Draft
nkitsaini wants to merge 1 commit into
mainfrom
MCDB-101082/asaini/udf-trace-context
Draft

nkitsaini wants to merge 1 commit into
mainfrom
MCDB-101082/asaini/udf-trace-context

Conversation

@nkitsaini

Copy link
Copy Markdown
Collaborator

Summary

Continues the caller's W3C trace context through the Python UDF server and forwards it on outbound AI calls. This lets one UDF batch request be followed with a single trace ID across nova-gateway nginx, nova-gateway, the UDF, and the Unified Model Gateway (UMG), which already continues traceparent.

UDF server (functions/ext/asgi.py)

  • Extracts traceparent from the incoming request using the OpenTelemetry propagator and attaches it as the current context for the duration of the request. Async UDF tasks and sync UDFs (run via to_thread under copy_context()) inherit it, so outbound calls the UDF makes carry the caller's traceparent.
  • Adds trace_id to the UDF request log records (initiated, timeout, cancelled, error, metrics), and allows trace_id in the JSON log formatter.
  • A missing or invalid traceparent leaves behaviour unchanged. The internal request_id is unchanged.

AI clients

  • New singlestoredb/utils/otel.py holds the header injection helpers that were previously local to ai/chat.py.
  • SingleStoreEmbeddingsFactory now injects traceparent on the OpenAI/Azure httpx clients (sync and async, including caller-provided clients) and on Bedrock requests.
  • Fixes SingleStoreChatFactory with httpx.AsyncClient. AsyncClient awaits its event hooks, so the existing sync request hook raised TypeError on async calls (ainvoke, abatch). Async clients now get a coroutine hook.

Only opentelemetry-api is used: no SDK, exporter, or spans are created. Everything is a no-op when opentelemetry-api is not installed, so it is not a runtime dependency. It is added to the test dependency group so CI runs the new tests. The nova-notebook runtime image already has opentelemetry-api through Pulse.

Behaviour change to note: for OpenAI/Azure embeddings, when the caller does not pass http_client / http_async_client, the factory now builds httpx clients itself (600s timeout, 5s connect), mirroring the chat factory. These use httpx default connection-pool limits instead of the openai SDK defaults.

Test Plan

  • New singlestoredb/tests/test_udf_trace_context.py covers async and sync UDFs: an incoming traceparent is propagated to outbound calls made inside the UDF and trace_id is logged; no header and an invalid header propagate nothing. Both propagation tests fail on main.
  • New tests in test_embeddings.py (OpenAI default and custom clients, Bedrock) and test_ai_chat.py (async client injects traceparent; this fails on main with TypeError).
  • pytest singlestoredb/tests/test_udf_trace_context.py test_embeddings.py test_ai_chat.py test_udf_event_loop.py: 40 passed. pre-commit passes.
  • Checked manually against the real libraries: embeddings sync and async, chat ainvoke, and chat .batch all sent the exact incoming traceparent (5 of 5 requests).

Not yet validated end to end on a cell. Still to check: that UDF container logs reach Loki and that info-level request logs are emitted in the deployed UDF runtime.

Deployment Plan

  • Customer impact: No functional change. It adds a traceparent header on outbound AI requests and a trace_id log field.
  • Rollout Plan: Release singlestoredb, then bump the version in the Helios nova-notebook image (MCDB-103305).
  • Rollback Plan: Revert, or pin the previous singlestoredb version in the image.

JIRA Issues

Closes MCDB-103303 (parent story MCDB-101082). Related: MCDB-103304 (nginx starts the trace, https://github.com/singlestore/helios/pull/28456) and MCDB-103305 (image bump).

…er and AI clients MCDB-103303 #closes

Co-authored-by: Cursor <cursoragent@cursor.com>

This branch was successfully deployed

1 active deployment
Base — 2992d457 Deployed Oct 11, 2026 by nkitsaini via test-coverage #542
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant