feat(llmobs): support distributed tracing

[Yun-Kim]

TLDR for Reviewers

This PR adds support for distributed tracing for LLM Observability, by propagating a LLMObs parent ID tag on distributed requests. Note that the notion of a LLMObs parent ID != APM parent ID because we only submit llm type spans to LLM Observability.

The files to review are (in order of importance):

• ddtrace/propagation/http.py - inject LLMobs parent ID in distributed contexts
• ddtrace/llmobs/_utils.py - implementation details of how we determine / inject LLMObs parent ID
• ddtrace/llmobs/_llmobs.py - add edge case handling, explained below
• tests/llmobs/test_propagation.py - test cases showing (hopefully) what is expected behavior
• tests/tracer/test_propagation.py - test cases for general propagation cases, mostly ensuring that our behavior can be turned on/off.

Context

LLM Obs parent ID

LLM Observability workflows involve marking tracer-generated spans with a specific llm span type, then extracting information from the given span (such as span ID, trace ID, start/end timestamps, errors, etc) to create a LLMObs-specific span event to be submitted to LLM Obs.

However since LLM Obs only accepts spans of type llm (or spans generated by the openai/langchain/bedrock integrations), we do not send other auto-generated spans to LLM Obs. This means that parent IDs can get tricky, especially in cases with non-LLMObs spans sprinkled between LLMObs spans (see example below):

Span 1 (LLM Obs)
|--> Span 2 (Non-LLM Obs)
        | --> Span 3 (LLM Obs) 

Technically span 3's parent is span 2, but since we only submit LLMObs type spans to LLM Obs (i.e. span 1 and span 3), LLM Observability's trace structure breaks down as it has no information about span 2. Therefore, we need to set span 3's llmobs_parent_id to be span 1, in other words the nearest ancestor of type llm.

The current solution is to go up the span's ancestor tree (which is connected via span._parent) until we hit a span of type llm, and use that span's span ID as the llmobs_parent_id. This works fine in most cases, except for distributed scenarios.

Problem in distributed tracing

In distributed scenarios, a trace can encompass multiple services connected via requests. However, while the immediate trace/parent IDs are automatically propagated on request headers, we have no access to the spans beyond the immediate parent, i.e. span._parent is None. This means that in distributed scenarios involving LLMObs spans, the first LLMObs span in each service will consider itself the root.

===================== Service A
Span 1 (LLM Obs)
|--> Span 2 (Non-LLM Obs)
===================== Service B (No access to spans from service A, just distributed request headers)
        | --> Span 3 (LLM Obs) 

Solution

This PR adds three things:

1. Injects the context propagated in a distributed request with the llmobs_parent_id.

By injecting the context with _dd.p.llmobs_parent_id, this will automatically be propagated on distributed request headers/tracecontexts, then be tagged on all spans in the next called service at span start. Note that at span finish time in the original service, the _dd.p.llmobs_parent_id tag also gets set on the local root span as well.

2. On LLMObs manual span startup, if there are no propagated parent IDs on the span, we assume that it is the first service in a distributed trace, i.e. set the span's llmobs_parent_id parent ID manually as a tag.

Since _dd.p.llmobs_parent_id tags are set at span finish time in the original service, we need to add this manual check to avoid relying exclusively on the propagated _dd.p.llmobs_parent_id tag for non-distributed cases or spans in the original service.

3. On span processing to be exported to LLMObs, do three checks:

• Check if the span has manually set _ml_obs.parent_id tags. If so, use that as the parent ID.
• Go up the span's ancestor tree to find the nearest LLMObs type ancestor span. If we find one, then use its span ID as the parent ID.
• If no spans are available in the ancestor tree, then it must be a distributed case. Use the local root span (i.e. the first span in the service)'s propagated _dd.p.llmobs_parent_id tag as the parent ID.

Testing

Unit tests were added to ensure injection is performed as expected when called explicitly. However, there are some manual local testing that was performed to ensure that propagation does indeed happen as expected, with 3 test services (A which calls B, which calls C), all of which were instrumented with an LLMObs span, running using FastAPI and making requests via the requests library.

APM Trace (for context)

This trace includes non-LLMObs spans, including FastAPI spans, requests spans, and manually constructed `APM` spans.

LLMObs trace(s) before changes

All LLMObs spans in each service are the roots of their own trace.

LLMObs trace after changes

All spans in the distributed trace are collected together in LLMObs.

Considerations

This solution relies on using the _dd.p.* context tag propagation for distributed traces in the tracer internals. However, one caveat of using this internal functionality is that the context tag not only gets propagated in the request to the spans in the next service at span start time, but this also gets set on the local root span in the original service at span finish time. This is the reason for manually adding a _ml_obs.parent_id tag in these cases, which takes precedence over the propagated _dd.p.llmobs_parent_id tag.

Additionally, this PR ensures that when LLMObs.enable() is called, we implicitly patch the below distributed tracing integrations, as they are required for propagating llmobs parent IDs via distributed request headers. Note that we only support distributed tracing through our current list of supported request integrations, including:

• "aiohttp",
• "asgi",
• "bottle",
• "celery",
• "cherrypy",
• "django",
• "falcon",
• "fastapi",
• "flask",
• "grpc",
• "httplib",
• "httpx",
• "molten",
• "pyramid",
• "requests",
• "rq",
• "sanic",
• "starlette",
• "tornado",
• "urllib3",
• "wsgi"

For users that make requests outside of these libraries, we do not provide support for distributed tracing currently. We will need to create a custom LLMObsPropagator class to enable them to inject/extract the parent IDs, similar to how ddtrace documents here: https://ddtrace.readthedocs.io/en/stable/advanced_usage.html#custom

Checklist

• Change(s) are motivated and described in the PR description
• Testing strategy is described if automated tests are not included in the PR
• Risks are described (performance impact, potential for breakage, maintainability)
• Change is maintainable (easy to change, telemetry, documentation)
• Library release note guidelines are followed or label changelog/no-changelog is set
• Documentation is included (in-code, generated user docs, public corp docs)
• Backport labels are set (if applicable)
• If this PR changes the public interface, I've notified @DataDog/apm-tees.

Reviewer Checklist

• Title is accurate
• All changes are related to the pull request's stated goal
• Description motivates each change
• Avoids breaking API changes
• Testing strategy adequately addresses listed risks
• Change is maintainable (easy to change, telemetry, documentation)
• Release note makes sense to a user of the library
• Author has acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment
• Backport labels are set in a manner that is consistent with the release branch maintenance policy

[pr-commenter]

Benchmarks

Benchmark execution time: 2024-05-23 20:40:34

Comparing candidate commit 8786fb5 in PR branch yunkim/llmobs-parent-id-propagation with baseline commit 3de0cf5 in branch main.

Found 0 performance improvements and 37 performance regressions! Performance is the same for 172 metrics, 9 unstable metrics.

scenario:httppropagationextract-all_styles_all_headers

• 🟥 max_rss_usage [+2.248MB; +2.303MB] or [+10.658%; +10.920%]

scenario:httppropagationextract-b3_headers

• 🟥 max_rss_usage [+1.556MB; +1.937MB] or [+7.219%; +8.987%]

scenario:httppropagationextract-datadog_tracecontext_tracestate_not_propagated_on_trace_id_no_match

• 🟥 max_rss_usage [+1.586MB; +1.636MB] or [+7.287%; +7.518%]

scenario:httppropagationextract-datadog_tracecontext_tracestate_propagated_on_trace_id_match

• 🟥 max_rss_usage [+1.788MB; +2.261MB] or [+8.353%; +10.564%]

scenario:httppropagationextract-full_t_id_datadog_headers

• 🟥 max_rss_usage [+1.675MB; +1.965MB] or [+7.784%; +9.133%]

scenario:httppropagationextract-invalid_priority_header

• 🟥 max_rss_usage [+1.554MB; +1.625MB] or [+7.131%; +7.457%]

scenario:httppropagationextract-invalid_span_id_header

• 🟥 max_rss_usage [+1.590MB; +1.652MB] or [+7.298%; +7.580%]

scenario:httppropagationextract-large_valid_headers_all

• 🟥 max_rss_usage [+1.600MB; +1.655MB] or [+7.339%; +7.593%]

scenario:httppropagationextract-medium_valid_headers_all

• 🟥 max_rss_usage [+1.608MB; +1.660MB] or [+7.381%; +7.619%]

scenario:httppropagationextract-none_propagation_style

• 🟥 max_rss_usage [+1.565MB; +1.610MB] or [+7.423%; +7.640%]

scenario:httppropagationextract-valid_headers_all

• 🟥 max_rss_usage [+2.231MB; +2.316MB] or [+10.559%; +10.960%]

scenario:httppropagationextract-wsgi_invalid_span_id_header

• 🟥 max_rss_usage [+1.524MB; +1.589MB] or [+7.222%; +7.527%]

scenario:httppropagationextract-wsgi_invalid_trace_id_header

• 🟥 max_rss_usage [+1.563MB; +1.605MB] or [+7.428%; +7.628%]

scenario:httppropagationextract-wsgi_large_header_no_matches

• 🟥 max_rss_usage [+1.772MB; +2.091MB] or [+8.245%; +9.727%]

scenario:httppropagationextract-wsgi_large_valid_headers_all

• 🟥 max_rss_usage [+1.595MB; +1.650MB] or [+7.312%; +7.564%]

scenario:httppropagationextract-wsgi_medium_header_no_matches

• 🟥 max_rss_usage [+1.639MB; +1.922MB] or [+7.588%; +8.897%]

scenario:httppropagationextract-wsgi_medium_valid_headers_all

• 🟥 max_rss_usage [+1.618MB; +1.666MB] or [+7.424%; +7.647%]

scenario:httppropagationextract-wsgi_valid_headers_all

• 🟥 max_rss_usage [+1.598MB; +1.650MB] or [+7.333%; +7.571%]

scenario:httppropagationextract-wsgi_valid_headers_basic

• 🟥 max_rss_usage [+1.615MB; +1.736MB] or [+7.701%; +8.278%]

scenario:httppropagationinject-with_all

• 🟥 max_rss_usage [+2.294MB; +2.332MB] or [+10.895%; +11.077%]

scenario:httppropagationinject-with_priority_and_origin

• 🟥 max_rss_usage [+1.526MB; +1.560MB] or [+7.007%; +7.164%]

scenario:httppropagationinject-with_tags

• 🟥 max_rss_usage [+2.264MB; +2.301MB] or [+10.765%; +10.940%]

scenario:httppropagationinject-with_tags_invalid

• 🟥 max_rss_usage [+2.597MB; +2.641MB] or [+12.518%; +12.730%]

scenario:httppropagationinject-with_tags_max_size

• 🟥 max_rss_usage [+1.637MB; +1.674MB] or [+7.801%; +7.976%]

scenario:sethttpmeta-all-disabled

• 🟥 max_rss_usage [+1.621MB; +1.674MB] or [+7.336%; +7.573%]

scenario:sethttpmeta-all-enabled

• 🟥 max_rss_usage [+1.594MB; +1.644MB] or [+7.210%; +7.436%]

scenario:sethttpmeta-no-collectipvariant

• 🟥 max_rss_usage [+2.066MB; +2.315MB] or [+9.629%; +10.790%]

scenario:sethttpmeta-obfuscation-disabled

• 🟥 max_rss_usage [+1.638MB; +1.688MB] or [+7.388%; +7.617%]

scenario:sethttpmeta-obfuscation-no-query

• 🟥 max_rss_usage [+2.064MB; +2.314MB] or [+9.642%; +10.811%]

scenario:sethttpmeta-obfuscation-send-querystring-disabled

• 🟥 max_rss_usage [+1.609MB; +1.683MB] or [+7.546%; +7.891%]

scenario:sethttpmeta-obfuscation-worst-case-explicit-query

• 🟥 max_rss_usage [+1.602MB; +1.839MB] or [+7.390%; +8.482%]

scenario:sethttpmeta-obfuscation-worst-case-implicit-query

• 🟥 max_rss_usage [+1.650MB; +1.703MB] or [+7.736%; +7.985%]

scenario:sethttpmeta-useragentvariant_exists_1

• 🟥 max_rss_usage [+1.539MB; +1.706MB] or [+7.171%; +7.949%]

scenario:sethttpmeta-useragentvariant_exists_2

• 🟥 max_rss_usage [+1.636MB; +1.681MB] or [+7.642%; +7.852%]

scenario:sethttpmeta-useragentvariant_exists_3

• 🟥 max_rss_usage [+1.618MB; +1.668MB] or [+7.552%; +7.784%]

scenario:sethttpmeta-useragentvariant_not_exists_1

• 🟥 max_rss_usage [+1.667MB; +1.721MB] or [+7.546%; +7.790%]

scenario:sethttpmeta-useragentvariant_not_exists_2

• 🟥 max_rss_usage [+2.117MB; +2.233MB] or [+9.824%; +10.363%]

[datadog-dd-trace-py-rkomorn]

Datadog Report

Branch report: yunkim/llmobs-parent-id-propagation
Commit report: a65b91f
Test service: dd-trace-py

✅ 0 Failed, 18622 Passed, 41085 Skipped, 3h 31m 35.56s Total duration (4h 39m 4.86s time saved)

[Yun-Kim]

Note: The http propagation benchmark tests have reported a significant increase in memory usage from this commit: #9152. However the changes from that commit do not justify such a memory increase, and LLMObs should not be enabled in benchmark tests. Additionally, the benchmarks have not run on this branch since that commit. This will need more investigation but I do not believe that this PR is responsible for such a memory increase.

EDIT: after further investigation I've found that importing from ddtrace.llmobs could be responsible for memory usage increase due to indirectly initializing the LLMObs service (even though it is disabled.).