Propagate trace context to webhook and upload requests

Based on the implementation in #1698 for sync cog.

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

This allows observability systems to link requests passing through cog.

[^1]: https://www.w3.org/TR/trace-context/

Signed-off-by: technillogue <technillogue@gmail.com>

Author: aron

python/cog/server/clients.py

@@ -46,14 +46,25 @@ def _get_version() -> str:
 WebhookSenderType = Callable[[Any, WebhookEvent], Awaitable[None]]
 
 
-def webhook_headers() -> "dict[str, str]":
+def common_headers() -> "dict[str, str]":
     headers = {"user-agent": _user_agent}
+    return headers
+
+
+def webhook_headers() -> "dict[str, str]":
+    headers = common_headers()
     auth_token = os.environ.get("WEBHOOK_AUTH_TOKEN")
     if auth_token:
         headers["authorization"] = "Bearer " + auth_token
+
     return headers
 
 
+async def on_request_trace_context_hook(request: httpx.Request) -> None:
+    ctx = current_trace_context() or {}
+    request.headers.update(ctx)
+
+
 def httpx_webhook_client() -> httpx.AsyncClient:
     return httpx.AsyncClient(headers=webhook_headers(), follow_redirects=True)
 
@@ -69,7 +80,10 @@ def httpx_retry_client() -> httpx.AsyncClient:
         retryable_methods=["POST"],
     )
     return httpx.AsyncClient(
-        headers=webhook_headers(), transport=transport, follow_redirects=True
+        event_hooks={"request": [on_request_trace_context_hook]},
+        headers=webhook_headers(),
+        transport=transport,
+        follow_redirects=True,
     )
 
 
@@ -91,6 +105,8 @@ def httpx_file_client() -> httpx.AsyncClient:
     headers["User-Agent"] = _user_agent
 
     return httpx.AsyncClient(
+        event_hooks={"request": [on_request_trace_context_hook]},
+        headers=common_headers(),
         transport=transport,
         follow_redirects=True,
         timeout=timeout,

python/cog/server/http.py

@@ -341,10 +341,7 @@ async def predict_idempotent(
         respond_async = prefer == "respond-async"
 
         with trace_context(make_trace_context(traceparent, tracestate)):
-            return shared_predict(
-                request=request,
-                respond_async=respond_async,
-            )
+            return await shared_predict(request=request, respond_async=respond_async)
 
     async def shared_predict(
         *, request: Optional[PredictionRequest], respond_async: bool = False

python/tests/server/test_http.py

@@ -1,5 +1,7 @@
 import base64
+import httpx
 import io
+import respx
 import time
 import unittest.mock as mock
 
@@ -629,6 +631,64 @@ def test_asynchronous_prediction_endpoint_with_trace_context(client, match):
     assert webhook.call_count == 1
 
 
+# End-to-end test for passing tracing headers on to downstream services.
+@pytest.mark.asyncio
+@pytest.mark.respx(base_url="https://example.com")
+@uses_predictor_with_client_options(
+    "output_file", upload_url="https://example.com/upload"
+)
+async def test_asynchronous_prediction_endpoint_with_trace_context(
+    respx_mock: respx.MockRouter, client, match
+):
+    webhook = respx_mock.post(
+        "/webhook",
+        json__id="12345abcde",
+        json__status="succeeded",
+        json__output="https://example.com/upload/file",
+        headers={
+            "traceparent": "traceparent-123",
+            "tracestate": "tracestate-123",
+        },
+    ).respond(200)
+    uploader = respx_mock.put(
+        "/upload/file",
+        headers={
+            "content-type": "application/octet-stream",
+            "traceparent": "traceparent-123",
+            "tracestate": "tracestate-123",
+        },
+    ).respond(200)
+
+    resp = client.post(
+        "/predictions",
+        json={
+            "id": "12345abcde",
+            "input": {},
+            "webhook": "https://example.com/webhook",
+            "webhook_events_filter": ["completed"],
+        },
+        headers={
+            "Prefer": "respond-async",
+            "traceparent": "traceparent-123",
+            "tracestate": "tracestate-123",
+        },
+    )
+    assert resp.status_code == 202
+
+    assert resp.json() == match(
+        {"status": "processing", "output": None, "started_at": mock.ANY}
+    )
+    assert resp.json()["started_at"] is not None
+
+    n = 0
+    while webhook.call_count < 1 and n < 10:
+        time.sleep(0.1)
+        n += 1
+
+    assert webhook.call_count == 1
+    assert uploader.call_count == 1
+
+
 @uses_predictor("sleep")
 def test_prediction_cancel(client):
     resp = client.post("/predictions/123/cancel")