ray-llm

1
import os
2
import secrets
3
import socket
4

5
from opentelemetry import (
6
    trace,
7
)
8
from opentelemetry.instrumentation.aiohttp_client import AioHttpClientInstrumentor
9
from opentelemetry.instrumentation.botocore import BotocoreInstrumentor
10
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
11
from opentelemetry.instrumentation.redis import RedisInstrumentor
12
from opentelemetry.sdk.resources import Resource
13
from opentelemetry.sdk.trace import TracerProvider
14
from opentelemetry.util._once import Once
15

16
from rayllm.backend.observability.tracing.baggage_span_processor import (
17
    BaggageSpanProcessor,
18
)
19
from rayllm.backend.observability.tracing.fastapi import FastAPIInstrumentor
20
from rayllm.backend.observability.tracing.threading_propagator import (
21
    ThreadPoolExecutorInstrumentor,
22
)
23

24
has_setup_tracing = Once()
25

26

27
def _setup_tracing():
28
    tracer_provider = TracerProvider(
29
        resource=Resource(
30
            attributes={
31
                "service.name": "aviary_endpoints",
32
                "meta.local_hostname": socket.gethostname(),
33
                "meta.process_id": os.getpid(),
34
                "meta.process_unique_id": secrets.token_urlsafe(),
35
            }
36
        ),
37
    )
38

39
    tracer_provider.add_span_processor(  # type: ignore
40
        # Ensure we set baggage entries as attributes on all spans
41
        BaggageSpanProcessor()
42
    )
43

44
    trace.set_tracer_provider(tracer_provider)
45

46
    # Not really effective right now, still need to `instrument_app` each app
47
    FastAPIInstrumentor().instrument()
48

49
    AioHttpClientInstrumentor().instrument()
50
    BotocoreInstrumentor().instrument()
51
    HTTPXClientInstrumentor().instrument()
52
    ThreadPoolExecutorInstrumentor().instrument()
53
    RedisInstrumentor().instrument()
54

55
    tracer = trace.get_tracer(__name__)
56
    with tracer.start_as_current_span("startup"):
57
        pass
58

59

60
def setup_tracing() -> None:
61
    has_setup_tracing.do_once(_setup_tracing)
62