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

@@ -43,14 +43,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)
 
@@ -66,7 +77,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,
     )
 
 
@@ -84,15 +98,14 @@ def httpx_file_client() -> httpx.AsyncClient:
     # requests has no write timeout, keep that
     # httpx default for pool is 5, use that
     timeout = httpx.Timeout(connect=10, read=15, write=None, pool=5)
-    headers = {key: str(value) for key, value in (current_trace_context() or {})}
-    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,
         http2=True,
-        headers=headers,
     )
 
 

python/cog/server/http.py

@@ -287,10 +287,7 @@ async def predict(
         respond_async = prefer == "respond-async"
 
         with trace_context(make_trace_context(traceparent, tracestate)):
-            return _predict(
-                request=request,
-                respond_async=respond_async
-            )
+            return await shared_predict(request=request, respond_async=respond_async)
 
     @limited
     @app.put(
@@ -328,13 +325,10 @@ async def predict_idempotent(
         respond_async = prefer == "respond-async"
 
         with trace_context(make_trace_context(traceparent, tracestate)):
-            return _predict(
-                request=request,
-                respond_async=respond_async,
-            )
+            return await shared_predict(request=request, respond_async=respond_async)
 
 
-    async def _predict(
+    async def shared_predict(
         *, request: Optional[PredictionRequest], respond_async: bool = False
     ) -> Response:
         # [compat] If no body is supplied, assume that this model can be run

python/tests/cog/test_files.py

@@ -10,7 +10,6 @@
 # async def test_put_file_to_signed_endpoint():
 #     mock_fh = io.BytesIO()
 #     mock_client = Mock()
-
 #     mock_response = Mock(spec=requests.Response)
 #     mock_response.status_code = 201
 #     mock_response.text = ""

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
 
@@ -604,6 +606,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")