ray-llm

1
# This is a lightweight fork of Opentelemetry's Threading instrmentor which
2
# since it hasn't been merged yet, we're pulling in directly. Pulled from
3
# https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1582 on
4
# 2023-09-12 by thomas@.
5

6
# Copyright The OpenTelemetry Authors
7
#
8
# Licensed under the Apache License, Version 2.0 (the "License");
9
# you may not use this file except in compliance with the License.
10
# You may obtain a copy of the License at
11
#
12
#     http://www.apache.org/licenses/LICENSE-2.0
13
#
14
# Unless required by applicable law or agreed to in writing, software
15
# distributed under the License is distributed on an "AS IS" BASIS,
16
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17
# See the License for the specific language governing permissions and
18
# limitations under the License.
19

20
# pylint: disable=empty-docstring,no-value-for-parameter,no-member,no-name-in-module
21

22
import threading  # pylint: disable=import-self
23
from typing import Collection
24

25
from opentelemetry import context, trace
26
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
27
from opentelemetry.trace import (
28
    get_current_span,
29
    get_tracer,
30
    get_tracer_provider,
31
)
32

33

34
class _InstrumentedThread(threading.Thread):
35
    _tracer: trace.Tracer
36
    _parent_span: trace.Span
37

38
    def start(self):
39
        self._parent_span = get_current_span()
40
        super().start()
41

42
    def run(self):
43
        parent_span = self._parent_span or get_current_span()
44
        ctx = trace.set_span_in_context(parent_span)
45
        context.attach(ctx)
46
        super().run()
47

48

49
class ThreadingInstrumentor(BaseInstrumentor):  # pylint: disable=empty-docstring
50
    original_threadcls = threading.Thread
51

52
    def instrumentation_dependencies(self) -> Collection[str]:
53
        return ()
54

55
    def _instrument(self, *args, **kwargs):
56
        tracer_provider = kwargs.get("tracer_provider", None) or get_tracer_provider()
57

58
        tracer = get_tracer(__name__, "0.0.1", tracer_provider)
59
        threading.Thread = _InstrumentedThread
60
        _InstrumentedThread._tracer = tracer
61

62
    def _uninstrument(self, **kwargs):
63
        threading.Thread = self.original_threadcls
64