genesis-flow 1.0.0__py3-none-any.whl

mlflow/tracing/export/inference_table.py

@@ -0,0 +1,138 @@
+import logging
+from typing import Any, Optional, Sequence
+
+from cachetools import TTLCache
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter
+
+from mlflow.entities.model_registry import PromptVersion
+from mlflow.entities.trace import Trace
+from mlflow.environment_variables import (
+    _MLFLOW_ENABLE_TRACE_DUAL_WRITE_IN_MODEL_SERVING,
+    MLFLOW_TRACE_BUFFER_MAX_SIZE,
+    MLFLOW_TRACE_BUFFER_TTL_SECONDS,
+)
+from mlflow.tracing.client import TracingClient
+from mlflow.tracing.export.async_export_queue import AsyncTraceExportQueue, Task
+from mlflow.tracing.export.utils import try_link_prompts_to_trace
+from mlflow.tracing.fluent import _set_last_active_trace_id
+from mlflow.tracing.trace_manager import InMemoryTraceManager
+from mlflow.tracing.utils import add_size_stats_to_trace_metadata
+
+_logger = logging.getLogger(__name__)
+
+
+def pop_trace(request_id: str) -> Optional[dict[str, Any]]:
+    """
+    Pop the completed trace data from the buffer. This method is used in
+    the Databricks model serving so please be careful when modifying it.
+    """
+    return _TRACE_BUFFER.pop(request_id, None)
+
+
+# For Inference Table, we use special TTLCache to store the finished traces
+# so that they can be retrieved by Databricks model serving. The values
+# in the buffer are not Trace dataclass, but rather a dictionary with the schema
+# that is used within Databricks model serving.
+def _initialize_trace_buffer():  # Define as a function for testing purposes
+    return TTLCache(
+        maxsize=MLFLOW_TRACE_BUFFER_MAX_SIZE.get(),
+        ttl=MLFLOW_TRACE_BUFFER_TTL_SECONDS.get(),
+    )
+
+
+_TRACE_BUFFER = _initialize_trace_buffer()
+
+
+class InferenceTableSpanExporter(SpanExporter):
+    """
+    An exporter implementation that logs the traces to Inference Table.
+
+    Currently the Inference Table does not use collector to receive the traces,
+    but rather actively fetches the trace during the prediction process. In the
+    future, we may consider using collector-based approach and this exporter should
+    send the traces instead of storing them in the buffer.
+    """
+
+    def __init__(self):
+        self._trace_manager = InMemoryTraceManager.get_instance()
+
+        # NB: When this env var is set to true, MLflow will export traces to both inference
+        #  table and the Databricks Tracing Server.
+        self._should_write_to_mlflow_backend = (
+            _MLFLOW_ENABLE_TRACE_DUAL_WRITE_IN_MODEL_SERVING.get()
+        )
+        if self._should_write_to_mlflow_backend:
+            self._client = TracingClient("databricks")
+            self._async_queue = AsyncTraceExportQueue()
+
+    def export(self, spans: Sequence[ReadableSpan]):
+        """
+        Export the spans to Inference Table via the TTLCache buffer.
+
+        Args:
+            spans: A sequence of OpenTelemetry ReadableSpan objects passed from
+                a span processor. Only root spans for each trace should be exported.
+        """
+        for span in spans:
+            if span._parent is not None:
+                _logger.debug("Received a non-root span. Skipping export.")
+                continue
+
+            manager_trace = self._trace_manager.pop_trace(span.context.trace_id)
+            if manager_trace is None:
+                _logger.debug(f"Trace for span {span} not found. Skipping export.")
+                continue
+
+            trace = manager_trace.trace
+            _set_last_active_trace_id(trace.info.trace_id)
+
+            # Add the trace to the in-memory buffer so it can be retrieved by upstream
+            # The key is Databricks request ID.
+            _TRACE_BUFFER[trace.info.client_request_id] = trace.to_dict()
+
+            if self._should_write_to_mlflow_backend:
+                if trace.info.experiment_id is None:
+                    # NB: The experiment ID is set based on the MLFLOW_EXPERIMENT_ID env var
+                    #   populated in the scoring server by Agent Framework. If the model is not
+                    #   deployed via agents.deploy(), the env var will not be set and the
+                    #   experiment will be empty, even if the dual write itself is enabled.
+                    _logger.warning(
+                        "Dual write to MLflow backend is enabled, but experiment ID is not set "
+                        "for the trace. Skipping trace export to MLflow backend."
+                    )
+                    continue
+
+                try:
+                    # Log the trace to the MLflow backend asynchronously
+                    self._async_queue.put(
+                        task=Task(
+                            handler=self._log_trace_to_mlflow_backend,
+                            args=(trace, manager_trace.prompts),
+                            error_msg=f"Failed to log trace {trace.info.trace_id}.",
+                        )
+                    )
+                except Exception as e:
+                    _logger.warning(
+                        f"Failed to export trace to MLflow backend. Error: {e}",
+                        stack_info=_logger.isEnabledFor(logging.DEBUG),
+                    )
+
+    def _log_trace_to_mlflow_backend(self, trace: Trace, prompts: Sequence[PromptVersion]):
+        add_size_stats_to_trace_metadata(trace)
+
+        returned_trace_info = self._client.start_trace(trace.info)
+        self._client._upload_trace_data(returned_trace_info, trace.data)
+
+        # Link prompt versions to the trace. Prompt linking is not critical for trace export
+        # (if the prompt fails to link, the user's workflow is minorly affected), so we handle
+        # errors gracefully without failing the entire trace export
+        try_link_prompts_to_trace(
+            client=self._client,
+            trace_id=returned_trace_info.trace_id,
+            prompts=prompts,
+            synchronous=True,  # Run synchronously since we're already in an async task
+        )
+        _logger.debug(
+            f"Finished logging trace to MLflow backend. TraceInfo: {returned_trace_info.to_dict()} "
+        )

mlflow/tracing/export/mlflow_v3.py

@@ -0,0 +1,137 @@
+import logging
+from typing import Optional, Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter
+
+from mlflow.entities.model_registry import PromptVersion
+from mlflow.entities.trace import Trace
+from mlflow.environment_variables import (
+    MLFLOW_ENABLE_ASYNC_TRACE_LOGGING,
+)
+from mlflow.tracing.client import TracingClient
+from mlflow.tracing.constant import TraceTagKey
+from mlflow.tracing.display import get_display_handler
+from mlflow.tracing.export.async_export_queue import AsyncTraceExportQueue, Task
+from mlflow.tracing.export.utils import try_link_prompts_to_trace
+from mlflow.tracing.fluent import _EVAL_REQUEST_ID_TO_TRACE_ID, _set_last_active_trace_id
+from mlflow.tracing.trace_manager import InMemoryTraceManager
+from mlflow.tracing.utils import add_size_stats_to_trace_metadata, maybe_get_request_id
+from mlflow.utils.databricks_utils import is_in_databricks_notebook
+from mlflow.utils.uri import is_databricks_uri
+
+_logger = logging.getLogger(__name__)
+
+
+class MlflowV3SpanExporter(SpanExporter):
+    """
+    An exporter implementation that logs the traces to MLflow Tracking Server
+    using the V3 trace schema and API.
+    """
+
+    def __init__(self, tracking_uri: Optional[str] = None):
+        self._client = TracingClient(tracking_uri)
+        self._is_async_enabled = self._should_enable_async_logging()
+        if self._is_async_enabled:
+            self._async_queue = AsyncTraceExportQueue()
+
+        # Display handler is no-op when running outside of notebooks.
+        self._display_handler = get_display_handler()
+
+    def export(self, spans: Sequence[ReadableSpan]):
+        """
+        Export the spans to the destination.
+
+        Args:
+            spans: A sequence of OpenTelemetry ReadableSpan objects passed from
+                a span processor. Only root spans for each trace should be exported.
+        """
+        for span in spans:
+            if span._parent is not None:
+                _logger.debug("Received a non-root span. Skipping export.")
+                continue
+
+            manager_trace = InMemoryTraceManager.get_instance().pop_trace(span.context.trace_id)
+            if manager_trace is None:
+                _logger.debug(f"Trace for span {span} not found. Skipping export.")
+                continue
+
+            trace = manager_trace.trace
+            _set_last_active_trace_id(trace.info.request_id)
+
+            # Store mapping from eval request ID to trace ID so that the evaluation
+            # harness can access to the trace using mlflow.get_trace(eval_request_id)
+            if eval_request_id := trace.info.tags.get(TraceTagKey.EVAL_REQUEST_ID):
+                _EVAL_REQUEST_ID_TO_TRACE_ID[eval_request_id] = trace.info.trace_id
+
+            if not maybe_get_request_id(is_evaluate=True):
+                self._display_handler.display_traces([trace])
+
+            if self._should_log_async():
+                self._async_queue.put(
+                    task=Task(
+                        handler=self._log_trace,
+                        args=(trace, manager_trace.prompts),
+                        error_msg="Failed to log trace to the trace server.",
+                    )
+                )
+            else:
+                self._log_trace(trace, prompts=manager_trace.prompts)
+
+    def _log_trace(self, trace: Trace, prompts: Sequence[PromptVersion]):
+        """
+        Handles exporting a trace to MLflow using the V3 API and blob storage.
+        Steps:
+        1. Create the trace in MLflow
+        2. Upload the trace data to blob storage using the returned trace info.
+        """
+        try:
+            if trace:
+                add_size_stats_to_trace_metadata(trace)
+                returned_trace_info = self._client.start_trace(trace.info)
+                self._client._upload_trace_data(returned_trace_info, trace.data)
+            else:
+                _logger.warning("No trace or trace info provided, unable to export")
+        except Exception as e:
+            _logger.warning(f"Failed to send trace to MLflow backend: {e}")
+
+        try:
+            # Always run prompt linking asynchronously since (1) prompt linking API calls
+            # would otherwise add latency to the export procedure and (2) prompt linking is not
+            # critical for trace export (if the prompt fails to link, the user's workflow is
+            # minorly affected), so we don't have to await successful linking
+            try_link_prompts_to_trace(
+                client=self._client,
+                trace_id=trace.info.trace_id,
+                prompts=prompts,
+                synchronous=False,
+            )
+        except Exception as e:
+            _logger.warning(f"Failed to link prompts to trace: {e}")
+
+    def _should_enable_async_logging(self):
+        if (
+            is_in_databricks_notebook()
+            # NB: Not defaulting OSS backend to async logging for now to reduce blast radius.
+            or not is_databricks_uri(self._client.tracking_uri)
+        ):
+            # NB: We don't turn on async logging in Databricks notebook by default
+            # until we are confident that the async logging is working on the
+            # offline workload on Databricks, to derisk the inclusion to the
+            # standard image. When it is enabled explicitly via the env var, we
+            # will respect that.
+            return (
+                MLFLOW_ENABLE_ASYNC_TRACE_LOGGING.get()
+                if MLFLOW_ENABLE_ASYNC_TRACE_LOGGING.is_set()
+                else False
+            )
+
+        return MLFLOW_ENABLE_ASYNC_TRACE_LOGGING.get()
+
+    def _should_log_async(self):
+        # During evaluate, the eval harness relies on the generated trace objects,
+        # so we should not log traces asynchronously.
+        if maybe_get_request_id(is_evaluate=True):
+            return False
+
+        return self._is_async_enabled

mlflow/tracing/export/utils.py

@@ -0,0 +1,70 @@
+"""
+Utility functions for prompt linking in trace exporters.
+"""
+
+import logging
+import threading
+import uuid
+from typing import Sequence
+
+from mlflow.entities.model_registry import PromptVersion
+from mlflow.tracing.client import TracingClient
+
+_logger = logging.getLogger(__name__)
+
+
+def try_link_prompts_to_trace(
+    client: TracingClient,
+    trace_id: str,
+    prompts: Sequence[PromptVersion],
+    synchronous: bool = True,
+) -> None:
+    """
+    Attempt to link prompt versions to a trace with graceful error handling.
+
+    This function provides a reusable way to link prompts to traces with consistent
+    error handling across different exporters. Errors are caught and logged but do
+    not propagate, ensuring that prompt linking failures don't affect trace export.
+
+    Args:
+        client: The TracingClient instance to use for linking.
+        trace_id: The ID of the trace to link prompts to.
+        prompts: Sequence of PromptVersion objects to link.
+        synchronous: If True, run the linking synchronously. If False, run in a separate thread.
+    """
+    if not prompts:
+        return
+
+    if synchronous:
+        _link_prompts_sync(client, trace_id, prompts)
+    else:
+        threading.Thread(
+            target=_link_prompts_sync,
+            args=(client, trace_id, prompts),
+            name=f"link_prompts_from_exporter-{uuid.uuid4().hex[:8]}",
+        ).start()
+
+
+def _link_prompts_sync(
+    client: TracingClient,
+    trace_id: str,
+    prompts: Sequence[PromptVersion],
+) -> None:
+    """
+    Synchronously link prompt versions to a trace with error handling.
+
+    This is the core implementation that handles the actual API call and error logging.
+
+    Args:
+        client: The TracingClient instance to use for linking.
+        trace_id: The ID of the trace to link prompts to.
+        prompts: Sequence of PromptVersion objects to link.
+    """
+    try:
+        client.link_prompt_versions_to_trace(
+            trace_id=trace_id,
+            prompts=prompts,
+        )
+        _logger.debug(f"Successfully linked {len(prompts)} prompts to trace {trace_id}")
+    except Exception as e:
+        _logger.warning(f"Failed to link prompts to trace {trace_id}: {e}")