andras/great-ai
1
0
Fork 0
Code Issues Pull requests 1 Projects Releases 4 Packages Wiki Activity Actions

Fix async tracing context

Browse source
This commit is contained in:
Andras Schmelczer 2022-07-04 08:47:17 +02:00
parent 26a91293f2
commit ee1ab71df2
No known key found for this signature in database
GPG key ID: 0EA1BC97D0AB076E
1 changed files with 12 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

35
src/great_ai/great_ai/tracing/tracing_context.py
View file

@ -1,18 +1,7 @@
import threading from contextvars import ContextVar
from collections import defaultdict
from datetime import datetime from datetime import datetime
from types import TracebackType from types import TracebackType
from typing import ( from typing import Any, Dict, Generic, List, Literal, Optional, Type, TypeVar
Any,
DefaultDict,
Dict,
Generic,
List,
Literal,
Optional,
Type,
TypeVar,
)
from ..constants import DEVELOPMENT_TAG_NAME, ONLINE_TAG_NAME, PRODUCTION_TAG_NAME from ..constants import DEVELOPMENT_TAG_NAME, ONLINE_TAG_NAME, PRODUCTION_TAG_NAME
from ..context import get_context from ..context import get_context
@ -22,8 +11,6 @@ T = TypeVar("T")
class TracingContext(Generic[T]): class TracingContext(Generic[T]):
_contexts: DefaultDict[int, List["TracingContext"]] = defaultdict(lambda: [])
def __init__(self, function_name: str, do_not_persist_traces: bool) -> None: def __init__(self, function_name: str, do_not_persist_traces: bool) -> None:
self._do_not_persist_traces = do_not_persist_traces self._do_not_persist_traces = do_not_persist_traces
self._models: List[Model] = [] self._models: List[Model] = []
@ -62,14 +49,12 @@ class TracingContext(Generic[T]):
return self._trace return self._trace
@classmethod @staticmethod
def get_current_tracing_context(cls) -> Optional["TracingContext"]: def get_current_tracing_context() -> Optional["TracingContext"]:
if cls._contexts[threading.get_ident()]: return _current_tracing_context.get()
return cls._contexts[threading.get_ident()][-1]
return None
def __enter__(self) -> "TracingContext": def __enter__(self) -> "TracingContext":
self._contexts[threading.get_ident()].append(self) _current_tracing_context.set(self)
return self return self
def __exit__( def __exit__(
@ -78,8 +63,7 @@ class TracingContext(Generic[T]):
exception: Optional[BaseException], exception: Optional[BaseException],
traceback: Optional[TracebackType], traceback: Optional[TracebackType],
) -> Literal[False]: ) -> Literal[False]:
assert self._contexts[threading.get_ident()][-1] == self _current_tracing_context.set(None)
self._contexts[threading.get_ident()].remove(self)
if exception is not None and type is not None: if exception is not None and type is not None:
self.finalise(exception=exception) self.finalise(exception=exception)
@ -95,3 +79,8 @@ class TracingContext(Generic[T]):
get_context().tracing_database.save(self._trace) get_context().tracing_database.save(self._trace)
return False return False
_current_tracing_context: ContextVar[Optional[TracingContext]] = ContextVar(
"_current_tracing_context"
)