diff --git a/great_ai/tracing/tracing_context.py b/great_ai/tracing/tracing_context.py index 9180061..5f5c548 100644 --- a/great_ai/tracing/tracing_context.py +++ b/great_ai/tracing/tracing_context.py @@ -74,7 +74,7 @@ class TracingContext(Generic[T]): return _current_tracing_context.get() def __enter__(self) -> "TracingContext": - _current_tracing_context.set(self) + self._token = _current_tracing_context.set(self) return self def __exit__( @@ -83,7 +83,7 @@ class TracingContext(Generic[T]): exception: Optional[BaseException], traceback: Optional[TracebackType], ) -> Literal[False]: - _current_tracing_context.set(None) + _current_tracing_context.reset(self._token) if exception is not None and type is not None: self.finalise(exception=exception) diff --git a/tests/test_tracing_regressions.py b/tests/test_tracing_regressions.py new file mode 100644 index 0000000..dc0db70 --- /dev/null +++ b/tests/test_tracing_regressions.py @@ -0,0 +1,27 @@ +import pytest +from great_ai import GreatAI, log_metric +from great_ai.tracing import TracingContext + + +@pytest.mark.parametrize("fail", [False, True]) +def test_nested_predictions_restore_outer_trace( + isolated_context: None, fail: bool +) -> None: + @GreatAI.create + def inner(value: int) -> int: + if fail: + raise ValueError("inner failed") + return value + + @GreatAI.create + def outer(value: int) -> int: + try: + inner(value) + except ValueError: + pass + log_metric("after_inner", 7, disable_logging=True) + return value + + trace = outer(1) + assert trace.logged_values["metric:outer:after_inner"] == 7 + assert TracingContext.get_current_tracing_context() is None