genesis-flow 1.0.0__py3-none-any.whl

mlflow/tracing/constant.py

@@ -0,0 +1,105 @@
+# NB: These keys are placeholders and subject to change
+class TraceMetadataKey:
+    INPUTS = "mlflow.traceInputs"
+    OUTPUTS = "mlflow.traceOutputs"
+    SOURCE_RUN = "mlflow.sourceRun"
+    MODEL_ID = "mlflow.modelId"
+    # Trace size statistics including total size, number of spans, and max span size
+    SIZE_STATS = "mlflow.trace.sizeStats"
+    # Aggregated token usage information in a single trace, stored as a dumped JSON string.
+    TOKEN_USAGE = "mlflow.trace.tokenUsage"
+    # Store the user ID/name of the application request. Do not confuse this with mlflow.user
+    # tag, which stores "who created the trace" i.e. developer or system name.
+    TRACE_USER = "mlflow.trace.user"
+    # Store the session ID of the application request.
+    TRACE_SESSION = "mlflow.trace.session"
+
+    # Total size of the trace in bytes. Deprecated, use SIZE_STATS instead.
+    SIZE_BYTES = "mlflow.trace.sizeBytes"
+
+
+class TraceTagKey:
+    TRACE_NAME = "mlflow.traceName"
+    EVAL_REQUEST_ID = "eval.requestId"
+
+
+class TokenUsageKey:
+    """Key for the token usage information in the `mlflow.chat.tokenUsage` span attribute."""
+
+    INPUT_TOKENS = "input_tokens"
+    OUTPUT_TOKENS = "output_tokens"
+    TOTAL_TOKENS = "total_tokens"
+
+    @classmethod
+    def all_keys(cls):
+        return [cls.INPUT_TOKENS, cls.OUTPUT_TOKENS, cls.TOTAL_TOKENS]
+
+
+class TraceSizeStatsKey:
+    TOTAL_SIZE_BYTES = "total_size_bytes"
+    NUM_SPANS = "num_spans"
+    MAX_SPAN_SIZE_BYTES = "max"
+    P25_SPAN_SIZE_BYTES = "p25"
+    P50_SPAN_SIZE_BYTES = "p50"
+    P75_SPAN_SIZE_BYTES = "p75"
+
+
+# A set of reserved attribute keys
+class SpanAttributeKey:
+    EXPERIMENT_ID = "mlflow.experimentId"
+    REQUEST_ID = "mlflow.traceRequestId"
+    INPUTS = "mlflow.spanInputs"
+    OUTPUTS = "mlflow.spanOutputs"
+    SPAN_TYPE = "mlflow.spanType"
+    FUNCTION_NAME = "mlflow.spanFunctionName"
+    START_TIME_NS = "mlflow.spanStartTimeNs"
+    # these attributes are for standardized chat messages and tool definitions
+    # in CHAT_MODEL and LLM spans. they are used for rendering the rich chat
+    # display in the trace UI, as well as downstream consumers of trace data
+    # such as evaluation
+    CHAT_MESSAGES = "mlflow.chat.messages"
+    CHAT_TOOLS = "mlflow.chat.tools"
+    # This attribute is used to store token usage information from LLM responses.
+    # Stored in {"input_tokens": int, "output_tokens": int, "total_tokens": int} format.
+    CHAT_USAGE = "mlflow.chat.tokenUsage"
+    # This attribute is used to populate `intermediate_outputs` property of a trace data
+    # representing intermediate outputs of the trace. This attribute is not empty only on
+    # the root span of a trace created by the `mlflow.log_trace` API. The `intermediate_outputs`
+    # property of the normal trace is generated by the outputs of non-root spans.
+    INTERMEDIATE_OUTPUTS = "mlflow.trace.intermediate_outputs"
+
+
+class AssessmentMetadataKey:
+    SOURCE_RUN_ID = "mlflow.assessment.sourceRunId"
+
+
+# All storage backends are guaranteed to support request_metadata key/value up to 250 characters
+MAX_CHARS_IN_TRACE_INFO_METADATA = 250
+# All storage backends are guaranteed to support tag keys up to 250 characters,
+# values up to 4096 characters
+MAX_CHARS_IN_TRACE_INFO_TAGS_KEY = 250
+MAX_CHARS_IN_TRACE_INFO_TAGS_VALUE = 4096
+TRUNCATION_SUFFIX = "..."
+
+TRACE_REQUEST_RESPONSE_PREVIEW_MAX_LENGTH_DBX = 10000
+TRACE_REQUEST_RESPONSE_PREVIEW_MAX_LENGTH_OSS = 1000
+
+# Trace request ID must have the prefix "tr-" appended to the OpenTelemetry trace ID
+TRACE_REQUEST_ID_PREFIX = "tr-"
+
+# Schema version of traces and spans.
+TRACE_SCHEMA_VERSION = 3
+
+# Key for the trace schema version in the trace. This key is also used in
+# Databricks model serving to be careful when modifying it.
+TRACE_SCHEMA_VERSION_KEY = "mlflow.trace_schema.version"
+
+
+STREAM_CHUNK_EVENT_NAME_FORMAT = "mlflow.chunk.item.{index}"
+STREAM_CHUNK_EVENT_VALUE_KEY = "mlflow.chunk.value"
+
+
+# Key for Databricks model serving options to return the trace in the response
+DATABRICKS_OPTIONS_KEY = "databricks_options"
+RETURN_TRACE_OPTION_KEY = "return_trace"
+DATABRICKS_OUTPUT_KEY = "databricks_output"

mlflow/tracing/destination.py

@@ -0,0 +1,81 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from mlflow.exceptions import MlflowException
+from mlflow.utils.annotations import experimental
+
+
+@experimental(version="2.21.0")
+@dataclass
+class TraceDestination:
+    """A configuration object for specifying the destination of trace data."""
+
+    @property
+    def type(self) -> str:
+        """Type of the destination."""
+        raise NotImplementedError
+
+
+@experimental(version="2.21.0")
+@dataclass
+class MlflowExperiment(TraceDestination):
+    """
+    A destination representing an MLflow experiment.
+
+    By setting this destination in the :py:func:`mlflow.tracing.set_destination` function,
+    MLflow will log traces to the specified experiment.
+
+    Attributes:
+        experiment_id: The ID of the experiment to log traces to. If not specified,
+            the current active experiment will be used.
+        tracking_uri: The tracking URI of the MLflow server to log traces to.
+            If not specified, the current tracking URI will be used.
+    """
+
+    experiment_id: Optional[str] = None
+    tracking_uri: Optional[str] = None
+
+    @property
+    def type(self) -> str:
+        return "experiment"
+
+
+@experimental(version="2.22.0")
+@dataclass
+class Databricks(TraceDestination):
+    """
+    A destination representing a Databricks tracing server.
+
+    By setting this destination in the :py:func:`mlflow.tracing.set_destination` function,
+    MLflow will log traces to the specified experiment.
+
+    If neither experiment_id nor experiment_name is specified, an active experiment
+    when traces are created will be used as the destination.
+    If both are specified, they must refer to the same experiment.
+
+    Attributes:
+        experiment_id: The ID of the experiment to log traces to.
+        experiment_name: The name of the experiment to log traces to.
+    """
+
+    experiment_id: Optional[str] = None
+    experiment_name: Optional[str] = None
+
+    def __post_init__(self):
+        if self.experiment_id is not None:
+            self.experiment_id = str(self.experiment_id)
+
+        if self.experiment_name is not None:
+            from mlflow.tracking._tracking_service.utils import _get_store
+
+            # NB: Use store directly rather than fluent API to avoid dependency on MLflowClient
+            experiment_id = _get_store().get_experiment_by_name(self.experiment_name).experiment_id
+            if self.experiment_id is not None and self.experiment_id != experiment_id:
+                raise MlflowException.invalid_parameter_value(
+                    "experiment_id and experiment_name must refer to the same experiment"
+                )
+            self.experiment_id = experiment_id
+
+    @property
+    def type(self) -> str:
+        return "databricks"

mlflow/tracing/display/__init__.py

@@ -0,0 +1,40 @@
+from mlflow.tracing.display.display_handler import (
+    IPythonTraceDisplayHandler,
+    get_notebook_iframe_html,
+    is_using_tracking_server,
+)
+
+__all__ = [
+    "IPythonTraceDisplayHandler",
+    "get_display_handler",
+    "is_using_tracking_server",
+    "get_notebook_iframe_html",
+]
+
+
+def get_display_handler() -> IPythonTraceDisplayHandler:
+    return IPythonTraceDisplayHandler.get_instance()
+
+
+def disable_notebook_display():
+    """
+    Disables displaying the MLflow Trace UI in notebook output cells.
+    Call :py:func:`mlflow.tracing.enable_notebook_display()` to re-enable display.
+    """
+    IPythonTraceDisplayHandler.disable()
+
+
+def enable_notebook_display():
+    """
+    Enables the MLflow Trace UI in notebook output cells. The display is on
+    by default, and the Trace UI will show up when any of the following operations
+    are executed:
+
+    * On trace completion (i.e. whenever a trace is exported)
+    * When calling the :py:func:`mlflow.search_traces` fluent API
+    * When calling the :py:meth:`mlflow.client.MlflowClient.get_trace`
+      or :py:meth:`mlflow.client.MlflowClient.search_traces` client APIs
+
+    To disable, please call :py:func:`mlflow.tracing.disable_notebook_display()`.
+    """
+    IPythonTraceDisplayHandler.enable()

mlflow/tracing/display/display_handler.py

@@ -0,0 +1,196 @@
+import html
+import json
+import logging
+from typing import TYPE_CHECKING
+from urllib.parse import urlencode, urljoin
+
+import mlflow
+from mlflow.environment_variables import MLFLOW_MAX_TRACES_TO_DISPLAY_IN_NOTEBOOK
+from mlflow.utils.databricks_utils import is_in_databricks_runtime
+from mlflow.utils.uri import is_http_uri
+
+_logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from mlflow.entities import Trace
+
+
+TRACE_RENDERER_ASSET_PATH = "/static-files/lib/notebook-trace-renderer/index.html"
+
+IFRAME_HTML = """
+<div>
+  <style scoped>
+  button {{
+    border: none;
+    border-radius: 4px;
+    background-color: rgb(34, 114, 180);
+    font-family: -apple-system, "system-ui", "Segoe UI", Roboto, "Helvetica Neue", Arial;
+    font-size: 13px;
+    color: white;
+    margin-top: 8px;
+    margin-bottom: 8px;
+    padding: 8px 16px;
+    cursor: pointer;
+  }}
+  button:hover {{
+    background-color: rgb(66, 153, 224);
+  }}
+  </style>
+  <button
+    onclick="
+        const display = this.nextElementSibling.style.display;
+        const isCollapsed = display === 'none';
+        this.nextElementSibling.style.display = isCollapsed ? null : 'none';
+
+        const verb = isCollapsed ? 'Collapse' : 'Expand';
+        this.innerText = `${{verb}} MLflow Trace`;
+    "
+  >Collapse MLflow Trace</button>
+  <iframe
+    id="trace-renderer"
+    style="width: 100%; height: 500px; border: none; resize: vertical;"
+    src="{src}"
+  />
+</div>
+"""
+
+
+def get_notebook_iframe_html(traces: list["Trace"]):
+    # fetch assets from tracking server
+    uri = urljoin(mlflow.get_tracking_uri(), TRACE_RENDERER_ASSET_PATH)
+    query_string = _get_query_string_for_traces(traces)
+
+    # include mlflow version to invalidate browser cache when mlflow updates
+    src = html.escape(f"{uri}?{query_string}&version={mlflow.__version__}")
+    return IFRAME_HTML.format(src=src)
+
+
+def _serialize_trace_list(traces: list["Trace"]):
+    return json.dumps(
+        # we can't just call trace.to_json() because this
+        # will cause the trace to be serialized twice (once
+        # by to_json and once by json.dumps)
+        [json.loads(trace._serialize_for_mimebundle()) for trace in traces],
+        ensure_ascii=False,
+    )
+
+
+def _get_query_string_for_traces(traces: list["Trace"]):
+    query_params = []
+
+    for trace in traces:
+        query_params.append(("trace_id", trace.info.request_id))
+        query_params.append(("experiment_id", trace.info.experiment_id))
+
+    return urlencode(query_params)
+
+
+def _is_jupyter():
+    try:
+        from IPython import get_ipython
+
+        return get_ipython() is not None
+    except ImportError:
+        return False
+
+
+def is_using_tracking_server():
+    return is_http_uri(mlflow.get_tracking_uri())
+
+
+def is_trace_ui_available():
+    # the notebook display feature only works in
+    # Databricks notebooks, or in Jupyter notebooks
+    # with a tracking server
+    return _is_jupyter() and (is_in_databricks_runtime() or is_using_tracking_server())
+
+
+class IPythonTraceDisplayHandler:
+    _instance = None
+    disabled = False
+
+    @classmethod
+    def get_instance(cls):
+        if cls._instance is None:
+            cls._instance = IPythonTraceDisplayHandler()
+        return cls._instance
+
+    @classmethod
+    def disable(cls):
+        cls.disabled = True
+
+    @classmethod
+    def enable(cls):
+        cls.disabled = False
+        if cls._instance is None:
+            cls._instance = IPythonTraceDisplayHandler()
+
+    def __init__(self):
+        self.traces_to_display = {}
+        if not _is_jupyter():
+            return
+
+        try:
+            from IPython import get_ipython
+
+            # Register a post-run cell display hook to display traces
+            # after the cell has executed. We don't validate that the
+            # user is using a tracking server at this step, because
+            # the user might set it later using mlflow.set_tracking_uri()
+            get_ipython().events.register("post_run_cell", self._display_traces_post_run)
+        except Exception:
+            # swallow exceptions. this function is called as
+            # a side-effect in a few other functions (e.g. log_trace,
+            # get_traces, search_traces), and we don't want to block
+            # the core functionality if the display fails.
+            _logger.debug("Failed to register post-run cell display hook", exc_info=True)
+
+    def _display_traces_post_run(self, result):
+        if self.disabled or not is_trace_ui_available():
+            self.traces_to_display = {}
+            return
+
+        try:
+            from IPython.display import display
+
+            MAX_TRACES_TO_DISPLAY = MLFLOW_MAX_TRACES_TO_DISPLAY_IN_NOTEBOOK.get()
+            traces_to_display = list(self.traces_to_display.values())[:MAX_TRACES_TO_DISPLAY]
+            if len(traces_to_display) == 0:
+                self.traces_to_display = {}
+                return
+
+            display(self.get_mimebundle(traces_to_display), raw=True)
+
+            # reset state
+            self.traces_to_display = {}
+        except Exception:
+            # swallow exceptions. this function is called as
+            # a side-effect in a few other functions (e.g. log_trace,
+            # get_traces, search_traces), and we don't want to block
+            # the core functionality if the display fails.
+            _logger.error("Failed to display traces", exc_info=True)
+            self.traces_to_display = {}
+
+    def get_mimebundle(self, traces: list["Trace"]):
+        if len(traces) == 1:
+            return traces[0]._repr_mimebundle_()
+        else:
+            bundle = {"text/plain": repr(traces)}
+            if is_in_databricks_runtime():
+                bundle["application/databricks.mlflow.trace"] = _serialize_trace_list(traces)
+            else:
+                bundle["text/html"] = get_notebook_iframe_html(traces)
+            return bundle
+
+    def display_traces(self, traces: list["Trace"]):
+        if self.disabled or not is_trace_ui_available():
+            return
+
+        try:
+            if len(traces) == 0:
+                return
+
+            traces_dict = {trace.info.request_id: trace for trace in traces}
+            self.traces_to_display.update(traces_dict)
+        except Exception:
+            _logger.debug("Failed to update traces", exc_info=True)

mlflow/tracing/export/async_export_queue.py

@@ -0,0 +1,186 @@
+import atexit
+import logging
+import threading
+import time
+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 Any, 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[..., Any]
+    args: Sequence[Any]
+    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()
+
+        self._last_full_queue_warning_time = None
+
+    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:
+            if self._last_full_queue_warning_time is None or (
+                time.time() - self._last_full_queue_warning_time > 30
+            ):
+                _logger.warning(
+                    "Trace export queue is full, trace will be discarded. "
+                    "Consider increasing the queue size through "
+                    "`MLFLOW_ASYNC_TRACE_LOGGING_MAX_QUEUE_SIZE` environment variable or "
+                    "number of workers through `MLFLOW_ASYNC_TRACE_LOGGING_MAX_WORKERS`"
+                    " environment variable."
+                )
+                self._last_full_queue_warning_time = time.time()
+
+    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()