phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-17 17:52:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

tracing for APIs

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-11-26 14:07:48 -08:00
parent c2a4850a79
commit af8a1fe5b3
8 changed files with 126 additions and 63 deletions
Download patch file Download diff file Expand all files Collapse all files

12
llama_stack/providers/utils/telemetry/tracing.py
View file

@ -204,6 +204,7 @@ class SpanContextManager:
def __init__(self, name: str, attributes: Dict[str, Any] = None):
self.name = name
self.attributes = attributes
self.span = None
def __enter__(self):
global CURRENT_TRACE_CONTEXT
@ -225,10 +226,17 @@ class SpanContextManager:
self.span.attributes[key] = value
async def __aenter__(self):
return self.__enter__()
global CURRENT_TRACE_CONTEXT
context = CURRENT_TRACE_CONTEXT
if context:
self.span = context.push_span(self.name, self.attributes)
return self
async def __aexit__(self, exc_type, exc_value, traceback):
self.__exit__(exc_type, exc_value, traceback)
global CURRENT_TRACE_CONTEXT
context = CURRENT_TRACE_CONTEXT
if context:
context.pop_span()
def __call__(self, func: Callable):
@wraps(func)