Support http header trace propagation

[codefromthecrypt]

🚀 Describe the new functionality needed

Right now, if the openai client is instrumented with opentelemetry, it doesn't join the server-side trace. Instead, you get two traces like this:

instrumented openai client

llama-stack with trace sinks

💡 Why is this needed? What if we don't build it?

The most typical use case of distributed tracing is to see traces that cross RPC boundaries. Until propagation flows (e.g. w3c traceparent and/or b3 headers), an application cannot see its side effects on the server, or understand latency or failures knowable on the server side.

Other thoughts

Here's a test script I use which when run with uv applies the env. e.g. uv run -q --env-file .env main.py

# /// script
# dependencies = [
#     "openai",
#     "elastic-opentelemetry",
#     "elastic-opentelemetry-instrumentation-openai",
#     "opentelemetry-instrumentation-httpx"
# ]
# ///
import argparse
import os

import openai
from opentelemetry.instrumentation import auto_instrumentation

auto_instrumentation.initialize()

CHAT_MODEL = os.environ["CHAT_MODEL"]


def main():
    parser = argparse.ArgumentParser(description="OpenTelemetry-Enabled OpenAI Test Client")
    parser.add_argument(
        "--use-responses-api", action="store_true", help="Use the responses API instead of chat completions."
    )
    args = parser.parse_args()

    client = openai.Client()

    messages = [
        {
            "role": "user",
            "content": "Answer in up to 3 words: Which ocean contains Bouvet Island?",
        }
    ]
    extra_body = {"chat_template_kwargs": {"enable_thinking": False}}
    if args.use_responses_api:
        response = client.responses.create(
            model=CHAT_MODEL, input=messages[0]["content"], temperature=0, extra_body=extra_body
        )
        print(response.output[0].content[0].text)
    else:
        chat_completion = client.chat.completions.create(
            model=CHAT_MODEL, messages=messages, temperature=0, extra_body=extra_body
        )
        print(chat_completion.choices[0].message.content)


if __name__ == "__main__":
    main()

Here's the ENV

# Variables used in test script (except CHAT_MODEL these are OpenAI defaults)
OPENAI_BASE_URL=http://localhost:8321/v1/openai/v1
OPENAI_API_KEY=unused
CHAT_MODEL=llama3.2:1b

# Variable name used by llama-stack
INFERENCE_MODEL=llama3.2:1b

# OpenTelemetry configuration
TELEMETRY_SINKS=otel_trace,otel_metric
OTEL_SERVICE_NAME=llama-stack

[frzifus]

That would be great. I faced the same trace fragmentation.

Some additional/other details..

Test client..

import httpx
from llama_stack_client import LlamaStackClient
from llama_stack_client.types import UserMessage

from opentelemetry import trace
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter

from opentelemetry.propagate import inject
from opentelemetry.context import get_current
from opentelemetry.propagators.textmap import default_setter

class OTelHttpxClient(httpx.Client):
    def send(self, request, *args, **kwargs):
        # Inject current trace context into the request headers
        inject(request.headers, context=get_current(), setter=default_setter)
        return super().send(request, *args, **kwargs)

http_client = OTelHttpxClient()
client = LlamaStackClient(base_url="http://localhost:8080", http_client=http_client)

trace.set_tracer_provider(
    TracerProvider(
        resource=Resource.create({"service.name": "llamastack-client-otel"})
    )
)

otlp_trace_exporter = OTLPSpanExporter( # NOTE: Testing backend.
    endpoint="http://backend.observability.svc.cluster.local:4317",
    insecure=True
)

trace.get_tracer_provider().add_span_processor(
    SimpleSpanProcessor(otlp_trace_exporter)
)

tracer = trace.get_tracer(__name__)

with tracer.start_as_current_span("llama_chat_completion"):
    response = client.inference.chat_completion(
        messages=[
            UserMessage(
                content="Write me a 2 sentence poem about the moon",
                role="user",
            ),
        ],
        model_id="vllm", # NOTE: Whatever model my vllm instance has loaded.
        stream=False,
    )
    print(response)

HTTP Request

POST /v1/inference/chat-completion HTTP/1.1
Host: localhost:8321
Accept-Encoding: gzip, deflate
Connection: keep-alive
Accept: application/json
Content-Type: application/json
User-Agent: LlamaStackClient/Python 0.2.4
X-Stainless-Lang: python
X-Stainless-Package-Version: 0.2.4
X-Stainless-OS: Linux
X-Stainless-Arch: x64
X-Stainless-Runtime: CPython
X-Stainless-Runtime-Version: 3.10.17
X-LlamaStack-Client-Version: 0.2.4
X-Stainless-Async: false
x-stainless-retry-count: 0
x-stainless-read-timeout: 60
Content-Length: 117
traceparent: 00-cd7088c08c5a37ba3fc0e27248981a71-0d36bf7af9548756-01

{"messages":[{"content":"Write me a 2 sentence poem about the moon","role":"user"}],"model_id":"vllm","stream":false}

While the trace header is set:

traceparent: 00-cd7088c08c5a37ba3fc0e27248981a71-0d36bf7af9548756-01

Llamastack creates its own root span.

[frzifus]

The same applies to outgoing requests. When llamastack sends a request to e.g. vllm, no trace information are propagated.

Request from llamastack -> vllm.

POST /v1/chat/completions HTTP/1.1
Host: localhost:8000
Accept: application/json
Accept-Encoding: gzip, deflate
Authorization: Bearer fake
Content-Length: 143
Content-Type: application/json
User-Agent: AsyncOpenAI/Python 1.76.1
X-Stainless-Arch: x64
X-Stainless-Async: async:asyncio
X-Stainless-Lang: python
X-Stainless-Os: Linux
X-Stainless-Package-Version: 1.76.1
X-Stainless-Read-Timeout: 600
X-Stainless-Retry-Count: 0
X-Stainless-Runtime: CPython
X-Stainless-Runtime-Version: 3.10.17

{"messages":[{"role":"user","content":[{"type":"text","text":"Berlin is"}]}],"model":"vllm","max_tokens":2048,"stream":false,"temperature":0.0}

[bbrowning]

#2153 addresses the original issue reported here. @frzifus for the issue of outgoing requests initiated by Llama Stack and joining those to the trace, that will take a bit more digging / implementation and is probably best to open that one as a separate issue. What complicates things here is Llama Stack has its own tracing API that knows nothing about opentelemetry, and then has implementations of that API. There's only a single implementation today, and that implementation does know about opentelemetry. However, it means usage of the opentelemetry python libraries and such can only happen within that provider's implementation code, so we can't just do things like auto-instrument in Llama Stack itself as the right contexts aren't available in the places where the calls are getting made.

So, to fix that one, we may have to re-think what the Tracing API provides and whether we want to take a dependency in the core of Llama Stack on opentelemetry or solve this some other way.

[frzifus]

Alright, thanks for this update @bbrowning! I will create a separate tracker for it 👍

[codefromthecrypt]

Thanks for the thinking here. I suspect right if you keep with the internal tracing API, you'll need a propagator function and a context strategy to move the incoming context left to right in the process.

Personally, I don't mind which choices you make on this, and glad to have the improvements meanwhile. FWIW, not many inference platforms do context propagation well, yet, so there's some time to think through the best way out ;)