Bug fixing

Author: Andrei Bratu

src/humanloop/otel/exporter.py

@@ -108,8 +108,9 @@ def is_evaluated_file(
                         ),
                     )
                     logger.debug(
-                        "Span %s with EvaluationContext %s added to upload queue",
-                        span.attributes,
+                        "[HumanloopSpanExporter] Span %s %s with EvaluationContext %s added to upload queue",
+                        span.context.span_id,
+                        span.name,
                         evaluation_context_copy,
                     )
             # Reset the EvaluationContext so run eval does not
@@ -119,24 +120,24 @@ def is_evaluated_file(
                 evaluation_context,
             ):
                 logger.debug(
-                    "EvaluationContext %s marked as exhausted for Log in Span %s",
+                    "[HumanloopSpanExporter] EvaluationContext %s marked as exhausted for Log in Span %s",
                     evaluation_context,
                     spans[0].attributes,
                 )
                 # Mark the EvaluationContext as used
                 self._client.evaluation_context_variable.set(None)
             return SpanExportResult.SUCCESS
         else:
-            logger.warning("HumanloopSpanExporter is shutting down, not accepting new spans")
+            logger.warning("[HumanloopSpanExporter] Shutting down, not accepting new spans")
             return SpanExportResult.FAILURE
 
     def shutdown(self) -> None:
         self._shutdown = True
         for thread in self._threads:
             thread.join()
-            logger.debug("Exporter Thread %s joined", thread.ident)
+            logger.debug("[HumanloopSpanExporter] Exporter Thread %s joined", thread.ident)
 
-    def force_flush(self, timeout_millis: int = 3000) -> bool:
+    def force_flush(self, timeout_millis: int = 10000) -> bool:
         self._shutdown = True
         for thread in self._threads:
             thread.join(timeout=timeout_millis)
@@ -179,14 +180,20 @@ def _do_work(self):
                 # Span is not part of a Flow Log
                 self._export_span_dispatch(span_to_export)
                 logger.debug(
-                    "_do_work on Thread %s: Starting to upload span %s",
+                    "[HumanloopSpanExporter] _do_work on Thread %s: Dispatching span %s %s",
                     threading.get_ident(),
-                    span_to_export,
+                    span_to_export.context.span_id,
+                    span_to_export.name,
                 )
             elif span_to_export.parent.span_id in self._span_id_to_uploaded_log_id:
                 # Span is part of a Flow and its parent has been uploaded
                 self._export_span_dispatch(span_to_export)
-                logger.debug("_do_work on Thread %s: Starting to upload span %s", threading.get_ident(), span_to_export)
+                logger.debug(
+                    "[HumanloopSpanExporter] _do_work on Thread %s: Dispatching span %s %s",
+                    threading.get_ident(),
+                    span_to_export.context.span_id,
+                    span_to_export.name,
+                )
             else:
                 # Requeue the Span and upload after its parent
                 self._upload_queue.put((span_to_export, evaluation_context))
@@ -200,16 +207,28 @@ def _complete_flow_log(self, span_id: int) -> None:
                     # All logs in the Trace have been uploaded, mark the Flow Log as complete
                     flow_log_id = self._span_id_to_uploaded_log_id[flow_log_span_id]
                     self._client.flows.update_log(log_id=flow_log_id, trace_status="complete")
-            break
+                break
 
     def _export_span_dispatch(self, span: ReadableSpan) -> None:
         hl_file = read_from_opentelemetry_span(span, key=HUMANLOOP_FILE_KEY)
         file_type = span._attributes.get(HUMANLOOP_FILE_TYPE_KEY)  # type: ignore
         parent_span_id = span.parent.span_id if span.parent else None
 
         while parent_span_id and self._span_id_to_uploaded_log_id.get(parent_span_id) is None:
+            logger.debug(
+                "[HumanloopSpanExporter] Span %s %s waiting for parent %s to be uploaded",
+                span.context.span_id,
+                span.name,
+                parent_span_id,
+            )
             time.sleep(0.1)
 
+        logger.debug(
+            "[HumanloopSpanExporter] Exporting span %s with file type %s",
+            span,
+            file_type,
+        )
+
         if file_type == "prompt":
             export_func = self._export_prompt
         elif file_type == "tool":

src/humanloop/otel/processor.py

@@ -1,6 +1,8 @@
+from concurrent.futures import ThreadPoolExecutor
 import logging
 from collections import defaultdict
-from typing import Any
+import time
+from typing import Any, TypedDict
 
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter
@@ -23,6 +25,11 @@
 logger = logging.getLogger("humanloop.sdk")
 
 
+class CompletableSpan(TypedDict):
+    span: ReadableSpan
+    complete: bool
+
+
 class HumanloopSpanProcessor(SimpleSpanProcessor):
     """Enrich Humanloop spans with data from their children spans.
 
@@ -42,43 +49,93 @@ class HumanloopSpanProcessor(SimpleSpanProcessor):
 
     def __init__(self, exporter: SpanExporter) -> None:
         super().__init__(exporter)
-        # Span parent to Span children map
-        self._children: dict[int, list] = defaultdict(list)
+        # span parent to span children map
+        self._children: dict[int, list[CompletableSpan]] = defaultdict(list)
+        # List of all span IDs that are contained in a Flow trace
+        # They are passed to the Exporter as a span attribute
+        # so the Exporter knows when to complete a trace
         self._prerequisites: dict[int, list[int]] = {}
+        self._executor = ThreadPoolExecutor(max_workers=4)
+
+    def shutdown(self):
+        self._executor.shutdown()
+        return super().shutdown()
 
     def on_start(self, span, parent_context=None):
         span_id = span.context.span_id
+        parent_span_id = span.parent.span_id if span.parent else None
         if span.name == "humanloop.flow":
             self._prerequisites[span_id] = []
-        if span.parent and is_humanloop_span(span):
-            parent_span_id = span.parent.span_id
+        if parent_span_id and is_humanloop_span(span):
             for trace_head, all_trace_nodes in self._prerequisites.items():
                 if parent_span_id == trace_head or parent_span_id in all_trace_nodes:
                     all_trace_nodes.append(span_id)
+                    break
+        # Handle stream case: when Prompt instrumented function calls a provider with streaming: true
+        # The instrumentor span will end only when the ChunksResponse is consumed, which can happen
+        # after the span created by the Prompt utility finishes. To handle this, we register all instrumentor
+        # spans belonging to a Humanloop span, and their parent will wait for them to complete in onEnd before
+        # exporting the Humanloop span.
+        if parent_span_id and _is_instrumentor_span(span):
+            if parent_span_id not in self._children:
+                self._children[parent_span_id] = []
+            self._children[parent_span_id].append(
+                {
+                    "span": span,
+                    "complete": False,
+                }
+            )
 
     def on_end(self, span: ReadableSpan) -> None:
         if is_humanloop_span(span=span):
-            _process_span_dispatch(span, self._children[span.context.span_id])
-            # Release the reference to the Spans as they've already
-            # been sent to the Exporter
-            del self._children[span.context.span_id]
+            # Wait for children to complete asynchronously
+            self._executor.submit(self._wait_for_children, span=span)
+        elif span.parent is not None and _is_instrumentor_span(span):
+            # If this is one of the children spans waited upon, update its completion status
+
+            # Updating the child span status
+            self._children[span.parent.span_id] = [
+                child if child["span"].context.span_id != span.context.span_id else {"span": span, "complete": True}
+                for child in self._children[span.parent.span_id]
+            ]
+
+            # Export the instrumentor span
+            self.span_exporter.export([span])
         else:
-            if span.parent is not None and _is_instrumentor_span(span):
-                # Copy the Span and keep it until the Humanloop Span
-                # arrives in order to enrich it
-                self._children[span.parent.span_id].append(span)
-        # Pass the Span to the Exporter
+            # Unknown span, pass it to the Exporter
+            self.span_exporter.export([span])
+
+    def _wait_for_children(self, span: ReadableSpan):
+        """Wait for all children spans to complete before processing the Humanloop span."""
+        span_id = span.context.span_id
+        while not all(child["complete"] for child in self._children[span_id]):
+            logger.debug(
+                "[HumanloopSpanProcessor] Span %s %s waiting for children to complete: %s",
+                span_id,
+                span.name,
+                self._children[span_id],
+            )
+            time.sleep(0.1)
+        # All instrumentor spans have arrived, we can process the
+        # Humanloop parent span owning them
         if span.name == "humanloop.flow":
             write_to_opentelemetry_span(
                 span=span,
                 key=HUMANLOOP_FLOW_PREREQUISITES_KEY,
-                value=self._prerequisites[span.context.span_id],
+                value=self._prerequisites[span_id],
             )
+            del self._prerequisites[span_id]
+        logger.debug("[HumanloopSpanProcessor] Dispatching span %s %s", span_id, span.name)
+        _process_span_dispatch(span, [child["span"] for child in self._children[span_id]])
+        # Release references
+        del self._children[span_id]
+        # Pass Humanloop span to Exporter
+        logger.debug("[HumanloopSpanProcessor] Sending span %s %s to exporter", span_id, span.name)
         self.span_exporter.export([span])
 
 
 def _is_instrumentor_span(span: ReadableSpan) -> bool:
-    """Determine if the Span contains information of interest for Spans created by Humanloop decorators."""
+    """Determine if the span contains information of interest for Spans created by Humanloop decorators."""
     # At the moment we only enrich Spans created by the Prompt decorators
     # As we add Instrumentors for other libraries, this function must
     # be expanded
@@ -104,7 +161,11 @@ def _process_span_dispatch(span: ReadableSpan, children_spans: list[ReadableSpan
     elif file_type == "flow":
         pass
     else:
-        logger.error("Unknown Humanloop File Span %s", span)
+        logger.error(
+            "[HumanloopSpanProcessor] Unknown Humanloop File span %s %s",
+            span.context.span_id,
+            span.name,
+        )
 
 
 def _process_prompt(prompt_span: ReadableSpan, children_spans: list[ReadableSpan]):
@@ -150,9 +211,14 @@ def _enrich_prompt_kernel(prompt_span: ReadableSpan, llm_provider_call_span: Rea
         # Validate the Prompt Kernel
         PromptKernelRequest.model_validate(obj=prompt)
     except PydanticValidationError as e:
-        logger.error("Could not validate Prompt Kernel extracted from Span: %s", e)
-
-    # Write the enriched Prompt Kernel back to the Span
+        logger.error(
+            "[HumanloopSpanProcessor] Could not validate Prompt Kernel extracted from span: %s %s. Error: %s",
+            prompt_span.context.span_id,
+            prompt_span.name,
+            e,
+        )
+
+    # Write the enriched Prompt Kernel back to the span
     hl_file["prompt"] = prompt
     write_to_opentelemetry_span(
         span=prompt_span,