Propagate trace context to webhook and upload requests

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: Aron Carroll <aron@replicate.com>

Author: aron

python/cog/server/http.py

@@ -260,6 +260,8 @@ async def healthcheck() -> Any:
     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
         """
         Run a single prediction on the model
@@ -272,7 +274,12 @@ async def predict(
         # TODO: spec-compliant parsing of Prefer header.
         respond_async = prefer == "respond-async"
 
-        return _predict(request=request, respond_async=respond_async)
+        return _predict(
+            request=request,
+            respond_async=respond_async,
+            traceparent=traceparent,
+            tracestate=tracestate,
+        )
 
     @limited
     @app.put(
@@ -284,6 +291,8 @@ async def predict_idempotent(
         prediction_id: str = Path(..., title="Prediction ID"),
         request: PredictionRequest = Body(..., title="Prediction Request"),
         prefer: Union[str, None] = Header(default=None),
+        traceparent: Union[str, None] = Header(default=None),
+        tracestate: Union[str, None] = Header(default=None),
     ) -> Any:
         """
         Run a single prediction on the model (idempotent creation).
@@ -307,10 +316,19 @@ async def predict_idempotent(
         # TODO: spec-compliant parsing of Prefer header.
         respond_async = prefer == "respond-async"
 
-        return _predict(request=request, respond_async=respond_async)
+        return _predict(
+            request=request,
+            respond_async=respond_async,
+            traceparent=traceparent,
+            tracestate=tracestate,
+        )
 
     def _predict(
-        *, request: PredictionRequest, respond_async: bool = False
+        *,
+        request: PredictionRequest,
+        respond_async: bool = False,
+        traceparent: Optional[str] = None,
+        tracestate: Optional[str] = None,
     ) -> Response:
         # [compat] If no body is supplied, assume that this model can be run
         # with empty input. This will throw a ValidationError if that's not
@@ -327,7 +345,10 @@ def _predict(
             # async predictions. This is unfortunate but required to ensure
             # backwards-compatible behaviour for synchronous predictions.
             initial_response, async_result = runner.predict(
-                request, upload=respond_async
+                request,
+                upload=respond_async,
+                traceparent=traceparent,
+                tracestate=tracestate,
             )
         except RunnerBusyError:
             return JSONResponse(

python/cog/server/runner.py

@@ -98,7 +98,11 @@ def handle_error(error: BaseException) -> None:
     # TODO: Make the return type AsyncResult[schema.PredictionResponse] when we
     # no longer have to support Python 3.8
     def predict(
-        self, prediction: schema.PredictionRequest, upload: bool = True
+        self,
+        prediction: schema.PredictionRequest,
+        upload: bool = True,
+        traceparent: Optional[str] = None,
+        tracestate: Optional[str] = None,
     ) -> Tuple[schema.PredictionResponse, PredictionTask]:
         # It's the caller's responsibility to not call us if we're busy.
         if self.is_busy():
@@ -119,9 +123,14 @@ def predict(
 
         self._should_cancel.clear()
         upload_url = self._upload_url if upload else None
-        event_handler = create_event_handler(prediction, upload_url=upload_url)
+        event_handler = create_event_handler(
+            prediction,
+            upload_url=upload_url,
+            traceparent=traceparent,
+            tracestate=tracestate,
+        )
 
-        def cleanup(_: schema.PredictionResponse = None) -> None:
+        def cleanup(_: Optional[schema.PredictionResponse] = None) -> None:
             input = cast(Any, prediction.input)
             if hasattr(input, "cleanup"):
                 input.cleanup()
@@ -178,7 +187,10 @@ def cancel(self, prediction_id: Optional[str] = None) -> None:
 
 
 def create_event_handler(
-    prediction: schema.PredictionRequest, upload_url: Optional[str] = None
+    prediction: schema.PredictionRequest,
+    upload_url: Optional[str] = None,
+    traceparent: Optional[str] = None,
+    tracestate: Optional[str] = None,
 ) -> "PredictionEventHandler":
     response = schema.PredictionResponse(**prediction.dict())
 
@@ -189,7 +201,9 @@ def create_event_handler(
 
     webhook_sender = None
     if webhook is not None:
-        webhook_sender = webhook_caller_filtered(webhook, set(events_filter))
+        webhook_sender = webhook_caller_filtered(
+            webhook, set(events_filter), traceparent=traceparent, tracestate=tracestate
+        )
 
     file_uploader = None
     if upload_url is not None:

python/cog/server/webhook.py

@@ -1,5 +1,5 @@
 import os
-from typing import Any, Callable, Set
+from typing import Any, Callable, Set, Optional
 
 import requests
 import structlog
@@ -39,8 +39,12 @@ def _get_version() -> str:
 def webhook_caller_filtered(
     webhook: str,
     webhook_events_filter: Set[WebhookEvent],
+    traceparent: Optional[str] = None,
+    tracestate: Optional[str] = None,
 ) -> Callable[[Any, WebhookEvent], None]:
-    upstream_caller = webhook_caller(webhook)
+    upstream_caller = webhook_caller(
+        webhook, traceparent=traceparent, tracestate=tracestate
+    )
 
     def caller(response: Any, event: WebhookEvent) -> None:
         if event in webhook_events_filter:
@@ -49,13 +53,17 @@ def caller(response: Any, event: WebhookEvent) -> None:
     return caller
 
 
-def webhook_caller(webhook: str) -> Callable[[Any], None]:
+def webhook_caller(
+    webhook: str, traceparent: Optional[str] = None, tracestate: Optional[str] = None
+) -> Callable[[Any], None]:
     # TODO: we probably don't need to create new sessions and new throttlers
     # for every prediction.
     throttler = ResponseThrottler(response_interval=_response_interval)
 
-    default_session = requests_session()
-    retry_session = requests_session_with_retries()
+    default_session = requests_session(traceparent=traceparent, tracestate=tracestate)
+    retry_session = requests_session_with_retries(
+        traceparent=traceparent, tracestate=tracestate
+    )
 
     def caller(response: Any) -> None:
         if throttler.should_send_response(response):
@@ -73,23 +81,33 @@ def caller(response: Any) -> None:
     return caller
 
 
-def requests_session() -> requests.Session:
+def requests_session(
+    traceparent: Optional[str | None] = None, tracestate: Optional[str | None] = None
+) -> requests.Session:
     session = requests.Session()
     session.headers["user-agent"] = (
         _user_agent + " " + str(session.headers["user-agent"])
     )
+
+    if traceparent:
+        session.headers["traceparent"] = traceparent
+    if tracestate:
+        session.headers["tracestate"] = tracestate
+
     auth_token = os.environ.get("WEBHOOK_AUTH_TOKEN")
     if auth_token:
         session.headers["authorization"] = "Bearer " + auth_token
 
     return session
 
 
-def requests_session_with_retries() -> requests.Session:
+def requests_session_with_retries(
+    traceparent: Optional[str | None] = None, tracestate: Optional[str | None] = None
+) -> requests.Session:
     # This session will retry requests up to 12 times, with exponential
     # backoff. In total it'll try for up to roughly 320 seconds, providing
     # resilience through temporary networking and availability issues.
-    session = requests_session()
+    session = requests_session(traceparent=traceparent, tracestate=tracestate)
     adapter = HTTPAdapter(
         max_retries=Retry(
             total=12,

python/tests/server/test_http.py

@@ -531,6 +531,75 @@ def test_asynchronous_prediction_endpoint(client, match):
         time.sleep(0.1)
         n += 1
 
+
+# End-to-end test for passing tracing headers on to downstream services.
+@responses.activate
+@uses_predictor_with_client_options(
+    "output_file", upload_url="https://example.com/upload"
+)
+def test_asynchronous_prediction_endpoint_with_trace_context(client, match):
+    webhook = responses.post(
+        "https://example.com/webhook",
+        match=[
+            matchers.json_params_matcher(
+                {
+                    "id": "12345abcde",
+                    "status": "succeeded",
+                    "output": "https://example.com/upload/file",
+                },
+                strict_match=False,
+            ),
+            matchers.header_matcher(
+                {
+                    "traceparent": "traceparent-123",
+                    "tracestate": "tracestate-123",
+                },
+                strict_match=False,
+            ),
+        ],
+        status=200,
+    )
+    uploader = responses.put(
+        "https://example.com/upload/file",
+        match=[
+            matchers.header_matcher(
+                {
+                    "traceparent": "traceparent-123",
+                    "tracestate": "tracestate-123",
+                },
+                strict_match=False,
+            ),
+        ],
+        status=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 uploader.call_count == 1
     assert webhook.call_count == 1