From 930710ffc87a3787f6932f4bcd3a9c16b453e146 Mon Sep 17 00:00:00 2001 From: Dhruv Gupta Date: Thu, 7 May 2026 00:03:33 +0000 Subject: [PATCH] test_vector_search_retriever_tool: flush async traces before assertion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit mlflow 3.12 (released 2026-05-05) flipped OSS-default trace logging from synchronous to asynchronous (mlflow/mlflow#22304). After the bump, ``get_trace(get_last_active_trace_id())`` reads the span store before the BatchSpanProcessor has flushed the just-ended span and returns None, so the next ``trace.search_spans(...)`` call raises: AttributeError: 'NoneType' object has no attribute 'search_spans' This shows up on both ``openai_test`` and ``langchain_test`` matrix variants on ``uv: highest`` (where mlflow resolves to 3.12); the ``lowest-direct`` and pinned-version variants resolve an older mlflow and pass. mlflow's own fix is to pass ``flush=True`` to ``get_trace``, which calls ``_flush_pending_async_trace_writes()`` and retries before returning. Two-line patch — no version pin, no production-code change. Surfaced first by https://github.com/databricks/databricks-ai-bridge/pull/427 because it was the first CI run after the mlflow bump; main hasn't run CI since 2026-05-01 so it'd manifest there too. Co-authored-by: Isaac Signed-off-by: Dhruv Gupta --- .../tests/unit_tests/test_vector_search_retriever_tool.py | 1 + .../openai/tests/unit_tests/test_vector_search_retriever_tool.py | 1 + 2 files changed, 2 insertions(+) diff --git a/integrations/langchain/tests/unit_tests/test_vector_search_retriever_tool.py b/integrations/langchain/tests/unit_tests/test_vector_search_retriever_tool.py index 7c79ad01e..d2373dc0d 100644 --- a/integrations/langchain/tests/unit_tests/test_vector_search_retriever_tool.py +++ b/integrations/langchain/tests/unit_tests/test_vector_search_retriever_tool.py @@ -214,6 +214,7 @@ def test_vector_search_retriever_tool_description_generation(index_name: str) -> def test_vs_tool_tracing(index_name: str, tool_name: str | None) -> None: vector_search_tool = init_vector_search_tool(index_name, tool_name=tool_name) vector_search_tool._run("Databricks Agent Framework") + mlflow.flush_trace_async_logging() trace = mlflow.get_trace(mlflow.get_last_active_trace_id()) spans = trace.search_spans(name=tool_name or index_name, span_type=SpanType.RETRIEVER) assert len(spans) == 1 diff --git a/integrations/openai/tests/unit_tests/test_vector_search_retriever_tool.py b/integrations/openai/tests/unit_tests/test_vector_search_retriever_tool.py index a87698c60..459130f05 100644 --- a/integrations/openai/tests/unit_tests/test_vector_search_retriever_tool.py +++ b/integrations/openai/tests/unit_tests/test_vector_search_retriever_tool.py @@ -187,6 +187,7 @@ def test_vector_search_retriever_tool_init( assert all(["id" in d["metadata"] for d in docs]) # Ensure tracing works properly + mlflow.flush_trace_async_logging() trace = mlflow.get_trace(mlflow.get_last_active_trace_id()) spans = trace.search_spans(name=tool_name or index_name, span_type=SpanType.RETRIEVER) assert len(spans) == 1