star

venv/lib/python3.9/site-packages/mlflow/tracing/export/async_export_queue.py

@@ -0,0 +1,176 @@
+import atexit
+import logging
+import threading
+from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
+from dataclasses import dataclass
+from queue import Empty, Queue
+from queue import Full as queue_Full
+from typing import Callable, Sequence
+
+from mlflow.environment_variables import (
+    MLFLOW_ASYNC_TRACE_LOGGING_MAX_QUEUE_SIZE,
+    MLFLOW_ASYNC_TRACE_LOGGING_MAX_WORKERS,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Task:
+    """A dataclass to represent a simple task."""
+
+    handler: Callable
+    args: Sequence
+    error_msg: str = ""
+
+    def handle(self) -> None:
+        """Handle the task execution. This method must not raise any exception."""
+        try:
+            self.handler(*self.args)
+        except Exception as e:
+            _logger.warning(
+                f"{self.error_msg} Error: {e}.",
+                exc_info=_logger.isEnabledFor(logging.DEBUG),
+            )
+
+
+class AsyncTraceExportQueue:
+    """A queue-based asynchronous tracing export processor."""
+
+    def __init__(self):
+        self._queue: Queue[Task] = Queue(maxsize=MLFLOW_ASYNC_TRACE_LOGGING_MAX_QUEUE_SIZE.get())
+        self._lock = threading.RLock()
+        self._max_workers = MLFLOW_ASYNC_TRACE_LOGGING_MAX_WORKERS.get()
+
+        # Thread event that indicates the queue should stop processing tasks
+        self._stop_event = threading.Event()
+        self._is_active = False
+        self._atexit_callback_registered = False
+
+        self._active_tasks = set()
+
+    def put(self, task: Task):
+        """Put a new task to the queue for processing."""
+        if not self.is_active():
+            self.activate()
+
+        # If stop event is set, wait for the queue to be drained before putting the task
+        if self._stop_event.is_set():
+            self._stop_event.wait()
+
+        try:
+            # Do not block if the queue is full, it will block the main application
+            self._queue.put(task, block=False)
+        except queue_Full:
+            _logger.warning(
+                "Trace export queue is full, trace will be discarded. "
+                "Consider increasing the queue size or number of workers."
+            )
+
+    def _consumer_loop(self) -> None:
+        while not self._stop_event.is_set():
+            self._dispatch_task()
+
+        # Drain remaining tasks when stopping
+        while not self._queue.empty():
+            self._dispatch_task()
+
+    def _dispatch_task(self) -> None:
+        """Dispatch a task from the queue to the worker thread pool."""
+        # NB: Monitor number of active tasks being processed by the workers. If the all
+        #   workers are busy, wait for one of them to finish before draining a new task
+        #   from the queue. This is because ThreadPoolExecutor does not have a built-in
+        #   mechanism to limit the number of pending tasks in the internal queue.
+        #   This ruins the purpose of having a size bound for self._queue, because the
+        #   TPE's internal queue can grow indefinitely and potentially run out of memory.
+        #   Therefore, we should only dispatch a new task when there is a worker available,
+        #   and pend the new tasks in the self._queue which has a size bound.
+        if len(self._active_tasks) >= self._max_workers:
+            _, self._active_tasks = wait(self._active_tasks, return_when=FIRST_COMPLETED)
+
+        try:
+            task = self._queue.get(timeout=1)
+        except Empty:
+            return
+
+        def _handle(task):
+            task.handle()
+            self._queue.task_done()
+
+        try:
+            future = self._worker_threadpool.submit(_handle, task)
+            self._active_tasks.add(future)
+        except Exception as e:
+            # In case it fails to submit the task to the worker thread pool
+            # such as interpreter shutdown, handle the task in this thread
+            _logger.debug(
+                f"Failed to submit task to worker thread pool. Error: {e}",
+                exc_info=True,
+            )
+            _handle(task)
+
+    def activate(self) -> None:
+        """Activate the async queue to accept and handle incoming tasks."""
+        with self._lock:
+            if self._is_active:
+                return
+
+            self._set_up_threads()
+
+            # Callback to ensure remaining tasks are processed before program exit
+            if not self._atexit_callback_registered:
+                atexit.register(self._at_exit_callback)
+                self._atexit_callback_registered = True
+
+            self._is_active = True
+
+    def is_active(self) -> bool:
+        return self._is_active
+
+    def _set_up_threads(self) -> None:
+        """Set up the consumer and worker threads."""
+        with self._lock:
+            self._worker_threadpool = ThreadPoolExecutor(
+                max_workers=self._max_workers,
+                thread_name_prefix="MlflowTraceLoggingWorker",
+            )
+            self._consumer_thread = threading.Thread(
+                target=self._consumer_loop,
+                name="MLflowTraceLoggingConsumer",
+                daemon=True,
+            )
+            self._consumer_thread.start()
+
+    def _at_exit_callback(self) -> None:
+        """Callback function executed when the program is exiting."""
+        try:
+            _logger.info(
+                "Flushing the async trace logging queue before program exit. "
+                "This may take a while..."
+            )
+            self.flush(terminate=True)
+        except Exception as e:
+            _logger.error(f"Error while finishing trace export requests: {e}")
+
+    def flush(self, terminate=False) -> None:
+        """
+        Flush the async logging queue.
+
+        Args:
+            terminate: If True, shut down the logging threads after flushing.
+        """
+        if not self.is_active():
+            return
+
+        self._stop_event.set()
+        self._consumer_thread.join()
+
+        # Wait for all tasks to be processed
+        self._queue.join()
+
+        self._worker_threadpool.shutdown(wait=True)
+        self._is_active = False
+        # Restart threads to listen to incoming requests after flushing, if not terminating
+        if not terminate:
+            self._stop_event.clear()
+            self.activate()

venv/lib/python3.9/site-packages/mlflow/tracing/export/databricks.py

@@ -0,0 +1,94 @@
+import logging
+from typing import Sequence
+
+from google.protobuf.json_format import MessageToDict
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter
+
+from mlflow.entities.trace import Trace
+from mlflow.environment_variables import (
+    MLFLOW_ASYNC_TRACE_LOGGING_RETRY_TIMEOUT,
+    MLFLOW_ENABLE_ASYNC_TRACE_LOGGING,
+)
+from mlflow.protos.databricks_trace_server_pb2 import CreateTrace, DatabricksTracingServerService
+from mlflow.tracing.export.async_export_queue import AsyncTraceExportQueue, Task
+from mlflow.tracing.fluent import _set_last_active_trace_id
+from mlflow.tracing.trace_manager import InMemoryTraceManager
+from mlflow.utils.databricks_utils import get_databricks_host_creds
+from mlflow.utils.rest_utils import (
+    _REST_API_PATH_PREFIX,
+    extract_api_info_for_service,
+    http_request,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+_METHOD_TO_INFO = extract_api_info_for_service(
+    DatabricksTracingServerService, _REST_API_PATH_PREFIX
+)
+
+
+class DatabricksSpanExporter(SpanExporter):
+    """
+    An exporter implementation that logs the traces to Databricks Tracing Server.
+    """
+
+    def __init__(self):
+        self._is_async = MLFLOW_ENABLE_ASYNC_TRACE_LOGGING.get()
+        if self._is_async:
+            _logger.info("MLflow is configured to log traces asynchronously.")
+            self._async_queue = AsyncTraceExportQueue()
+
+    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
+
+            trace = InMemoryTraceManager.get_instance().pop_trace(span.context.trace_id)
+            if trace is None:
+                _logger.debug(f"Trace for span {span} not found. Skipping export.")
+                continue
+
+            _set_last_active_trace_id(trace.info.request_id)
+
+            if self._is_async:
+                self._async_queue.put(
+                    task=Task(
+                        handler=self._log_trace,
+                        args=(trace,),
+                        error_msg="Failed to log trace to the trace server.",
+                    )
+                )
+            else:
+                self._log_trace(trace)
+
+    def _log_trace(self, trace: Trace):
+        """Create a new Trace record in the Databricks Tracing Server."""
+        request_body = MessageToDict(trace.to_proto(), preserving_proto_field_name=True)
+        endpoint, method = _METHOD_TO_INFO[CreateTrace]
+
+        # NB: Using Databricks SDK's built-in retry logic, which simply retries until the timeout
+        #    is reached, with linearly increasing backoff. Since it doesn't expose additional
+        #    configuration options, we might want to implement our own retry logic in the future.
+        # NB: If async logging is disabled, we don't retry to avoid blocking the application.
+        timeout = MLFLOW_ASYNC_TRACE_LOGGING_RETRY_TIMEOUT.get() if self._is_async else 0
+
+        # Use context manager to ensure the request is closed properly
+        with http_request(
+            host_creds=get_databricks_host_creds(),
+            endpoint=endpoint,
+            method=method,
+            json=request_body,
+            retry_timeout_seconds=timeout,
+        ) as res:
+            if res.status_code != 200:
+                _logger.warning(f"Failed to log trace to the trace server. Response: {res.text}")

venv/lib/python3.9/site-packages/mlflow/tracing/export/databricks_agent_legacy.py

@@ -0,0 +1,53 @@
+import logging
+from typing import Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter
+
+from mlflow.deployments import get_deploy_client
+from mlflow.tracing.destination import TraceDestination
+from mlflow.tracing.trace_manager import InMemoryTraceManager
+
+_logger = logging.getLogger(__name__)
+
+
+class DatabricksAgentSpanExporter(SpanExporter):
+    """
+    An exporter implementation that logs the traces to Databricks Agent Monitoring.
+
+    Args:
+        trace_destination: The destination of the traces.
+
+    TODO: This class should be deprecated in favor of DatabricksSpanExporter, once
+        the Databricks Agent Monitoring is fully migrated to the new tracing server.
+    """
+
+    def __init__(self, trace_destination: TraceDestination):
+        self._databricks_monitor_id = trace_destination.databricks_monitor_id
+        self._trace_manager = InMemoryTraceManager.get_instance()
+        self._deploy_client = get_deploy_client("databricks")
+
+    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
+
+            trace = self._trace_manager.pop_trace(span.context.trace_id)
+            if trace is None:
+                _logger.debug(f"Trace for span {span} not found. Skipping export.")
+                continue
+
+            # Traces are exported via a serving endpoint that accepts trace JSON as
+            # an input payload, and then will be written to the Inference Table.
+            self._deploy_client.predict(
+                endpoint=self._databricks_monitor_id,
+                inputs={"inputs": [trace.to_json()]},
+            )

venv/lib/python3.9/site-packages/mlflow/tracing/export/inference_table.py

@@ -0,0 +1,74 @@
+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.environment_variables import (
+    MLFLOW_TRACE_BUFFER_MAX_SIZE,
+    MLFLOW_TRACE_BUFFER_TTL_SECONDS,
+)
+from mlflow.tracing.fluent import _set_last_active_trace_id
+from mlflow.tracing.trace_manager import InMemoryTraceManager
+
+_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()
+
+    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
+
+            trace = self._trace_manager.pop_trace(span.context.trace_id)
+            if trace is None:
+                _logger.debug(f"Trace for span {span} not found. Skipping export.")
+                continue
+
+            _set_last_active_trace_id(trace.info.request_id)
+
+            # Add the trace to the in-memory buffer so it can be retrieved by upstream
+            _TRACE_BUFFER[trace.info.request_id] = trace.to_dict()

venv/lib/python3.9/site-packages/mlflow/tracing/export/mlflow.py

@@ -0,0 +1,100 @@
+import logging
+from typing import Optional, Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter
+
+from mlflow.entities.trace import Trace
+from mlflow.environment_variables import MLFLOW_ENABLE_ASYNC_LOGGING
+from mlflow.tracing.constant import TraceTagKey
+from mlflow.tracing.display import get_display_handler
+from mlflow.tracing.display.display_handler import IPythonTraceDisplayHandler
+from mlflow.tracing.export.async_export_queue import AsyncTraceExportQueue, Task
+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 maybe_get_request_id
+from mlflow.tracking.client import MlflowClient
+
+_logger = logging.getLogger(__name__)
+
+
+class MlflowSpanExporter(SpanExporter):
+    """
+    An exporter implementation that logs the traces to MLflow.
+
+    MLflow backend (will) only support logging the complete trace, not incremental updates
+    for spans, so this exporter is designed to aggregate the spans into traces in memory.
+    Therefore, this only works within a single process application and not intended to work
+    in a distributed environment. For the same reason, this exporter should only be used with
+    SimpleSpanProcessor.
+
+    If we want to support distributed tracing, we should first implement an incremental trace
+    logging in MLflow backend, then we can get rid of the in-memory trace aggregation.
+
+    :meta private:
+    """
+
+    def __init__(
+        self,
+        client: Optional[MlflowClient] = None,
+        display_handler: Optional[IPythonTraceDisplayHandler] = None,
+    ):
+        self._client = client or MlflowClient()
+        self._display_handler = display_handler or get_display_handler()
+        self._trace_manager = InMemoryTraceManager.get_instance()
+        self._async_queue = AsyncTraceExportQueue()
+
+    def export(self, spans: Sequence[ReadableSpan]):
+        """
+        Export the spans to MLflow backend.
+
+        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
+
+            trace = self._trace_manager.pop_trace(span.context.trace_id)
+            if trace is None:
+                _logger.debug(f"TraceInfo for span {span} not found. Skipping export.")
+                continue
+
+            _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.request_id
+
+            if not maybe_get_request_id(is_evaluate=True):
+                # Display the trace in the UI if the trace is not generated from within
+                # an MLflow model evaluation context
+                self._display_handler.display_traces([trace])
+
+            self._log_trace(trace)
+
+    def _log_trace(self, trace: Trace):
+        """Log the trace to MLflow backend."""
+        upload_trace_data_task = Task(
+            handler=self._client._upload_trace_data,
+            args=(trace.info, trace.data),
+            error_msg="Failed to log trace to MLflow backend.",
+        )
+
+        upload_ended_trace_info_task = Task(
+            handler=self._client._upload_ended_trace_info,
+            args=(trace.info,),
+            error_msg="Failed to log trace to MLflow backend.",
+        )
+
+        # TODO: Use MLFLOW_ENABLE_ASYNC_TRACE_LOGGING instead and default to async
+        # logging once the async logging implementation becomes stable.
+        if MLFLOW_ENABLE_ASYNC_LOGGING.get():
+            self._async_queue.put(upload_trace_data_task)
+            self._async_queue.put(upload_ended_trace_info_task)
+        else:
+            upload_trace_data_task.handle()
+            upload_ended_trace_info_task.handle()