Change recording to create a span of type "record_root" (#1703)

examples/experimental/otel_exporter.ipynb

@@ -38,7 +38,6 @@
     "root.setLevel(logging.DEBUG)\n",
     "handler = logging.StreamHandler(sys.stdout)\n",
     "handler.setLevel(logging.DEBUG)\n",
-    "handler.addFilter(logging.Filter(\"trulens\"))\n",
     "formatter = logging.Formatter(\n",
     "    \"%(asctime)s - %(name)s - %(levelname)s - %(message)s\"\n",
     ")\n",
@@ -53,10 +52,11 @@
    "outputs": [],
    "source": [
     "from trulens.experimental.otel_tracing.core.instrument import instrument\n",
+    "from trulens.otel.semconv.trace import SpanAttributes\n",
     "\n",
     "\n",
     "class TestApp:\n",
-    "    @instrument()\n",
+    "    @instrument(span_type=SpanAttributes.SpanType.MAIN)\n",
     "    def respond_to_query(self, query: str) -> str:\n",
     "        return f\"answer: {self.nested(query)}\"\n",
     "\n",
@@ -109,6 +109,7 @@
     "session = TruSession()\n",
     "session.experimental_enable_feature(\"otel_tracing\")\n",
     "session.reset_database()\n",
+    "\n",
     "init(session, debug=True)"
    ]
   },

src/core/trulens/core/app.py

@@ -3,6 +3,7 @@
 from abc import ABC
 from abc import ABCMeta
 from abc import abstractmethod
+import contextlib
 import contextvars
 import datetime
 import inspect
@@ -422,6 +423,18 @@ def tru(self) -> core_connector.DBConnector:
         pydantic.PrivateAttr(default_factory=dict)
     )
 
+    tokens: List[object] = []
+    """
+    OTEL context tokens for the current context manager. These tokens are how the OTEL
+    context api keeps track of what is changed in the context, and used to undo the changes.
+    """
+
+    span_context: Optional[contextlib.AbstractContextManager] = None
+    """
+    Span context manager. Required to help keep track of the appropriate span context
+    to enter/exit.
+    """
+
     def __init__(
         self,
         connector: Optional[core_connector.DBConnector] = None,
@@ -884,27 +897,29 @@ def __enter__(self):
         if self.session.experimental_feature(
             core_experimental.Feature.OTEL_TRACING
         ):
-            from trulens.experimental.otel_tracing.core.app import _App
+            from trulens.experimental.otel_tracing.core.instrument import (
+                App as OTELApp,
+            )
 
-            return _App.__enter__(self)
+            return OTELApp.__enter__(self)
 
         ctx = core_instruments._RecordingContext(app=self)
 
         token = self.recording_contexts.set(ctx)
         ctx.token = token
 
-        # self._set_context_vars()
-
         return ctx
 
     # For use as a context manager.
     def __exit__(self, exc_type, exc_value, exc_tb):
         if self.session.experimental_feature(
             core_experimental.Feature.OTEL_TRACING
         ):
-            from trulens.experimental.otel_tracing.core.app import _App
+            from trulens.experimental.otel_tracing.core.instrument import (
+                App as OTELApp,
+            )
 
-            return _App.__exit__(self, exc_type, exc_value, exc_tb)
+            return OTELApp.__exit__(self, exc_type, exc_value, exc_tb)
 
         ctx = self.recording_contexts.get()
         self.recording_contexts.reset(ctx.token)
@@ -921,9 +936,11 @@ async def __aenter__(self):
         if self.session.experimental_feature(
             core_experimental.Feature.OTEL_TRACING
         ):
-            from trulens.experimental.otel_tracing.core.app import _App
+            from trulens.experimental.otel_tracing.core.instrument import (
+                App as OTELApp,
+            )
 
-            return await _App.__aenter__(self)
+            return OTELApp.__enter__(self)
 
         ctx = core_instruments._RecordingContext(app=self)
 
@@ -939,9 +956,11 @@ async def __aexit__(self, exc_type, exc_value, exc_tb):
         if self.session.experimental_feature(
             core_experimental.Feature.OTEL_TRACING
         ):
-            from trulens.experimental.otel_tracing.core.app import _App
+            from trulens.experimental.otel_tracing.core.instrument import (
+                App as OTELApp,
+            )
 
-            return await _App.__aexit__(self, exc_type, exc_value, exc_tb)
+            return OTELApp.__exit__(self, exc_type, exc_value, exc_tb)
 
         ctx = self.recording_contexts.get()
         self.recording_contexts.reset(ctx.token)

src/core/trulens/experimental/otel_tracing/core/instrument.py

@@ -1,8 +1,13 @@
 from functools import wraps
 import logging
-from typing import Callable, Optional
+from typing import Any, Callable, Dict, Optional
+import uuid
 
 from opentelemetry import trace
+from opentelemetry.baggage import remove_baggage
+from opentelemetry.baggage import set_baggage
+import opentelemetry.context as context_api
+from trulens.core import app as core_app
 from trulens.experimental.otel_tracing.core.init import TRULENS_SERVICE_NAME
 from trulens.experimental.otel_tracing.core.span import Attributes
 from trulens.experimental.otel_tracing.core.span import (
@@ -11,6 +16,9 @@
 from trulens.experimental.otel_tracing.core.span import (
     set_user_defined_attributes,
 )
+from trulens.experimental.otel_tracing.core.span import (
+    set_main_span_attributes,
+)
 from trulens.otel.semconv.trace import SpanAttributes
 
 logger = logging.getLogger(__name__)
@@ -52,6 +60,14 @@ def wrapper(*args, **kwargs):
                 set_general_span_attributes(span, span_type)
                 attributes_exception = None
 
+                if span_type == SpanAttributes.SpanType.MAIN:
+                    # Only an exception in calling the function should determine whether
+                    # to set the main error. Errors in setting attributes should not be classified
+                    # as main errors.
+                    set_main_span_attributes(
+                        span, func, args, kwargs, ret, func_exception
+                    )
+
                 try:
                     set_user_defined_attributes(
                         span,
@@ -77,3 +93,65 @@ def wrapper(*args, **kwargs):
         return wrapper
 
     return inner_decorator
+
+
+class App(core_app.App):
+    # For use as a context manager.
+    def __enter__(self):
+        logger.debug("Entering the OTEL app context.")
+
+        # Note: This is not the same as the record_id in the core app since the OTEL
+        # tracing is currently separate from the old records behavior
+        otel_record_id = str(uuid.uuid4())
+
+        tracer = trace.get_tracer_provider().get_tracer(TRULENS_SERVICE_NAME)
+
+        # Calling set_baggage does not actually add the baggage to the current context, but returns a new one
+        # To avoid issues with remembering to add/remove the baggage, we attach it to the runtime context.
+        self.tokens.append(
+            context_api.attach(
+                set_baggage(SpanAttributes.RECORD_ID, otel_record_id)
+            )
+        )
+        self.tokens.append(
+            context_api.attach(set_baggage(SpanAttributes.APP_ID, self.app_id))
+        )
+
+        # Use start_as_current_span as a context manager
+        self.span_context = tracer.start_as_current_span("root")
+        root_span = self.span_context.__enter__()
+
+        # Set general span attributes
+        root_span.set_attribute("name", "root")
+        set_general_span_attributes(
+            root_span, SpanAttributes.SpanType.RECORD_ROOT
+        )
+
+        # Set record root specific attributes
+        root_span.set_attribute(
+            SpanAttributes.RECORD_ROOT.APP_NAME, self.app_name
+        )
+        root_span.set_attribute(
+            SpanAttributes.RECORD_ROOT.APP_VERSION, self.app_version
+        )
+        root_span.set_attribute(SpanAttributes.RECORD_ROOT.APP_ID, self.app_id)
+        root_span.set_attribute(
+            SpanAttributes.RECORD_ROOT.RECORD_ID, otel_record_id
+        )
+
+        return root_span
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        remove_baggage(SpanAttributes.RECORD_ID)
+        remove_baggage(SpanAttributes.APP_ID)
+
+        logger.debug("Exiting the OTEL app context.")
+
+        while self.tokens:
+            # Clearing the context once we're done with this root span.
+            # See https://github.com/open-telemetry/opentelemetry-python/issues/2432#issuecomment-1593458684
+            context_api.detach(self.tokens.pop())
+
+        if self.span_context:
+            # TODO[SNOW-1854360]: Add in feature function spans.
+            self.span_context.__exit__(exc_type, exc_value, exc_tb)

src/core/trulens/experimental/otel_tracing/core/span.py

@@ -6,6 +6,7 @@
 import logging
 from typing import Any, Callable, Dict, Optional, Union
 
+from opentelemetry.baggage import get_baggage
 from opentelemetry.trace.span import Span
 from trulens.core.utils import signature as signature_utils
 from trulens.otel.semconv.trace import SpanAttributes
@@ -87,6 +88,12 @@ def set_general_span_attributes(
 ) -> Span:
     span.set_attribute("kind", "SPAN_KIND_TRULENS")
     span.set_attribute(SpanAttributes.SPAN_TYPE, span_type)
+    span.set_attribute(
+        SpanAttributes.APP_ID, str(get_baggage(SpanAttributes.APP_ID))
+    )
+    span.set_attribute(
+        SpanAttributes.RECORD_ID, str(get_baggage(SpanAttributes.RECORD_ID))
+    )
 
     return span
 
@@ -138,3 +145,26 @@ def get_main_input(func: Callable, args: tuple, kwargs: dict) -> str:
     sig = signature(func)
     bindings = signature(func).bind(*args, **kwargs)
     return signature_utils.main_input(func, sig, bindings)
+
+
+def set_main_span_attributes(
+    span: Span,
+    /,
+    func: Callable,
+    args: tuple,
+    kwargs: dict,
+    ret: Any,
+    exception: Optional[Exception],
+) -> None:
+    span.set_attribute(
+        SpanAttributes.MAIN.MAIN_INPUT, get_main_input(func, args, kwargs)
+    )
+
+    if exception:
+        span.set_attribute(SpanAttributes.MAIN.MAIN_ERROR, str(exception))
+
+    if ret is not None:
+        span.set_attribute(
+            SpanAttributes.MAIN.MAIN_OUTPUT,
+            signature_utils.main_output(func, ret),
+        )

src/otel/semconv/trulens/otel/semconv/trace.py

@@ -56,6 +56,15 @@ class SpanAttributes:
     User-defined selector name for the current span.
     """
 
+    RECORD_ID = BASE + "record_id"
+    """ID of the record that the span belongs to."""
+
+    APP_ID = BASE + "app_id"
+    """ID of the app that the span belongs to."""
+
+    ROOT_SPAN_ID = BASE + "root_span_id"
+    """ID of the root span of the record that the span belongs to."""
+
     class SpanType(str, Enum):
         """Span type attribute values.
 
@@ -90,6 +99,9 @@ class SpanType(str, Enum):
         RECORD_ROOT = "record_root"
         """Spans as collected by tracing system."""
 
+        MAIN = "main"
+        """The main span of a record."""
+
         EVAL_ROOT = "eval_root"
         """Feedback function evaluation span."""
 
@@ -153,6 +165,25 @@ class UNKNOWN:
 
         base = "trulens.unknown"
 
+    class MAIN:
+        """Attributes for the main span of a record."""
+
+        base = "trulens.main"
+
+        SPAN_NAME_PREFIX = base + "."
+
+        MAIN_INPUT = base + ".main_input"
+        """Main input to the app."""
+
+        MAIN_OUTPUT = base + ".main_output"
+        """Main output of the app."""
+
+        MAIN_ERROR = base + ".main_error"
+        """Main error of the app.
+
+        Exclusive with main output.
+        """
+
     class RECORD_ROOT:
         """Attributes for the root span of a record.
 
@@ -182,18 +213,6 @@ class RECORD_ROOT:
         all those costs.
         """
 
-        MAIN_INPUT = base + ".main_input"
-        """Main input to the app."""
-
-        MAIN_OUTPUT = base + ".main_output"
-        """Main output of the app."""
-
-        MAIN_ERROR = base + ".main_error"
-        """Main error of the app.
-
-        Exclusive with main output.
-        """
-
     class EVAL_ROOT:
         """Attributes for the root span of a feedback evaluation.
 

tests/unit/test_otel_instrument.py

@@ -11,6 +11,7 @@
 from trulens.core.session import TruSession
 from trulens.experimental.otel_tracing.core.init import init
 from trulens.experimental.otel_tracing.core.instrument import instrument
+from trulens.otel.semconv.trace import SpanAttributes
 
 from tests.test import TruTestCase
 from tests.util.df_comparison import (
@@ -19,7 +20,7 @@
 
 
 class _TestApp:
-    @instrument()
+    @instrument(span_type=SpanAttributes.SpanType.MAIN)
     def respond_to_query(self, query: str) -> str:
         return f"answer: {self.nested(query)}"
 
@@ -124,7 +125,7 @@ def test_instrument_decorator(self) -> None:
         # Compare results to expected.
         GOLDEN_FILENAME = "tests/unit/static/golden/test_otel_instrument__test_instrument_decorator.csv"
         actual = self._get_events()
-        self.assertEqual(len(actual), 8)
+        self.assertEqual(len(actual), 10)
         self.write_golden(GOLDEN_FILENAME, actual)
         expected = self.load_golden(GOLDEN_FILENAME)
         self._convert_column_types(expected)
@@ -134,7 +135,7 @@ def test_instrument_decorator(self) -> None:
             actual,
             ignore_locators=[
                 f"df.iloc[{i}][resource_attributes][telemetry.sdk.version]"
-                for i in range(8)
+                for i in range(10)
             ],
         )