Propagate trace context to webhook and upload requests

[aron]

If the request to /predict contains headers traceparent and tracestate defined by w3c Trace Context¹ then these headers are forwarded on to the webhook and upload calls.

This allows observability systems to link requests passing through cog.

Footnotes

1. https://www.w3.org/TR/trace-context/ ↩

[nickstenning]

I think this is ok, but this seems like a clear-cut case where we could use a ContextVar and a couple of helper functions (in, say, a telemetry module) to make this a lot less messy.

The idea would be something like:

# telemetry.py
from contextlib import contextmanager
from contextvars import ContextVar

TRACE_CONTEXT = ContextVar("trace_context", default=None)

def current_trace_context():
    return TRACE_CONTEXT.get()

@contextmanager
def trace_context(ctx):
    t = TRACE_CONTEXT.set(ctx)
    try:
        yield
    finally:
        TRACE_CONTEXT.reset(t)

# http.py
async def predict(
    request: PredictionRequest = Body(default=None),
    prefer: Union[str, None] = Header(default=None),
    traceparent: Union[str, None] = Header(default=None),
    tracestate: Union[str, None] = Header(default=None),
) -> Any:  # type: ignore
    ...
    ctx = {}
    if traceparent:
        ctx['traceparent'] = traceparent
    if tracestate:
        ctx['tracestate'] = tracestate

    if ctx:
        with telemetry.trace_context(ctx):
            return predict(...)
    
    return predict(...)

[aron]

Okay, nice, I'll refactor the code.

[nickstenning]

Generally LGTM other than the handling of absent/None values. The branch diff here also introduces a bunch of spurious whitespace and formatting changes which it would maybe be nice to rebase out, but that's not a huge deal.

[aron]

The branch diff here also introduces a bunch of spurious whitespace and formatting changes which it would maybe be nice to rebase out, but that's not a huge deal.

Hmm, I used ruff format and ruff check to format and validate the files respectively. Perhaps we're not enforcing consistent formatting?

[aron]

Hmm, @nickstenning the ContextVar doesn't seem to propagate across the threads, does this need to be handled manually?

Nope, I'm a dope and wasn't adding the headers to the upload client.

[aron]

I've now rebased on main and run the python directory through ruff:

 python -m ruff check python --fix

I'm leaving that as the canonical whitespace.