ray-llm

1
import concurrent.futures.thread
2
from functools import wraps
3
from typing import Collection
4

5
from opentelemetry import trace
6
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
7
from opentelemetry.trace import (
8
    get_current_span,
9
)
10

11
from rayllm.backend.observability.tracing.context import use_context
12

13

14
class _InstrumentedThreadPoolExecutorWorkItem(concurrent.futures.thread._WorkItem):
15
    def __init__(self, future, fn, args, kwargs):
16
        super().__init__(
17
            future,
18
            self.wrap_with_context(fn, parent_span=get_current_span()),
19
            args,
20
            kwargs,
21
        )
22

23
    @staticmethod
24
    def wrap_with_context(fn, parent_span: trace.Span):
25
        @wraps(fn)
26
        def wrapper(*args, **kwargs):
27
            with use_context(trace.set_span_in_context(parent_span)):
28
                return fn(*args, **kwargs)
29

30
        return wrapper
31

32

33
class ThreadPoolExecutorInstrumentor(BaseInstrumentor):
34
    original_workitem = concurrent.futures.thread._WorkItem
35

36
    def instrumentation_dependencies(self) -> Collection[str]:
37
        return ()
38

39
    def _instrument(self, *args, **kwargs):
40
        concurrent.futures.thread._WorkItem = _InstrumentedThreadPoolExecutorWorkItem
41

42
    def _uninstrument(self, **kwargs):
43
        concurrent.futures.thread._WorkItem = self.original_workitem
44