ray-llm

1
import logging
2
import os
3
import time
4
from contextlib import contextmanager
5
from typing import Dict, Optional
6

7
from opentelemetry import trace
8
from opentelemetry.util.types import AttributeValue
9

10
from rayllm.backend.observability.tracing.baggage import baggage as set_baggage
11

12
# Standard loggers
13
serve_logger = logging.getLogger("ray.serve")
14
aviary_logger = logging.getLogger("aviary")
15

16
# If the AVIARY_DEBUG flag is set, then set the log level to debug
17
if os.getenv("AVIARY_DEBUG"):
18
    aviary_logger.setLevel(logging.DEBUG)
19

20
# Standard tracers
21
tracer = trace.get_tracer(__name__)
22

23

24
@contextmanager
25
def step(
26
    step_name: str,
27
    request_id: Optional[str] = None,
28
    attrs: Optional[Dict[str, AttributeValue]] = None,
29
    baggage: Optional[Dict[str, AttributeValue]] = None,
30
):
31
    if baggage is None:
32
        baggage = {}
33

34
    # TODO(tchordia): Add tracing here
35
    t = time.monotonic()
36
    try:
37
        aviary_logger.debug(f"Starting {step_name} at {t}. Id: {request_id}")
38
        with set_baggage(baggage), tracer.start_as_current_span(
39
            step_name, attributes=attrs
40
        ) as span:
41
            yield span
42
    finally:
43
        took = time.monotonic() - t
44
        aviary_logger.debug(f"Completed {step_name}, took {took}s. Id: {request_id}")
45