OTEL improvements

shepherd_server/base_routes.py

@@ -414,7 +414,12 @@ async def callback(
     # adds otel trace to carrier for next worker
     parent_ctx = extract(json.loads(original_query[1]))
     with tracer.start_as_current_span("callback", context=parent_ctx) as span:
-        span.set_attribute("callback_id", callback_id)
+        kgraph = response["message"]["knowledge_graph"]
+        span.set_attribute("callback.id", callback_id)
+        span.set_attribute("callback.results", len(response["message"]["results"]))
+        span.set_attribute("callback.kg_nodes", len(kgraph.get("nodes", {})))
+        span.set_attribute("callback.kg_edges", len(kgraph.get("edges", {})))
+        span.set_attribute("callback.payload_bytes", len(raw))
         span_carrier = {}
         inject(span_carrier)
         # add new task to merge callback response into original message

shepherd_server/server.py

@@ -65,7 +65,14 @@ async def lifespan(app: FastAPI):
 
 APP.mount("/static", StaticFiles(directory="shepherd_server/static"), name="static")
 
-FastAPIInstrumentor.instrument_app(APP, excluded_urls="docs,openapi.json")
+FastAPIInstrumentor.instrument_app(
+    APP,
+    excluded_urls="docs,openapi.json",
+    # Drop the per-ASGI-message receive/send spans.
+    # They represent individual events that are part of a larger message
+    # and can flood the OTEL backend with traces that aren't interesting.
+    exclude_spans=["receive", "send"],
+)
 
 
 @APP.get("/docs", include_in_schema=False)

workers/aragorn_lookup/worker.py

@@ -95,7 +95,7 @@ async def run_async_lookup(
     """Return an async lookup response with callback id."""
     callback_id = str(uuid.uuid4())[:8]
     with tracer.start_as_current_span("aragorn.lookup") as span:
-        span.set_attribute("callback_id", callback_id)
+        span.set_attribute("callback.id", callback_id)
         lookup_carrier = {}
         inject(lookup_carrier)
         # Put callback UID and query ID in postgres
@@ -167,7 +167,7 @@ async def aragorn_lookup(task, logger: logging.Logger):
             f"[{callback_id}] Sending lookup query to {settings.kg_retrieval_url}"
         )
         with tracer.start_as_current_span("aragorn.lookup") as span:
-            span.set_attribute("callback_id", callback_id)
+            span.set_attribute("callback.id", callback_id)
             async with httpx.AsyncClient(timeout=100) as client:
                 await client.post(
                     settings.kg_retrieval_url,

workers/aragorn_pathfinder/worker.py

@@ -56,7 +56,7 @@ async def shadowfax(task, logger: logging.Logger) -> str:
             f"""Sending pathfinder rehydration to {settings.kg_rehydrate_url}."""
         )
         with tracer.start_as_current_span("aragorn.pathfinder.rehydrate") as span:
-            span.set_attribute("query_id", query_id)
+            span.set_attribute("query.id", query_id)
             async with httpx.AsyncClient(timeout=210) as client:
                 # send a sync rehydrate query that "should" be very quick
                 rehydrated_response = await client.post(
@@ -243,7 +243,7 @@ async def shadowfax(task, logger: logging.Logger) -> str:
         f"[{callback_id}] Sending pathfinder query to {settings.kg_retrieval_url}"
     )
     with tracer.start_as_current_span("aragorn.pathfinder") as span:
-        span.set_attribute("callback_id", callback_id)
+        span.set_attribute("callback.id", callback_id)
         async with httpx.AsyncClient(timeout=100) as client:
             retriever_async_response = await client.post(
                 settings.kg_retrieval_url,

workers/bte_lookup/worker.py

@@ -95,7 +95,7 @@ async def run_async_lookup(
     """Return an async lookup response with callback id."""
     callback_id = str(uuid.uuid4())[:8]
     with tracer.start_as_current_span("bte.lookup") as span:
-        span.set_attribute("callback_id", callback_id)
+        span.set_attribute("callback.id", callback_id)
         lookup_carrier = {}
         inject(lookup_carrier)
         # Put callback UID and query ID in postgres
@@ -160,7 +160,7 @@ async def bte_lookup(task, logger: logging.Logger):
             f"[{callback_id}] Sending lookup query to {settings.kg_retrieval_url}"
         )
         with tracer.start_as_current_span("bte.lookup") as span:
-            span.set_attribute("callback_id", callback_id)
+            span.set_attribute("callback.id", callback_id)
             async with httpx.AsyncClient(timeout=100) as client:
                 await client.post(
                     settings.kg_retrieval_url,

workers/finish_query/worker.py

@@ -130,6 +130,8 @@ async def send_callback(
     payload_size = len(message_bytes)
     delivered = False
     attempts = 0
+    wait = 0.0
+    backoff = 0.0
     for attempt in range(1, CALLBACK_RETRIES + 1):
         attempts = attempt
         attempt_start = time.time()
@@ -142,6 +144,7 @@ async def send_callback(
                 )
                 response.raise_for_status()
                 elapsed = time.time() - attempt_start
+                wait += elapsed
                 logger.info(
                     f"Sent response back to {callback_url} in {elapsed:.3f}s "
                     f"({len(message_bytes)} bytes, "
@@ -151,22 +154,28 @@ async def send_callback(
                 break
         except Exception as e:
             elapsed = time.time() - attempt_start
+            wait += elapsed
             failure = (
                 f"Failed to send callback to {callback_url} after {elapsed:.3f}s "
                 f"(attempt {attempt}/{CALLBACK_RETRIES}, "
                 f"{len(message_bytes)} bytes): {_describe_callback_failure(e)}"
             )
             logger.error(failure)
             span.add_event(
-                "callback_attempt_failed",
-                {"attempt": attempt, "duration_ms": int(elapsed * 1000)},
+                "callback.attempt_failed",
+                {
+                    "callback.attempt": attempt,
+                    "callback.attempt_duration_ms": int(elapsed * 1000),
+                },
             )
             if attempt < CALLBACK_RETRIES:
                 if len(message_bytes) <= RETRY_LOG_SPLICE_MAX_BYTES:
                     message_bytes = _append_log_entry(
                         message_bytes, _log_entry(failure)
                     )
-                await asyncio.sleep(1 * (2 ** (attempt - 1)))
+                sleep_for = 1 * (2 ** (attempt - 1))
+                backoff += sleep_for
+                await asyncio.sleep(sleep_for)
 
     total = time.time() - started
     if not delivered:
@@ -183,6 +192,8 @@ async def send_callback(
     # Attributes rather than another log line: same numbers, no per-query log
     # storage, and they're queryable alongside the rest of the trace.
     span.set_attribute("callback.duration_ms", int(total * 1000))
+    span.set_attribute("callback.wait_ms", int(wait * 1000))
+    span.set_attribute("callback.backoff_ms", int(backoff * 1000))
     span.set_attribute("callback.attempts", attempts)
     span.set_attribute("callback.payload_bytes", payload_size)
     span.set_attribute("callback.delivered", delivered)

workers/merge_message/worker.py

@@ -894,8 +894,8 @@ async def process_query(task, parent_ctx, logger, limiter):
         drained = 0
         try:
             with tracer.start_as_current_span(STREAM, context=parent_ctx) as span:
-                span.set_attribute("callback_id", callback_id)
-                span.set_attribute("response_id", response_id)
+                span.set_attribute("callback.id", callback_id)
+                span.set_attribute("response.id", response_id)
 
                 # Non-blocking: never wait on the lock. The worker that holds it
                 # drains the whole query, so a loser has nothing useful to add.
@@ -963,7 +963,7 @@ async def process_query(task, parent_ctx, logger, limiter):
                     logger.error(f"[{callback_id}] Process pool broken; re-enqueuing.")
                     await remove_lock(response_id, CONSUMER, logger)
                     await _reenqueue_wake_task(task, logger)
-                    span.set_attribute("drained_callbacks", drained)
+                    span.set_attribute("merge.drained_callbacks", drained)
                     return
                 except Exception:
                     logger.error(
@@ -972,10 +972,10 @@ async def process_query(task, parent_ctx, logger, limiter):
                     )
                     await remove_lock(response_id, CONSUMER, logger)
                     await _reenqueue_wake_task(task, logger)
-                    span.set_attribute("drained_callbacks", drained)
+                    span.set_attribute("merge.drained_callbacks", drained)
                     return
 
-                span.set_attribute("drained_callbacks", drained)
+                span.set_attribute("merge.drained_callbacks", drained)
                 logger.info(
                     f"[{callback_id}] Merged {drained} callback(s) in "
                     f"{time.time() - lock_time:.2f}s"