diff --git a/api/core/ops/mlflow_trace/mlflow_trace.py b/api/core/ops/mlflow_trace/mlflow_trace.py index 745157d108..355232a969 100644 --- a/api/core/ops/mlflow_trace/mlflow_trace.py +++ b/api/core/ops/mlflow_trace/mlflow_trace.py @@ -53,7 +53,7 @@ class MLflowDataTrace(BaseTraceInstance): """Lazy initialization to prevent blocking workflow when tracing service is unavailable.""" if self._initialized: return - + try: if isinstance(self._config, DatabricksConfig): self._setup_databricks(self._config) @@ -111,7 +111,7 @@ class MLflowDataTrace(BaseTraceInstance): """Simple dispatch to trace methods""" # Initialize on first use self._ensure_initialized() - + try: if isinstance(trace_info, WorkflowTraceInfo): self.workflow_trace(trace_info) diff --git a/api/tests/unit_tests/core/ops/test_ops_trace_manager.py b/api/tests/unit_tests/core/ops/test_ops_trace_manager.py index 351b8a01a2..886bc188f0 100644 --- a/api/tests/unit_tests/core/ops/test_ops_trace_manager.py +++ b/api/tests/unit_tests/core/ops/test_ops_trace_manager.py @@ -519,17 +519,15 @@ def test_extract_streaming_metrics_invalid_json(): def test_trace_queue_manager_add_and_collect(monkeypatch): # With lazy initialization, get_ops_trace_instance should NOT be called during __init__ mock_get_instance = MagicMock(return_value=True) - monkeypatch.setattr( - "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", mock_get_instance - ) + monkeypatch.setattr("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", mock_get_instance) manager = TraceQueueManager(app_id="app-id", user_id="user") # get_ops_trace_instance should NOT be called during initialization (lazy init) mock_get_instance.assert_not_called() - + # When trace_instance property is accessed, it should call get_ops_trace_instance _ = manager.trace_instance mock_get_instance.assert_called_once_with("app-id") - + task = TraceTask(trace_type=TraceTaskName.CONVERSATION_TRACE) manager.add_trace_task(task) tasks = manager.collect_tasks() @@ -539,14 +537,12 @@ def test_trace_queue_manager_add_and_collect(monkeypatch): def test_trace_queue_manager_lazy_init_error_handling(monkeypatch): """Test that TraceQueueManager handles initialization errors gracefully.""" mock_get_instance = MagicMock(side_effect=ConnectionError("Service unavailable")) - monkeypatch.setattr( - "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", mock_get_instance - ) - + monkeypatch.setattr("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", mock_get_instance) + manager = TraceQueueManager(app_id="app-id", user_id="user") # get_ops_trace_instance should NOT be called during initialization mock_get_instance.assert_not_called() - + # When trace_instance is accessed, it should handle the error gracefully instance = manager.trace_instance mock_get_instance.assert_called_once_with("app-id")