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>

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟢 Approval recommended

The propagation paths, optional-dependency behavior, logging integration, and synchronous/asynchronous clients are implemented consistently and covered by targeted tests.

0 open findings

What changed in this PR

Adds end-to-end W3C trace-context propagation through UDF execution and outbound AI requests.

Changes:

  • Extracts incoming trace context and adds trace IDs to UDF logs.
  • Adds reusable synchronous/asynchronous HTTP header-injection hooks.
  • Extends AI-client tests and test dependencies for OpenTelemetry.
File Description
singlestoredb/​utils/​otel.py Provides reusable OpenTelemetry injection hooks.
singlestoredb/​functions/​ext/​asgi.py Propagates request trace context through UDF execution.
singlestoredb/​functions/​ext/​utils.py Permits trace IDs in JSON logs.
singlestoredb/​ai/​chat.py Uses shared sync/async tracing hooks.
singlestoredb/​ai/​embeddings.py Adds tracing to OpenAI, Azure, and Bedrock requests.
singlestoredb/​tests/​test_udf_trace_context.py Tests UDF propagation and logging.
singlestoredb/​tests/​test_embeddings.py Tests embeddings-client propagation.
singlestoredb/​tests/​test_ai_chat.py Tests asynchronous chat-client propagation.
pyproject.toml Adds OpenTelemetry API to test dependencies.

🧠 Review effort: Balanced


💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

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.

2 participants