diff --git a/py/src/braintrust/test_trace.py b/py/src/braintrust/test_trace.py index a16ae042..4bfb2b59 100644 --- a/py/src/braintrust/test_trace.py +++ b/py/src/braintrust/test_trace.py @@ -1,5 +1,6 @@ """Tests for Trace functionality.""" +import asyncio import os import braintrust @@ -213,6 +214,409 @@ async def fetch_fn(filters): assert sorted(span_ids) == ["fn-1", "llm-1", "llm-2"] assert len(span_ids) == len(set(span_ids)), f"duplicate spans: {span_ids}" + @pytest.mark.parametrize( + ("first_types", "second_types"), + [ + (None, None), + (["llm"], ["llm"]), + (["llm", "tool"], ["tool", "function"]), + (None, ["llm"]), + ], + ) + @pytest.mark.parametrize("completion_order", [(0, 1), (1, 0)]) + @pytest.mark.asyncio + async def test_concurrent_fetches_do_not_append_duplicate_groups( + self, first_types, second_types, completion_order + ): + # Repeated IDs within one response are preserved; only duplication caused + # by merging overlapping responses should be removed. + all_spans = [ + make_span("llm-1", "llm"), + make_span("llm-1", "llm"), + make_span("tool-1", "tool"), + make_span("function-1", "function"), + ] + requests = asyncio.Queue() + + async def fetch_fn(filters): + response = asyncio.get_running_loop().create_future() + await requests.put((filters, response)) + return await response + + fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) + filters = [{"span_type": types} if types is not None else {} for types in (first_types, second_types)] + tasks = [asyncio.create_task(fetcher.get_spans(filters=value)) for value in filters] + try: + pending = [await asyncio.wait_for(requests.get(), timeout=5) for _ in tasks] + for index in completion_order: + requested_filters, response = pending[index] + assert requested_filters == filters[index] + expected = [span for span in all_spans if _matches_span_filters(span, filters[index])] + response.set_result(expected) + result = await asyncio.wait_for(tasks[index], timeout=5) + assert sorted(span.span_id for span in result) == sorted(span.span_id for span in expected) + + cached_types = None if first_types is None else sorted(set(first_types + second_types)) + cached_filters = {"span_type": cached_types} if cached_types is not None else {} + cached = await asyncio.wait_for(fetcher.get_spans(filters=cached_filters), timeout=5) + expected = [span for span in all_spans if _matches_span_filters(span, cached_filters)] + assert sorted(span.span_id for span in cached) == sorted(span.span_id for span in expected) + assert requests.empty() + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + @pytest.mark.parametrize("older_empty", [False, True], ids=["older-nonempty", "older-empty"]) + @pytest.mark.parametrize("new_span_type", ["llm", "tool"]) + @pytest.mark.parametrize("completion_order", [(0, 1), (1, 0)]) + @pytest.mark.asyncio + async def test_concurrent_full_fetches_keep_newer_snapshot(self, new_span_type, completion_order, older_empty): + older_spans = [] if older_empty else [make_span("llm-1", "llm")] + newer_spans = [*older_spans, make_span("new-1", new_span_type)] + snapshots = [older_spans, newer_spans] + requests = asyncio.Queue() + + async def fetch_fn(filters): + assert filters == {} + response = asyncio.get_running_loop().create_future() + await requests.put(response) + return await response + + fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) + tasks = [] + responses = [] + try: + for _ in snapshots: + tasks.append(asyncio.create_task(fetcher.get_spans())) + responses.append(await asyncio.wait_for(requests.get(), timeout=5)) + for index in completion_order: + responses[index].set_result(snapshots[index]) + await asyncio.wait_for(tasks[index], timeout=5) + + # A late older response must not make the newer span disappear from + # the complete cache or the last caller's result. + assert fetcher._span_cache + cached = await asyncio.wait_for(fetcher.get_spans(), timeout=5) + assert [tasks[completion_order[-1]].result(), cached] == [newer_spans, newer_spans] + expected_type = [span for span in newer_spans if span.span_attributes["type"] == new_span_type] + assert await fetcher.get_spans(filters={"span_type": [new_span_type]}) == expected_type + assert requests.empty() + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + @pytest.mark.parametrize( + ("first_types", "second_types", "expected"), + [ + (None, ["llm"], [("absent-1", "old"), ("llm-1", "new"), ("llm-2", "new"), ("tool-1", "old")]), + (None, ["function"], [("absent-1", "old"), ("function-1", "new"), ("llm-1", "old"), ("tool-1", "old")]), + ( + ["llm"], + None, + [("function-1", "new"), ("llm-1", "new"), ("llm-2", "new"), ("tool-1", "new")], + ), + ( + ["absent"], + None, + [("function-1", "new"), ("llm-1", "new"), ("llm-2", "new"), ("tool-1", "new")], + ), + (["llm"], ["llm"], [("llm-1", "new"), ("llm-2", "new")]), + (["llm", "tool"], ["tool", "function"], [("function-1", "new"), ("llm-1", "old"), ("tool-1", "new")]), + ], + ) + @pytest.mark.parametrize("completion_order", [(0, 1), (1, 0)]) + @pytest.mark.asyncio + async def test_overlapping_fetches_keep_newer_type_groups( + self, first_types, second_types, expected, completion_order + ): + snapshots = [ + [ + make_span("llm-1", "llm", metadata={"version": "old"}), + make_span("tool-1", "tool", metadata={"version": "old"}), + make_span("absent-1", "absent", metadata={"version": "old"}), + ], + [ + make_span("llm-1", "llm", metadata={"version": "new"}), + make_span("llm-2", "llm", metadata={"version": "new"}), + make_span("tool-1", "tool", metadata={"version": "new"}), + make_span("function-1", "function", metadata={"version": "new"}), + ], + ] + requests = asyncio.Queue() + + async def fetch_fn(filters): + response = asyncio.get_running_loop().create_future() + await requests.put((filters, response)) + return await response + + fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) + tasks = [] + pending = [] + try: + for types in (first_types, second_types): + filters = {"span_type": types} if types is not None else {} + tasks.append(asyncio.create_task(fetcher.get_spans(filters=filters))) + pending.append(await asyncio.wait_for(requests.get(), timeout=5)) + for index in completion_order: + filters, response = pending[index] + response.set_result([span for span in snapshots[index] if _matches_span_filters(span, filters)]) + await asyncio.wait_for(tasks[index], timeout=5) + + types = None if first_types is None or second_types is None else sorted(set(first_types + second_types)) + cached = await asyncio.wait_for( + fetcher.get_spans(filters={"span_type": types} if types is not None else {}), timeout=5 + ) + assert sorted((span.span_id, span.metadata["version"]) for span in cached) == expected + assert requests.empty() + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + @pytest.mark.parametrize("completion_order", [(0, 1), (1, 0)]) + @pytest.mark.asyncio + async def test_empty_full_fetch_keeps_newer_type_retriable(self, completion_order): + spans = [make_span("llm-1", "llm")] + full_spans = [*spans, make_span("tool-1", "tool")] + requests = asyncio.Queue() + calls = [] + + async def fetch_fn(filters): + calls.append(filters) + if len(calls) > 2: + return full_spans + response = asyncio.get_running_loop().create_future() + await requests.put(response) + return await response + + fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) + tasks = [] + responses = [] + try: + for filters in ({}, {"span_type": ["llm"]}): + tasks.append(asyncio.create_task(fetcher.get_spans(filters=filters))) + responses.append(await asyncio.wait_for(requests.get(), timeout=5)) + for index in completion_order: + responses[index].set_result([] if index == 0 else spans) + await asyncio.wait_for(tasks[index], timeout=5) + + assert await fetcher.get_spans(filters={"span_type": ["llm"]}) == spans + assert calls == [{}, {"span_type": ["llm"]}] + assert await fetcher.get_spans() == full_spans + assert calls == [{}, {"span_type": ["llm"]}, {}] + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + @pytest.mark.parametrize("cancel", [False, True], ids=["failure", "cancellation"]) + @pytest.mark.parametrize("interrupted_first", [False, True]) + @pytest.mark.parametrize( + "first_types,second_types", [(None, None), (None, ["llm"]), (["llm"], None), (["llm"], ["llm"])] + ) + @pytest.mark.asyncio + async def test_interrupted_newer_fetch_preserves_older_response( + self, cancel, interrupted_first, first_types, second_types + ): + spans = [make_span("llm-1", "llm")] + requests = asyncio.Queue() + + async def fetch_fn(_filters): + response = asyncio.get_running_loop().create_future() + await requests.put(response) + return await response + + fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) + tasks = [] + responses = [] + try: + for types in (first_types, second_types): + filters = {"span_type": types} if types is not None else {} + tasks.append(asyncio.create_task(fetcher.get_spans(filters=filters))) + responses.append(await asyncio.wait_for(requests.get(), timeout=5)) + for index in (1, 0) if interrupted_first else (0, 1): + if index == 0: + responses[0].set_result(spans) + assert await asyncio.wait_for(tasks[0], timeout=5) == spans + elif cancel: + tasks[1].cancel() + with pytest.raises(asyncio.CancelledError): + await tasks[1] + else: + responses[1].set_exception(RuntimeError("fetch failed")) + with pytest.raises(RuntimeError, match="fetch failed"): + await tasks[1] + + filters = {"span_type": first_types} if first_types is not None else {} + assert await asyncio.wait_for(fetcher.get_spans(filters=filters), timeout=5) == spans + assert requests.empty() + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + @pytest.mark.parametrize("first_types", [None, ["llm"]]) + @pytest.mark.parametrize( + "later_kwargs,empty", + [ + ({"filters": {"name": ["search"]}}, False), + ({"include_scorers": True}, False), + ({"filters": {"span_type": ["llm"]}}, True), + ], + ) + @pytest.mark.asyncio + async def test_uncached_response_does_not_supersede_older_fetch(self, first_types, later_kwargs, empty): + older_spans = [make_span("llm-1", "llm")] + later_spans = [] if empty else [make_span("llm-2", "llm", name="search")] + started = asyncio.Event() + response = asyncio.get_running_loop().create_future() + call_count = 0 + + async def fetch_fn(_filters): + nonlocal call_count + call_count += 1 + if call_count == 1: + started.set() + return await response + return later_spans + + fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) + filters = {"span_type": first_types} if first_types is not None else {} + task = asyncio.create_task(fetcher.get_spans(filters=filters)) + try: + await asyncio.wait_for(started.wait(), timeout=5) + assert await fetcher.get_spans(**later_kwargs) == later_spans + response.set_result(older_spans) + assert await asyncio.wait_for(task, timeout=5) == older_spans + assert await fetcher.get_spans(filters=filters) == older_spans + assert call_count == 2 + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + @pytest.mark.parametrize("first_types", [None, ["llm"]]) + @pytest.mark.parametrize("completion_order", [(0, 1), (1, 0)]) + @pytest.mark.asyncio + async def test_newer_empty_full_fetch_remains_retriable(self, first_types, completion_order): + older_spans = [make_span("llm-1", "llm")] + retry_spans = [make_span("tool-1", "tool")] + requests = asyncio.Queue() + call_count = 0 + + async def fetch_fn(_filters): + nonlocal call_count + call_count += 1 + if call_count > 2: + return retry_spans + response = asyncio.get_running_loop().create_future() + await requests.put(response) + return await response + + fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) + tasks = [] + responses = [] + try: + for types in (first_types, None): + filters = {"span_type": types} if types is not None else {} + tasks.append(asyncio.create_task(fetcher.get_spans(filters=filters))) + responses.append(await asyncio.wait_for(requests.get(), timeout=5)) + for index in completion_order: + responses[index].set_result(older_spans if index == 0 else []) + await asyncio.wait_for(tasks[index], timeout=5) + + assert tasks[completion_order[-1]].result() == [] + assert await fetcher.get_spans() == retry_spans + assert call_count == 3 + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + @pytest.mark.parametrize("cancel", [False, True], ids=["failure", "cancellation"]) + @pytest.mark.asyncio + async def test_interrupted_full_fetch_preserves_cached_types(self, cancel): + cached_span = make_span("llm-1", "llm") + full_spans = [cached_span, make_span("tool-1", "tool")] + started = asyncio.Event() + response = asyncio.get_running_loop().create_future() + calls = [] + + async def fetch_fn(filters): + calls.append(filters) + if len(calls) == 1: + return [cached_span] + if len(calls) == 2: + started.set() + return await response + return full_spans + + fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) + assert await fetcher.get_spans(filters={"span_type": ["llm"]}) == [cached_span] + task = asyncio.create_task(fetcher.get_spans()) + try: + await asyncio.wait_for(started.wait(), timeout=5) + if cancel: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + response.set_exception(RuntimeError("fetch failed")) + with pytest.raises(RuntimeError, match="fetch failed"): + await task + + assert await fetcher.get_spans(filters={"span_type": ["llm"]}) == [cached_span] + assert calls == [{"span_type": ["llm"]}, {}] + assert await fetcher.get_spans() == full_spans + assert calls == [{"span_type": ["llm"]}, {}, {}] + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + @pytest.mark.asyncio + async def test_concurrent_empty_full_fetch_remains_retriable(self): + spans = [make_span("llm-1", "llm")] + requests = asyncio.Queue() + call_count = 0 + + async def fetch_fn(_filters): + nonlocal call_count + call_count += 1 + if call_count > 2: + return spans + response = asyncio.get_running_loop().create_future() + await requests.put(response) + return await response + + fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) + tasks = [asyncio.create_task(fetcher.get_spans()) for _ in range(2)] + try: + responses = [await asyncio.wait_for(requests.get(), timeout=5) for _ in tasks] + responses[0].set_result(spans) + assert await tasks[0] == spans + responses[1].set_result([]) + assert await tasks[1] == [] + assert await fetcher.get_spans() == spans + assert call_count == 3 + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + @pytest.mark.asyncio + async def test_full_fetch_replaces_stale_types(self): + cached_span = make_span("llm-1", "llm") + full_spans = [make_span("tool-1", "tool")] + + async def fetch_fn(filters): + return [cached_span] if filters else full_spans + + fetcher = CachedSpanFetcher(fetch_fn=fetch_fn) + assert await fetcher.get_spans(filters={"span_type": ["llm"]}) == [cached_span] + assert await fetcher.get_spans() == full_spans + assert await fetcher.get_spans(filters={"span_type": ["llm"]}) == [] + def test_span_data_roundtrip(self): row = { "span_id": "tool-span", diff --git a/py/src/braintrust/trace.py b/py/src/braintrust/trace.py index cdcb6455..040b33c5 100644 --- a/py/src/braintrust/trace.py +++ b/py/src/braintrust/trace.py @@ -314,6 +314,9 @@ def __init__( ): self._span_cache: dict[str, list[SpanData]] = {} self._all_fetched = False + self._fetch_generation = 0 + self._committed_full_fetch_generation = 0 + self._span_cache_generations: dict[str, int] = {} if fetch_fn is not None: # Direct fetch function injection (for testing). Like the server, the injected @@ -393,13 +396,7 @@ async def get_spans( # If no filter requested, fetch everything. if not span_type: - # A full fetch is authoritative; reset the per-type cache first so a - # prior typed fetch's spans are not duplicated by re-fetching them - # (_fetch_spans appends). - self._span_cache = {} await self._fetch_spans(None) - if self._span_cache: # Only cache if we got results - self._all_fetched = True return self._get_from_cache(None) # Find which span types we don't have in cache yet. @@ -414,14 +411,42 @@ async def _fetch_spans(self, span_type: list[str] | None) -> None: Spans are filed by the type they report, not the type that was asked for, so a requested type that yields nothing leaves no entry and will be asked for again. """ + self._fetch_generation += 1 + generation = self._fetch_generation spans = await self._fetch_fn({"span_type": span_type} if span_type else {}, False) + fetched: dict[str, list[SpanData]] = {} for span in spans: span_attrs = span.span_attributes or {} span_type_str = span_attrs.get("type", "") - if span_type_str not in self._span_cache: - self._span_cache[span_type_str] = [] - self._span_cache[span_type_str].append(span) + fetched.setdefault(span_type_str, []).append(span) + + # Commit complete groups without awaiting, so overlapping requests never + # append the same spans twice. A failed or cancelled fetch leaves the + # previous cache intact. Generations order requests, not server revisions. + if span_type is None: + if generation < self._committed_full_fetch_generation: + return + # Replace stale groups, retaining only groups committed by requests + # started after this full fetch. An empty full response stays + # retriable even when a newer typed response is retained. + all_fetched = bool(fetched) + cache_generations = dict.fromkeys(fetched, generation) + for type_str, cached_generation in self._span_cache_generations.items(): + if cached_generation > generation: + fetched[type_str] = self._span_cache[type_str] + cache_generations[type_str] = cached_generation + self._span_cache = fetched + self._span_cache_generations = cache_generations + self._all_fetched = all_fetched + self._committed_full_fetch_generation = generation + else: + for type_str, group in fetched.items(): + if generation > max( + self._committed_full_fetch_generation, self._span_cache_generations.get(type_str, 0) + ): + self._span_cache[type_str] = group + self._span_cache_generations[type_str] = generation def _get_from_cache(self, span_type: list[str] | None) -> list[SpanData]: """Read spans back out of the cache, optionally narrowing to some types.