lmnr 0.4.53.dev0__py3-none-any.whl → 0.7.26__py3-none-any.whl

lmnr/opentelemetry_lib/opentelemetry/instrumentation/threading/__init__.py

@@ -0,0 +1,197 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Instrument threading to propagate OpenTelemetry context.
+
+Copied from opentelemetry-instrumentation-threading at commit:
+ad2fe813abb2ab0b6e25bedeebef5041ca3189f7
+https://github.com/open-telemetry/opentelemetry-python-contrib/blob/ad2fe813abb2ab0b6e25bedeebef5041ca3189f7/instrumentation/opentelemetry-instrumentation-threading/src/opentelemetry/instrumentation/threading/__init__.py
+
+Modified to use the Laminar isolated context.
+
+Usage
+-----
+
+.. code-block:: python
+
+    from opentelemetry.instrumentation.threading import ThreadingInstrumentor
+
+    ThreadingInstrumentor().instrument()
+
+This library provides instrumentation for the `threading` module to ensure that
+the OpenTelemetry context is propagated across threads. It is important to note
+that this instrumentation does not produce any telemetry data on its own. It
+merely ensures that the context is correctly propagated when threads are used.
+
+
+When instrumented, new threads created using threading.Thread, threading.Timer,
+or within futures.ThreadPoolExecutor will have the current OpenTelemetry
+context attached, and this context will be re-activated in the thread's
+run method or the executor's worker thread."
+"""
+
+from __future__ import annotations
+
+import threading
+from concurrent import futures
+from typing import TYPE_CHECKING, Any, Callable, Collection
+
+from wrapt import (
+    wrap_function_wrapper,  # type: ignore[reportUnknownVariableType]
+)
+
+from lmnr.opentelemetry_lib.tracing.context import (
+    get_current_context,
+    attach_context,
+    detach_context,
+)
+from opentelemetry import context
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import unwrap
+
+_instruments = ()
+
+if TYPE_CHECKING:
+    from typing import Protocol, TypeVar
+
+    R = TypeVar("R")
+
+    class HasOtelContext(Protocol):
+        _otel_context: context.Context
+
+
+class ThreadingInstrumentor(BaseInstrumentor):
+    __WRAPPER_START_METHOD = "start"
+    __WRAPPER_RUN_METHOD = "run"
+    __WRAPPER_SUBMIT_METHOD = "submit"
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs: Any):
+        self._instrument_thread()
+        self._instrument_timer()
+        self._instrument_thread_pool()
+
+    def _uninstrument(self, **kwargs: Any):
+        self._uninstrument_thread()
+        self._uninstrument_timer()
+        self._uninstrument_thread_pool()
+
+    @staticmethod
+    def _instrument_thread():
+        wrap_function_wrapper(
+            threading.Thread,
+            ThreadingInstrumentor.__WRAPPER_START_METHOD,
+            ThreadingInstrumentor.__wrap_threading_start,
+        )
+        wrap_function_wrapper(
+            threading.Thread,
+            ThreadingInstrumentor.__WRAPPER_RUN_METHOD,
+            ThreadingInstrumentor.__wrap_threading_run,
+        )
+
+    @staticmethod
+    def _instrument_timer():
+        wrap_function_wrapper(
+            threading.Timer,
+            ThreadingInstrumentor.__WRAPPER_START_METHOD,
+            ThreadingInstrumentor.__wrap_threading_start,
+        )
+        wrap_function_wrapper(
+            threading.Timer,
+            ThreadingInstrumentor.__WRAPPER_RUN_METHOD,
+            ThreadingInstrumentor.__wrap_threading_run,
+        )
+
+    @staticmethod
+    def _instrument_thread_pool():
+        wrap_function_wrapper(
+            futures.ThreadPoolExecutor,
+            ThreadingInstrumentor.__WRAPPER_SUBMIT_METHOD,
+            ThreadingInstrumentor.__wrap_thread_pool_submit,
+        )
+
+    @staticmethod
+    def _uninstrument_thread():
+        unwrap(threading.Thread, ThreadingInstrumentor.__WRAPPER_START_METHOD)
+        unwrap(threading.Thread, ThreadingInstrumentor.__WRAPPER_RUN_METHOD)
+
+    @staticmethod
+    def _uninstrument_timer():
+        unwrap(threading.Timer, ThreadingInstrumentor.__WRAPPER_START_METHOD)
+        unwrap(threading.Timer, ThreadingInstrumentor.__WRAPPER_RUN_METHOD)
+
+    @staticmethod
+    def _uninstrument_thread_pool():
+        unwrap(
+            futures.ThreadPoolExecutor,
+            ThreadingInstrumentor.__WRAPPER_SUBMIT_METHOD,
+        )
+
+    @staticmethod
+    def __wrap_threading_start(
+        call_wrapped: Callable[[], None],
+        instance: HasOtelContext,
+        args: tuple[()],
+        kwargs: dict[str, Any],
+    ) -> None:
+        instance._lmnr_otel_context = get_current_context()
+        return call_wrapped(*args, **kwargs)
+
+    @staticmethod
+    def __wrap_threading_run(
+        call_wrapped: Callable[..., R],
+        instance: HasOtelContext,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> R:
+        token = None
+        try:
+            # Genearally, this must be set in __wrap_threading_start, but it is
+            # possible to Thread().run() without Thread().start(), so in that case,
+            # we need to capture the context here.
+            # We still want to capture the context in __wrap_threading_start,
+            # in order to stay close to the original implementation.
+            if not hasattr(instance, "_lmnr_otel_context"):
+                instance._lmnr_otel_context = get_current_context()
+            token = attach_context(instance._lmnr_otel_context)
+            return call_wrapped(*args, **kwargs)
+        finally:
+            if token is not None:
+                detach_context(token)
+
+    @staticmethod
+    def __wrap_thread_pool_submit(
+        call_wrapped: Callable[..., R],
+        instance: futures.ThreadPoolExecutor,
+        args: tuple[Callable[..., Any], ...],
+        kwargs: dict[str, Any],
+    ) -> R:
+        # obtain the original function and wrapped kwargs
+        original_func = args[0]
+        otel_context = get_current_context()
+
+        def wrapped_func(*func_args: Any, **func_kwargs: Any) -> R:
+            token = None
+            try:
+                token = attach_context(otel_context)
+                return original_func(*func_args, **func_kwargs)
+            finally:
+                if token is not None:
+                    detach_context(token)
+
+        # replace the original function with the wrapped function
+        new_args: tuple[Callable[..., Any], ...] = (wrapped_func,) + args[1:]
+        return call_wrapped(*new_args, **kwargs)

lmnr/opentelemetry_lib/tracing/__init__.py

@@ -0,0 +1,263 @@
+import atexit
+import logging
+import threading
+
+from lmnr.opentelemetry_lib.tracing.processor import LaminarSpanProcessor
+from lmnr.sdk.client.asynchronous.async_client import AsyncLaminarClient
+from lmnr.sdk.types import SessionRecordingOptions
+from lmnr.sdk.log import VerboseColorfulFormatter
+from lmnr.opentelemetry_lib.tracing.instruments import (
+    Instruments,
+    init_instrumentations,
+)
+from lmnr.opentelemetry_lib.tracing.context import (
+    attach_context,
+    clear_context,
+    pop_span_context as ctx_pop_span_context,
+    get_current_context,
+    get_token_stack,
+    push_span_context as ctx_push_span_context,
+    set_token_stack,
+)
+
+from opentelemetry import trace
+from opentelemetry.context import Context
+
+# instead of importing from opentelemetry.instrumentation.threading,
+# we import from our modified copy to use Laminar's isolated context.
+from ..opentelemetry.instrumentation.threading import ThreadingInstrumentor
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider, SpanProcessor
+from opentelemetry.sdk.trace.export import SpanExporter
+
+TRACER_NAME = "lmnr.tracer"
+
+MAX_EVENTS_OR_ATTRIBUTES_PER_SPAN = 5000
+
+
+class TracerWrapper(object):
+    resource_attributes: dict = {}
+    enable_content_tracing: bool = True
+    session_recording_options: SessionRecordingOptions = {}
+    _lock = threading.Lock()
+    _tracer_provider: TracerProvider | None = None
+    _logger: logging.Logger
+    _async_client: AsyncLaminarClient
+    _resource: Resource
+    _span_processor: SpanProcessor
+    _original_thread_init = None
+
+    def __new__(
+        cls,
+        disable_batch=False,
+        exporter: SpanExporter | None = None,
+        instruments: set[Instruments] | None = None,
+        block_instruments: set[Instruments] | None = None,
+        base_url: str | None = None,
+        port: int = 8443,
+        http_port: int = 443,
+        project_api_key: str | None = None,
+        max_export_batch_size: int | None = None,
+        force_http: bool = False,
+        timeout_seconds: int = 30,
+        set_global_tracer_provider: bool = True,
+        otel_logger_level: int = logging.ERROR,
+        session_recording_options: SessionRecordingOptions | None = None,
+    ) -> "TracerWrapper":
+        # Silence some opentelemetry warnings
+        logging.getLogger("opentelemetry.trace").setLevel(otel_logger_level)
+
+        base_http_url = f"{base_url}:{http_port}" if base_url else None
+        with cls._lock:
+            if not hasattr(cls, "instance"):
+                cls._initialize_logger(cls)
+                obj = super(TracerWrapper, cls).__new__(cls)
+
+                # Store session recording options
+                cls.session_recording_options = session_recording_options or {}
+
+                if project_api_key:
+                    obj._async_client = AsyncLaminarClient(
+                        base_url=base_http_url or "https://api.lmnr.ai",
+                        project_api_key=project_api_key,
+                    )
+                else:
+                    obj._async_client = None
+
+                obj._resource = Resource(attributes=TracerWrapper.resource_attributes)
+
+                obj._span_processor = LaminarSpanProcessor(
+                    base_url=base_url,
+                    api_key=project_api_key,
+                    port=http_port if force_http else port,
+                    exporter=exporter,
+                    max_export_batch_size=max_export_batch_size,
+                    timeout_seconds=timeout_seconds,
+                    force_http=force_http,
+                    disable_batch=disable_batch,
+                )
+
+                lmnr_provider = TracerProvider(resource=obj._resource)
+                global_provider = trace.get_tracer_provider()
+                if set_global_tracer_provider and isinstance(
+                    global_provider, trace.ProxyTracerProvider
+                ):
+                    trace.set_tracer_provider(lmnr_provider)
+
+                obj._tracer_provider = lmnr_provider
+
+                obj._tracer_provider.add_span_processor(obj._span_processor)
+
+                # Setup threading context inheritance
+                obj._setup_threading_inheritance()
+
+                # This is not a real instrumentation and does not generate telemetry
+                # data, but it is required to ensure that OpenTelemetry context
+                # propagation is enabled.
+                # See the README at:
+                # https://pypi.org/project/opentelemetry-instrumentation-threading/
+                ThreadingInstrumentor().instrument()
+
+                init_instrumentations(
+                    tracer_provider=obj._tracer_provider,
+                    instruments=instruments,
+                    block_instruments=block_instruments,
+                    async_client=obj._async_client,
+                )
+
+                cls.instance = obj
+
+                # Force flushes for debug environments (e.g. local development)
+                atexit.register(obj.exit_handler)
+
+            return cls.instance
+
+    def _setup_threading_inheritance(self):
+        """Setup threading inheritance for isolated context."""
+        if TracerWrapper._original_thread_init is None:
+            # Monkey patch Thread.__init__ to capture context inheritance
+            TracerWrapper._original_thread_init = threading.Thread.__init__
+
+            def patched_thread_init(thread_self, *args, **kwargs):
+                # Capture current isolated context and token stack for inheritance
+                current_context = get_current_context()
+                current_token_stack = get_token_stack().copy()
+
+                # Get the original target function
+                original_target = kwargs.get("target")
+                if not original_target and args:
+                    original_target = args[0]
+
+                # Only inherit if we have a target function
+                if original_target:
+                    # Create a wrapper function that sets up context
+                    def thread_wrapper(*target_args, **target_kwargs):
+                        # Set inherited context and token stack in the new thread
+                        attach_context(current_context)
+                        set_token_stack(current_token_stack)
+                        # Run original target
+                        return original_target(*target_args, **target_kwargs)
+
+                    # Replace the target with our wrapper
+                    if "target" in kwargs:
+                        kwargs["target"] = thread_wrapper
+                    elif args:
+                        args = (thread_wrapper,) + args[1:]
+
+                # Call original init
+                TracerWrapper._original_thread_init(thread_self, *args, **kwargs)
+
+            threading.Thread.__init__ = patched_thread_init
+
+    def exit_handler(self):
+        if isinstance(self._span_processor, LaminarSpanProcessor):
+            self._span_processor.clear()
+        self.flush()
+
+    def _initialize_logger(self):
+        self._logger = logging.getLogger(__name__)
+        console_log_handler = logging.StreamHandler()
+        console_log_handler.setFormatter(VerboseColorfulFormatter())
+        self._logger.addHandler(console_log_handler)
+
+    def get_isolated_context(self) -> Context:
+        """Get the current isolated context."""
+        return get_current_context()
+
+    def push_span_context(self, span: trace.Span) -> Context:
+        """Push a new context with the given span onto the stack."""
+        current_ctx = get_current_context()
+        new_context = trace.set_span_in_context(span, current_ctx)
+        ctx_push_span_context(new_context)
+        return new_context
+
+    def pop_span_context(self) -> None:
+        """Pop the current span context from the stack."""
+        ctx_pop_span_context()
+
+    @staticmethod
+    def set_static_params(
+        resource_attributes: dict,
+        enable_content_tracing: bool,
+    ) -> None:
+        TracerWrapper.resource_attributes = resource_attributes
+        TracerWrapper.enable_content_tracing = enable_content_tracing
+
+    @classmethod
+    def verify_initialized(cls) -> bool:
+        # This is not using lock, but it is fine to return False from here
+        # even if initialization is going on.
+
+        # If we try to acquire the lock here, it may deadlock if an automatic
+        # instrumentation is importing a file that (at the top level) has a
+        # function annotated with Laminar's `observe` decorator.
+        # The decorator is evaluated at import time, inside `init_instrumentations`,
+        # which is called by the `TracerWrapper` constructor while holding the lock.
+        # Without the lock here, we will simply return False, which will cause
+        # the decorator to return the original function. This is fine, at runtime,
+        # the next import statement will re-evaluate the decorator, and Laminar will
+        # have been initialized by that time.
+        return hasattr(cls, "instance") and hasattr(cls.instance, "_span_processor")
+
+    @classmethod
+    def clear(cls):
+        if not cls.verify_initialized():
+            return
+        # Any state cleanup. Now used in between tests
+        if isinstance(cls.instance._span_processor, LaminarSpanProcessor):
+            cls.instance._span_processor.clear()
+        # Clear the isolated context state for clean test state
+        clear_context()
+
+    def shutdown(self):
+        if self._tracer_provider is None:
+            return
+        self._tracer_provider.shutdown()
+
+    def flush(self):
+        if not hasattr(self, "_span_processor"):
+            self._logger.warning("TracerWrapper not fully initialized, cannot flush")
+            return False
+        return self._span_processor.force_flush()
+
+    def force_reinit_processor(self):
+        if isinstance(self._span_processor, LaminarSpanProcessor):
+            self._span_processor.force_flush()
+            self._span_processor.force_reinit()
+            # Clear the isolated context to prevent subsequent invocations
+            # (e.g., in Lambda) from continuing traces from previous invocations
+            clear_context()
+        else:
+            self._logger.warning(
+                "Not using LaminarSpanProcessor, cannot force reinit processor"
+            )
+
+    @classmethod
+    def get_session_recording_options(cls) -> SessionRecordingOptions:
+        """Get the session recording options set during initialization."""
+        return cls.session_recording_options
+
+    def get_tracer(self) -> trace.Tracer:
+        if self._tracer_provider is None:
+            return trace.get_tracer_provider().get_tracer(TRACER_NAME)
+        return self._tracer_provider.get_tracer(TRACER_NAME)