Skip to content
Snippets Groups Projects
Unverified Commit 893c4057 authored by Andrei Fajardo's avatar Andrei Fajardo Committed by GitHub
Browse files

Thread-safe instrumentation (#12638)

parent 2ec36fcd
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
llama-index-core/llama_index/core/instrumentation/dispatcher.py
+ 136
− 104
View file @ 893c4057
from typing import Any, List, Optional, Dict, Protocol from typing import Any, List, Optional, Dict, Protocol
from functools import partial from functools import partial
from contextlib import contextmanager from collections import defaultdict
import asyncio import asyncio
import inspect import inspect
import threading
import uuid import uuid
from llama_index.core.bridge.pydantic import BaseModel, Field, PrivateAttr from llama_index.core.bridge.pydantic import BaseModel, Field, PrivateAttr
from llama_index.core.instrumentation.events import BaseEvent from llama_index.core.instrumentation.events import BaseEvent
...@@ -17,7 +18,11 @@ from contextvars import ContextVar ...@@ -17,7 +18,11 @@ from contextvars import ContextVar
import wrapt import wrapt
span_ctx = ContextVar("span_ctx", default={}) # ContextVar's for managing active spans
span_ctx_var = ContextVar(
"span_ctx_var", default=defaultdict(dict)
) # per thread >> async-task
DEFAULT_SYNC_KEY = "sync_tasks"
class EventDispatcher(Protocol): class EventDispatcher(Protocol):
...@@ -25,14 +30,21 @@ class EventDispatcher(Protocol): ...@@ -25,14 +30,21 @@ class EventDispatcher(Protocol):
... ...
class EventContext(BaseModel): class Dispatcher(BaseModel):
span_id: str = Field(default="") """Dispatcher class.
event_context = EventContext() Responsible for dispatching BaseEvent (and its subclasses) as well as
sending signals to enter/exit/drop a BaseSpan. It does so by sending
event and span signals to its attached BaseEventHandler as well as
BaseSpanHandler.
Concurrency:
- Dispatcher is async-task and thread safe in the sense that
spans of async coros will maintain its hieararchy or trace-trees and
spans which emanate from various threads will also maintain its
hierarchy.
"""
class Dispatcher(BaseModel):
name: str = Field(default_factory=str, description="Name of dispatcher") name: str = Field(default_factory=str, description="Name of dispatcher")
event_handlers: List[BaseEventHandler] = Field( event_handlers: List[BaseEventHandler] = Field(
default=[], description="List of attached handlers" default=[], description="List of attached handlers"
...@@ -51,10 +63,12 @@ class Dispatcher(BaseModel): ...@@ -51,10 +63,12 @@ class Dispatcher(BaseModel):
default=True, default=True,
description="Whether to propagate the event to parent dispatchers and their handlers", description="Whether to propagate the event to parent dispatchers and their handlers",
) )
current_span_id: Optional[str] = Field( current_span_ids: Optional[Dict[Any, str]] = Field(
default=None, description="Id of current span." default_factory=dict,
description="Id of current enclosing span. Used for creating `dispatch_event` partials.",
) )
_asyncio_lock: Optional[asyncio.Lock] = PrivateAttr() _asyncio_lock: Optional[asyncio.Lock] = PrivateAttr()
_lock: Optional[threading.Lock] = PrivateAttr()
def __init__( def __init__(
self, self,
...@@ -67,6 +81,7 @@ class Dispatcher(BaseModel): ...@@ -67,6 +81,7 @@ class Dispatcher(BaseModel):
propagate: bool = True, propagate: bool = True,
): ):
self._asyncio_lock = None self._asyncio_lock = None
self._lock = None
super().__init__( super().__init__(
name=name, name=name,
event_handlers=event_handlers, event_handlers=event_handlers,
...@@ -77,12 +92,29 @@ class Dispatcher(BaseModel): ...@@ -77,12 +92,29 @@ class Dispatcher(BaseModel):
propagate=propagate, propagate=propagate,
) )
@property
def current_span_id(self) -> Optional[str]:
current_thread = threading.get_ident()
if current_thread in self.current_span_ids:
return self.current_span_ids[current_thread]
return None
def set_current_span_id(self, value: str):
current_thread = threading.get_ident()
self.current_span_ids[current_thread] = value
@property @property
def asyncio_lock(self) -> asyncio.Lock: def asyncio_lock(self) -> asyncio.Lock:
if self._asyncio_lock is None: if self._asyncio_lock is None:
self._asyncio_lock = asyncio.Lock() self._asyncio_lock = asyncio.Lock()
return self._asyncio_lock return self._asyncio_lock
@property
def lock(self) -> threading.Lock:
if self._lock is None:
self._lock = threading.Lock()
return self._lock
@property @property
def parent(self) -> "Dispatcher": def parent(self) -> "Dispatcher":
return self.manager.dispatchers[self.parent_name] return self.manager.dispatchers[self.parent_name]
...@@ -191,95 +223,53 @@ class Dispatcher(BaseModel): ...@@ -191,95 +223,53 @@ class Dispatcher(BaseModel):
functions only. Otherwise, the span_id should not be trusted, as the functions only. Otherwise, the span_id should not be trusted, as the
span decorator sets the span_id. span decorator sets the span_id.
""" """
span_id = self.current_span_id with self.lock:
span_id = self.current_span_id
dispatch_event: EventDispatcher = partial(self.event, span_id=span_id) dispatch_event: EventDispatcher = partial(self.event, span_id=span_id)
return dispatch_event return dispatch_event
@contextmanager def _get_parent_update_span_ctx_var(self, id_: str, current_task_name: str):
def dispatch_event(self): """Helper method to get parent id from the appropriate async contextvar."""
"""Context manager for firing events within a span session. current_thread = threading.get_ident()
thread_span_ctx = span_ctx_var.get().copy()
This context manager should be used with @dispatcher.span decorated span_ctx = thread_span_ctx[current_thread]
functions only. Otherwise, the span_id should not be trusted, as the if current_task_name not in span_ctx:
span decorator sets the span_id. parent_id = None
""" span_ctx[current_task_name] = [id_]
span_id = self.current_span_id else:
dispatch_event: EventDispatcher = partial(self.event, span_id=span_id) parent_id = span_ctx[current_task_name][-1]
span_ctx[current_task_name].append(id_)
try: span_ctx_var.set(thread_span_ctx)
yield dispatch_event
finally:
del dispatch_event
def async_span_with_parent_id(self, parent_id: str):
"""This decorator should be used to span an async function nested in an outer span.
Primary example: llama_index.core.async_utils.run_jobs
Args:
parent_id (str): The span_id of the outer span.
"""
def outer(func):
@wrapt.decorator
async def async_wrapper(func, instance, args, kwargs):
bound_args = inspect.signature(func).bind(*args, **kwargs)
id_ = f"{func.__qualname__}-{uuid.uuid4()}"
async with self.asyncio_lock:
self.current_span_id = id_
async with self.root.asyncio_lock:
self.root.current_span_id = id_
current_task = asyncio.current_task() return parent_id
current_task_name = current_task.get_name()
span_ctx_dict = span_ctx.get().copy()
if current_task_name not in span_ctx_dict:
span_ctx_dict[current_task_name] = [id_]
else:
span_ctx_dict[current_task_name].append(id_)
span_ctx.set(span_ctx_dict)
self.span_enter( def _pop_span_from_ctx_var(self, current_task_name: str) -> None:
id_=id_, """Helper method to pop completed/dropped span from async contextvar."""
bound_args=bound_args, current_thread = threading.get_ident()
instance=instance, thread_span_ctx = span_ctx_var.get().copy()
parent_id=parent_id, span_ctx = thread_span_ctx[current_thread]
)
try:
result = await func(*args, **kwargs)
except BaseException as e:
self.event(SpanDropEvent(span_id=id_, err_str=str(e)))
self.span_drop(
id_=id_, bound_args=bound_args, instance=instance, err=e
)
raise
else:
self.span_exit(
id_=id_, bound_args=bound_args, instance=instance, result=result
)
return result
finally:
# clean up
current_task = asyncio.current_task()
current_task_name = current_task.get_name()
span_ctx_dict = span_ctx.get().copy()
span_ctx_dict[current_task_name].pop()
if len(span_ctx_dict[current_task_name]) == 0:
del span_ctx_dict[current_task_name]
span_ctx.set(span_ctx_dict)
return async_wrapper(func) span_ctx[current_task_name].pop()
if len(span_ctx[current_task_name]) == 0:
return outer del span_ctx[current_task_name]
span_ctx_var.set(thread_span_ctx)
def span(self, func): def span(self, func):
@wrapt.decorator @wrapt.decorator
def wrapper(func, instance, args, kwargs): def wrapper(func, instance, args, kwargs):
bound_args = inspect.signature(func).bind(*args, **kwargs) bound_args = inspect.signature(func).bind(*args, **kwargs)
id_ = f"{func.__qualname__}-{uuid.uuid4()}" id_ = f"{func.__qualname__}-{uuid.uuid4()}"
self.current_span_id = id_ with self.lock:
self.root.current_span_id = id_ self.set_current_span_id(id_)
self.span_enter(id_=id_, bound_args=bound_args, instance=instance) with self.root.lock:
self.root.set_current_span_id(id_)
# get parent_id (thread-safe)
parent_id = self._get_parent_update_span_ctx_var(id_, DEFAULT_SYNC_KEY)
self.span_enter(
id_=id_, bound_args=bound_args, instance=instance, parent_id=parent_id
)
try: try:
result = func(*args, **kwargs) result = func(*args, **kwargs)
except BaseException as e: except BaseException as e:
...@@ -291,27 +281,24 @@ class Dispatcher(BaseModel): ...@@ -291,27 +281,24 @@ class Dispatcher(BaseModel):
id_=id_, bound_args=bound_args, instance=instance, result=result id_=id_, bound_args=bound_args, instance=instance, result=result
) )
return result return result
finally:
# clean up
self._pop_span_from_ctx_var(DEFAULT_SYNC_KEY)
@wrapt.decorator @wrapt.decorator
async def async_wrapper(func, instance, args, kwargs): async def async_wrapper(func, instance, args, kwargs):
bound_args = inspect.signature(func).bind(*args, **kwargs) bound_args = inspect.signature(func).bind(*args, **kwargs)
id_ = f"{func.__qualname__}-{uuid.uuid4()}" id_ = f"{func.__qualname__}-{uuid.uuid4()}"
async with self.asyncio_lock: async with self.asyncio_lock:
self.current_span_id = id_ self.set_current_span_id(id_)
async with self.root.asyncio_lock: async with self.root.asyncio_lock:
self.root.current_span_id = id_ self.root.set_current_span_id(id_)
# get parent_id # get parent_id (thread and async-task safe)
# spans are managed in this hieararchy: thread > async task > async coros
current_task = asyncio.current_task() current_task = asyncio.current_task()
current_task_name = current_task.get_name() current_task_name = current_task.get_name()
span_ctx_dict = span_ctx.get().copy() parent_id = self._get_parent_update_span_ctx_var(id_, current_task_name)
if current_task_name not in span_ctx_dict:
parent_id = None
span_ctx_dict[current_task_name] = [id_]
else:
parent_id = span_ctx_dict[current_task_name][-1]
span_ctx_dict[current_task_name].append(id_)
span_ctx.set(span_ctx_dict)
self.span_enter( self.span_enter(
id_=id_, bound_args=bound_args, instance=instance, parent_id=parent_id id_=id_, bound_args=bound_args, instance=instance, parent_id=parent_id
...@@ -329,19 +316,64 @@ class Dispatcher(BaseModel): ...@@ -329,19 +316,64 @@ class Dispatcher(BaseModel):
return result return result
finally: finally:
# clean up # clean up
current_task = asyncio.current_task() self._pop_span_from_ctx_var(current_task_name)
current_task_name = current_task.get_name()
span_ctx_dict = span_ctx.get().copy()
span_ctx_dict[current_task_name].pop()
if len(span_ctx_dict[current_task_name]) == 0:
del span_ctx_dict[current_task_name]
span_ctx.set(span_ctx_dict)
if inspect.iscoroutinefunction(func): if inspect.iscoroutinefunction(func):
return async_wrapper(func) return async_wrapper(func)
else: else:
return wrapper(func) return wrapper(func)
def async_span_with_parent_id(self, parent_id: str):
"""This decorator should be used to span an async function nested in an outer span.
Primary example: llama_index.core.async_utils.run_jobs
Args:
parent_id (str): The span_id of the outer span.
"""
def outer(func):
@wrapt.decorator
async def async_wrapper(func, instance, args, kwargs):
bound_args = inspect.signature(func).bind(*args, **kwargs)
id_ = f"{func.__qualname__}-{uuid.uuid4()}"
async with self.asyncio_lock:
self.set_current_span_id(id_)
async with self.root.asyncio_lock:
self.root.set_current_span_id(id_)
# don't need parent_id but need to update span ctx var
current_task = asyncio.current_task()
current_task_name = current_task.get_name()
_ = self._get_parent_update_span_ctx_var(id_, current_task_name)
self.span_enter(
id_=id_,
bound_args=bound_args,
instance=instance,
parent_id=parent_id,
)
try:
result = await func(*args, **kwargs)
except BaseException as e:
self.event(SpanDropEvent(span_id=id_, err_str=str(e)))
self.span_drop(
id_=id_, bound_args=bound_args, instance=instance, err=e
)
raise
else:
self.span_exit(
id_=id_, bound_args=bound_args, instance=instance, result=result
)
return result
finally:
# clean up
self._pop_span_from_ctx_var(current_task_name)
return async_wrapper(func)
return outer
@property @property
def log_name(self) -> str: def log_name(self) -> str:
"""Name to be used in logging.""" """Name to be used in logging."""
......
llama-index-core/llama_index/core/instrumentation/span_handlers/base.py
+ 67
− 15
View file @ 893c4057
import inspect import inspect
import threading
from abc import abstractmethod from abc import abstractmethod
from typing import Any, Dict, List, Generic, Optional, TypeVar from typing import Any, Dict, List, Generic, Optional, TypeVar
from llama_index.core.bridge.pydantic import BaseModel, Field from llama_index.core.bridge.pydantic import BaseModel, Field, PrivateAttr
from llama_index.core.instrumentation.span.base import BaseSpan from llama_index.core.instrumentation.span.base import BaseSpan
T = TypeVar("T", bound=BaseSpan) T = TypeVar("T", bound=BaseSpan)
...@@ -18,17 +19,50 @@ class BaseSpanHandler(BaseModel, Generic[T]): ...@@ -18,17 +19,50 @@ class BaseSpanHandler(BaseModel, Generic[T]):
dropped_spans: List[T] = Field( dropped_spans: List[T] = Field(
default_factory=list, description="List of completed spans." default_factory=list, description="List of completed spans."
) )
current_span_id: Optional[str] = Field( current_span_ids: Dict[Any, Optional[str]] = Field(
default=None, description="Id of current span." default={}, description="Id of current spans in a given thread."
) )
_lock: Optional[threading.Lock] = PrivateAttr()
class Config: class Config:
arbitrary_types_allowed = True arbitrary_types_allowed = True
def __init__(
self,
open_spans: Dict[str, T] = {},
completed_spans: List[T] = [],
dropped_spans: List[T] = [],
current_span_ids: Dict[Any, str] = {},
):
self._lock = None
super().__init__(
open_spans=open_spans,
completed_spans=completed_spans,
dropped_spans=dropped_spans,
current_span_ids=current_span_ids,
)
def class_name(cls) -> str: def class_name(cls) -> str:
"""Class name.""" """Class name."""
return "BaseSpanHandler" return "BaseSpanHandler"
@property
def lock(self) -> threading.Lock:
if self._lock is None:
self._lock = threading.Lock()
return self._lock
@property
def current_span_id(self) -> Optional[str]:
current_thread = threading.get_ident()
if current_thread in self.current_span_ids:
return self.current_span_ids[current_thread]
return None
def set_current_span_id(self, value: str) -> None:
current_thread = threading.get_ident()
self.current_span_ids[current_thread] = value
def span_enter( def span_enter(
self, self,
id_: str, id_: str,
...@@ -49,8 +83,9 @@ class BaseSpanHandler(BaseModel, Generic[T]): ...@@ -49,8 +83,9 @@ class BaseSpanHandler(BaseModel, Generic[T]):
parent_span_id=parent_id or self.current_span_id, parent_span_id=parent_id or self.current_span_id,
) )
if span: if span:
self.open_spans[id_] = span with self.lock:
self.current_span_id = id_ self.open_spans[id_] = span
self.set_current_span_id(id_)
def span_exit( def span_exit(
self, self,
...@@ -65,11 +100,13 @@ class BaseSpanHandler(BaseModel, Generic[T]): ...@@ -65,11 +100,13 @@ class BaseSpanHandler(BaseModel, Generic[T]):
id_=id_, bound_args=bound_args, instance=instance, result=result id_=id_, bound_args=bound_args, instance=instance, result=result
) )
if span: if span:
if self.current_span_id == id_: with self.lock:
self.current_span_id = self.open_spans[id_].parent_id if self.current_span_id == id_:
del self.open_spans[id_] self.set_current_span_id(self.open_spans[id_].parent_id)
del self.open_spans[id_]
if not self.open_spans: # empty so flush if not self.open_spans: # empty so flush
self.current_span_id = None with self.lock:
self.set_current_span_id(None)
def span_drop( def span_drop(
self, self,
...@@ -84,9 +121,10 @@ class BaseSpanHandler(BaseModel, Generic[T]): ...@@ -84,9 +121,10 @@ class BaseSpanHandler(BaseModel, Generic[T]):
id_=id_, bound_args=bound_args, instance=instance, err=err id_=id_, bound_args=bound_args, instance=instance, err=err
) )
if span: if span:
if self.current_span_id == id_: with self.lock:
self.current_span_id = self.open_spans[id_].parent_id if self.current_span_id == id_:
del self.open_spans[id_] self.set_current_span_id(self.open_spans[id_].parent_id)
del self.open_spans[id_]
@abstractmethod @abstractmethod
def new_span( def new_span(
...@@ -97,7 +135,11 @@ class BaseSpanHandler(BaseModel, Generic[T]): ...@@ -97,7 +135,11 @@ class BaseSpanHandler(BaseModel, Generic[T]):
parent_span_id: Optional[str] = None, parent_span_id: Optional[str] = None,
**kwargs: Any, **kwargs: Any,
) -> Optional[T]: ) -> Optional[T]:
"""Create a span.""" """Create a span.
Subclasses of BaseSpanHandler should create the respective span type T
and return it. Only NullSpanHandler should return a None here.
"""
... ...
@abstractmethod @abstractmethod
...@@ -109,7 +151,12 @@ class BaseSpanHandler(BaseModel, Generic[T]): ...@@ -109,7 +151,12 @@ class BaseSpanHandler(BaseModel, Generic[T]):
result: Optional[Any] = None, result: Optional[Any] = None,
**kwargs: Any, **kwargs: Any,
) -> Optional[T]: ) -> Optional[T]:
"""Logic for preparing to exit a span.""" """Logic for preparing to exit a span.
Subclasses of BaseSpanHandler should return back the specific span T
that is to be exited. If None is returned, then the span won't actually
be exited.
"""
... ...
@abstractmethod @abstractmethod
...@@ -121,5 +168,10 @@ class BaseSpanHandler(BaseModel, Generic[T]): ...@@ -121,5 +168,10 @@ class BaseSpanHandler(BaseModel, Generic[T]):
err: Optional[BaseException] = None, err: Optional[BaseException] = None,
**kwargs: Any, **kwargs: Any,
) -> Optional[T]: ) -> Optional[T]:
"""Logic for preparing to drop a span.""" """Logic for preparing to drop a span.
Subclasses of BaseSpanHandler should return back the specific span T
that is to be dropped. If None is returned, then the span won't actually
be dropped.
"""
... ...
llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py
+ 60
− 34
View file @ 893c4057
...@@ -3,6 +3,7 @@ from typing import Any, cast, List, Optional, TYPE_CHECKING ...@@ -3,6 +3,7 @@ from typing import Any, cast, List, Optional, TYPE_CHECKING
from llama_index.core.instrumentation.span.simple import SimpleSpan from llama_index.core.instrumentation.span.simple import SimpleSpan
from llama_index.core.instrumentation.span_handlers.base import BaseSpanHandler from llama_index.core.instrumentation.span_handlers.base import BaseSpanHandler
from datetime import datetime from datetime import datetime
from functools import reduce
import warnings import warnings
if TYPE_CHECKING: if TYPE_CHECKING:
...@@ -40,7 +41,8 @@ class SimpleSpanHandler(BaseSpanHandler[SimpleSpan]): ...@@ -40,7 +41,8 @@ class SimpleSpanHandler(BaseSpanHandler[SimpleSpan]):
span = cast(SimpleSpan, span) span = cast(SimpleSpan, span)
span.end_time = datetime.now() span.end_time = datetime.now()
span.duration = (span.end_time - span.start_time).total_seconds() span.duration = (span.end_time - span.start_time).total_seconds()
self.completed_spans += [span] with self.lock:
self.completed_spans += [span]
return span return span
def prepare_to_drop_span( def prepare_to_drop_span(
...@@ -53,55 +55,79 @@ class SimpleSpanHandler(BaseSpanHandler[SimpleSpan]): ...@@ -53,55 +55,79 @@ class SimpleSpanHandler(BaseSpanHandler[SimpleSpan]):
) -> SimpleSpan: ) -> SimpleSpan:
"""Logic for droppping a span.""" """Logic for droppping a span."""
if id_ in self.open_spans: if id_ in self.open_spans:
span = self.open_spans[id_] with self.lock:
span.metadata = {"error": str(err)} span = self.open_spans[id_]
self.dropped_spans += [span] span.metadata = {"error": str(err)}
self.dropped_spans += [span]
return span return span
return None return None
def _get_parents(self) -> List[SimpleSpan]:
"""Helper method to get all parent/root spans."""
all_spans = self.completed_spans + self.dropped_spans
return [s for s in all_spans if s.parent_id is None]
def _build_tree_by_parent(
self, parent: SimpleSpan, acc: List[SimpleSpan], spans: List[SimpleSpan]
) -> List[SimpleSpan]:
"""Builds the tree by parent root."""
if not spans:
return acc
children = [s for s in spans if s.parent_id == parent.id_]
if not children:
return acc
updated_spans = [s for s in spans if s not in children]
children_trees = [
self._build_tree_by_parent(
parent=c, acc=[c], spans=[s for s in updated_spans if c != s]
)
for c in children
]
return acc + reduce(lambda x, y: x + y, children_trees)
def _get_trace_trees(self) -> List["Tree"]: def _get_trace_trees(self) -> List["Tree"]:
"""Method for getting trace trees.""" """Method for getting trace trees."""
try: try:
from treelib import Tree from treelib import Tree
from treelib.exceptions import NodeIDAbsentError
except ImportError as e: except ImportError as e:
raise ImportError( raise ImportError(
"`treelib` package is missing. Please install it by using " "`treelib` package is missing. Please install it by using "
"`pip install treelib`." "`pip install treelib`."
) )
sorted_spans = sorted(
self.completed_spans + self.dropped_spans, key=lambda x: x.start_time all_spans = self.completed_spans + self.dropped_spans
) for s in all_spans:
if s.parent_id is None:
continue
if not any(ns.id_ == s.parent_id for ns in all_spans):
warnings.warn("Parent with id {span.parent_id} missing from spans")
s.parent_id += "-MISSING"
all_spans.append(SimpleSpan(id_=s.parent_id, parent_id=None))
parents = self._get_parents()
span_groups = []
for p in parents:
this_span_group = self._build_tree_by_parent(
parent=p, acc=[p], spans=[s for s in all_spans if s != p]
)
sorted_span_group = sorted(this_span_group, key=lambda x: x.start_time)
span_groups.append(sorted_span_group)
trees = [] trees = []
tree = Tree() tree = Tree()
for span in sorted_spans: for grp in span_groups:
if span.parent_id is None: for span in grp:
# complete old tree unless its empty (i.e., start of loop) if span.parent_id is None:
if tree.all_nodes(): # complete old tree unless its empty (i.e., start of loop)
trees.append(tree) if tree.all_nodes():
# start new tree trees.append(tree)
tree = Tree() # start new tree
tree = Tree()
try:
tree.create_node(
tag=f"{span.id_} ({span.duration})",
identifier=span.id_,
parent=span.parent_id,
data=span.start_time,
)
except NodeIDAbsentError:
warnings.warn("Parent with id {span.parent_id} missing from spans")
# create new tree and fake parent node
trees.append(tree)
tree = Tree()
tree.create_node(
tag=f"{span.parent_id} (MISSING)",
identifier=span.parent_id,
parent=None,
data=span.start_time,
)
tree.create_node( tree.create_node(
tag=f"{span.id_} ({span.duration})", tag=f"{span.id_} ({span.duration})",
identifier=span.id_, identifier=span.id_,
......
llama-index-core/tests/instrumentation/test_dispatcher.py
+ 12
− 2
View file @ 893c4057
...@@ -126,7 +126,12 @@ def test_dispatcher_span_args(mock_uuid, mock_span_enter, mock_span_exit): ...@@ -126,7 +126,12 @@ def test_dispatcher_span_args(mock_uuid, mock_span_enter, mock_span_exit):
mock_span_enter.assert_called_once() mock_span_enter.assert_called_once()
args, kwargs = mock_span_enter.call_args args, kwargs = mock_span_enter.call_args
assert args == () assert args == ()
assert kwargs == {"id_": span_id, "bound_args": bound_args, "instance": None} assert kwargs == {
"id_": span_id,
"bound_args": bound_args,
"instance": None,
"parent_id": None,
}
# span_exit # span_exit
args, kwargs = mock_span_exit.call_args args, kwargs = mock_span_exit.call_args
...@@ -157,7 +162,12 @@ def test_dispatcher_span_args_with_instance(mock_uuid, mock_span_enter, mock_spa ...@@ -157,7 +162,12 @@ def test_dispatcher_span_args_with_instance(mock_uuid, mock_span_enter, mock_spa
mock_span_enter.assert_called_once() mock_span_enter.assert_called_once()
args, kwargs = mock_span_enter.call_args args, kwargs = mock_span_enter.call_args
assert args == () assert args == ()
assert kwargs == {"id_": span_id, "bound_args": bound_args, "instance": instance} assert kwargs == {
"id_": span_id,
"bound_args": bound_args,
"instance": instance,
"parent_id": None,
}
# span_exit # span_exit
args, kwargs = mock_span_exit.call_args args, kwargs = mock_span_exit.call_args
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment

:feet::heart: Strive to be the person your dogs believe you are :heart::feet: