fix: remove attaching and detaching of opentelemetry contexts (#1112)

Arize-ai · Nov 12, 2024 · 0f16ffb · 0f16ffb
1 parent e782f7b
commit 0f16ffb
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 20 deletions.
diff --git a/...nce-instrumentation-llama-index/src/openinference/instrumentation/llama_index/_handler.py b/...nce-instrumentation-llama-index/src/openinference/instrumentation/llama_index/_handler.py
@@ -29,7 +29,7 @@
 )
 
 from opentelemetry import context as context_api
-from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY, attach, detach
+from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY
 from opentelemetry.trace import Span, Status, StatusCode, Tracer, set_span_in_context
 from opentelemetry.util.types import AttributeValue
 from pydantic import BaseModel as PydanticBaseModel
@@ -694,13 +694,11 @@ def _sweep(self, q: "SimpleQueue[Optional[_QueueItem]]") -> None:
 
 
 class _SpanHandler(BaseSpanHandler[_Span], extra="allow"):
-    _context_tokens: Dict[str, object] = PrivateAttr()
     _otel_tracer: Tracer = PrivateAttr()
     _export_queue: _ExportQueue = PrivateAttr()
 
     def __init__(self, tracer: Tracer) -> None:
         super().__init__()
-        self._context_tokens: Dict[str, object] = {}
         self._otel_tracer = tracer
         self._export_queue = _ExportQueue()
 
@@ -723,8 +721,6 @@ def new_span(
             attributes=dict(get_attributes_from_context()),
             context=(parent.context if parent else None),
         )
-        with self.lock:
-            self._context_tokens[id_] = attach(set_span_in_context(otel_span))
         span = _Span(
             otel_span=otel_span,
             span_kind=_init_span_kind(instance),
@@ -748,9 +744,6 @@ def prepare_to_exit_span(
             return None
         with self.lock:
             span = self.open_spans.get(id_)
-            token = self._context_tokens.get(id_)
-        if token:
-            detach(token)
         if span:
             if isinstance(instance, (BaseLLM, MultiModalLLM)) and (
                 isinstance(result, Generator)
@@ -779,9 +772,6 @@ def prepare_to_drop_span(
             return None
         with self.lock:
             span = self.open_spans.get(id_)
-            token = self._context_tokens.get(id_)
-        if token:
-            detach(token)
         if span:
             if err and isinstance(err, WorkflowDone):
                 span.end()

diff --git a/...strumentation-llama-index/tests/openinference/instrumentation/llama_index/test_handler.py b/...strumentation-llama-index/tests/openinference/instrumentation/llama_index/test_handler.py
@@ -41,7 +41,6 @@
 
 from openinference.instrumentation import using_attributes
 from openinference.instrumentation.llama_index import LlamaIndexInstrumentor
-from openinference.instrumentation.openai import OpenAIInstrumentor
 from openinference.semconv.trace import (
     DocumentAttributes,
     EmbeddingAttributes,
@@ -154,12 +153,7 @@ def main() -> None:
     for span in spans:
         traces[span.context.trace_id][span.name] = span
 
-    if is_stream:
-        # OpenAIInstrumentor is on a separate trace because no span
-        # is open when the stream iteration starts.
-        assert len(traces) == n * 2
-    else:
-        assert len(traces) == n
+    assert len(traces) == n
     for spans_by_name in traces.values():
         spans_by_id = _spans_by_id(spans_by_name.values())
         if is_stream and len(spans_by_name) == 1:
@@ -461,9 +455,7 @@ def instrument(
     in_memory_span_exporter: InMemorySpanExporter,
 ) -> Generator[None, None, None]:
     LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)
-    OpenAIInstrumentor().instrument(tracer_provider=tracer_provider)
     yield
-    OpenAIInstrumentor().uninstrument()
     LlamaIndexInstrumentor().uninstrument()
     in_memory_span_exporter.clear()