genesis-flow 1.0.0__py3-none-any.whl

mlflow/tracing/utils/timeout.py

@@ -0,0 +1,250 @@
+import atexit
+import logging
+import threading
+import time
+from collections import OrderedDict
+
+from cachetools import Cache, TTLCache
+
+from mlflow.entities.span_event import SpanEvent
+from mlflow.entities.span_status import SpanStatusCode
+from mlflow.environment_variables import (
+    MLFLOW_TRACE_BUFFER_MAX_SIZE,
+    MLFLOW_TRACE_BUFFER_TTL_SECONDS,
+    MLFLOW_TRACE_TIMEOUT_CHECK_INTERVAL_SECONDS,
+    MLFLOW_TRACE_TIMEOUT_SECONDS,
+)
+from mlflow.exceptions import MlflowTracingException
+
+_logger = logging.getLogger(__name__)
+
+_TRACE_EXPIRATION_MSG = (
+    "Trace {request_id} is timed out after {ttl} seconds. The operation may be stuck or "
+    "taking too long to complete. To increase the timeout, set the environment variable "
+    "MLFLOW_TRACE_TIMEOUT_SECONDS to a larger value."
+)
+
+
+def get_trace_cache_with_timeout() -> Cache:
+    """
+    Return a cache object that stores traces in-memory while they are in-progress.
+
+    If the timeout is specified, this returns a customized cache that logs the
+    expired traces to the backend. Otherwise, this returns a regular cache.
+    """
+
+    if timeout := MLFLOW_TRACE_TIMEOUT_SECONDS.get():
+        return MlflowTraceTimeoutCache(
+            timeout=timeout,
+            maxsize=MLFLOW_TRACE_BUFFER_MAX_SIZE.get(),
+        )
+
+    # NB: Ideally we should return the vanilla Cache object only with maxsize.
+    # But we used TTLCache before introducing the timeout feature (that does not
+    # monitor timeout periodically nor log the expired traces). To keep the
+    # backward compatibility, we return TTLCache.
+    return TTLCache(
+        ttl=MLFLOW_TRACE_BUFFER_TTL_SECONDS.get(),
+        maxsize=MLFLOW_TRACE_BUFFER_MAX_SIZE.get(),
+    )
+
+
+class _TimedCache(Cache):
+    """
+    This code is ported from cachetools library to avoid depending on the private class.
+    https://github.com/tkem/cachetools/blob/d44c98407030d2e91cbe82c3997be042d9c2f0de/src/cachetools/__init__.py#L376
+    """
+
+    class _Timer:
+        def __init__(self, timer):
+            self.__timer = timer
+            self.__nesting = 0
+
+        def __call__(self):
+            if self.__nesting == 0:
+                return self.__timer()
+            else:
+                return self.__time
+
+        def __enter__(self):
+            if self.__nesting == 0:
+                self.__time = time = self.__timer()
+            else:
+                time = self.__time
+            self.__nesting += 1
+            return time
+
+        def __exit__(self, *exc):
+            self.__nesting -= 1
+
+        def __reduce__(self):
+            return _TimedCache._Timer, (self.__timer,)
+
+        def __getattr__(self, name):
+            return getattr(self.__timer, name)
+
+    def __init__(self, maxsize, timer=time.monotonic, getsizeof=None):
+        Cache.__init__(self, maxsize, getsizeof)
+        self.__timer = _TimedCache._Timer(timer)
+
+    def __repr__(self, cache_repr=Cache.__repr__):
+        with self.__timer as time:
+            self.expire(time)
+            return cache_repr(self)
+
+    def __len__(self, cache_len=Cache.__len__):
+        with self.__timer as time:
+            self.expire(time)
+            return cache_len(self)
+
+    @property
+    def currsize(self):
+        with self.__timer as time:
+            self.expire(time)
+            return super().currsize
+
+    @property
+    def timer(self):
+        """The timer function used by the cache."""
+        return self.__timer
+
+    def clear(self):
+        with self.__timer as time:
+            self.expire(time)
+            Cache.clear(self)
+
+    def get(self, *args, **kwargs):
+        with self.__timer:
+            return Cache.get(self, *args, **kwargs)
+
+    def pop(self, *args, **kwargs):
+        with self.__timer:
+            return Cache.pop(self, *args, **kwargs)
+
+    def setdefault(self, *args, **kwargs):
+        with self.__timer:
+            return Cache.setdefault(self, *args, **kwargs)
+
+
+class MlflowTraceTimeoutCache(_TimedCache):
+    """
+    A different implementation of cachetools.TTLCache that logs the expired traces to the backend.
+
+    NB: Do not use this class outside a singleton context. This class is not thread-safe.
+    """
+
+    def __init__(self, timeout: int, maxsize: int):
+        super().__init__(maxsize=maxsize)
+        self._timeout = timeout
+
+        # Set up the linked list ordered by expiration time
+        self._root = TTLCache._Link()
+        self._root.prev = self._root
+        self._root.next = self._root
+        self._links = OrderedDict()
+
+        self._start_expire_check_loop()
+
+    @property
+    def timeout(self) -> int:
+        # Timeout should not be changed after the cache is created
+        # because the linked list will not be updated accordingly.
+        return self._timeout
+
+    def __setitem__(self, key, value, cache_setitem=Cache.__setitem__):
+        """Set the item in the cache, and also in the linked list if it is a new key"""
+        with self.timer as time:
+            cache_setitem(self, key, value)
+
+        if key not in self._links:
+            # Add the new item to the tail of the linked list
+            # Inspired by https://github.com/tkem/cachetools/blob/d44c98407030d2e91cbe82c3997be042d9c2f0de/src/cachetools/__init__.py#L432
+            tail = self._root.prev
+            link = TTLCache._Link(key)
+            link.expires = time + self._timeout
+            link.next = self._root
+            link.prev = tail
+            tail.next = link
+            self._root.prev = link
+            self._links[key] = link
+
+    def __delitem__(self, key, cache_delitem=Cache.__delitem__):
+        """Delete the item from the cache and the linked list."""
+        cache_delitem(self, key)
+        link = self._links.pop(key)
+        link.unlink()
+
+    def _start_expire_check_loop(self):
+        # Close the daemon thread when the main thread exits
+        atexit.register(self.clear)
+
+        self._expire_checker_thread = threading.Thread(
+            target=self._expire_check_loop, daemon=True, name="TTLCacheExpireLoop"
+        )
+        self._expire_checker_stop_event = threading.Event()
+        self._expire_checker_thread.start()
+
+    def _expire_check_loop(self):
+        while not self._expire_checker_stop_event.is_set():
+            try:
+                self.expire()
+            except Exception as e:
+                _logger.debug(f"Failed to expire traces: {e}")
+                # If an error is raised from the expiration method, stop running the loop.
+                # Otherwise, the expire task might get heavier and heavier due to the
+                # increasing number of expired items.
+                break
+
+            time.sleep(MLFLOW_TRACE_TIMEOUT_CHECK_INTERVAL_SECONDS.get())
+
+    def expire(self, time=None):
+        """
+        Trigger the expiration of traces that have exceeded the timeout.
+
+        Args:
+            time: Unused. Only for compatibility with the parent class.
+        """
+        expired = self._get_expired_traces()
+
+        # End the expired traces and set the status to ERROR in background thread
+        for request_id in expired:
+            trace = self[request_id]
+            if root_span := trace.get_root_span():
+                try:
+                    root_span.set_status(SpanStatusCode.ERROR)
+                    msg = _TRACE_EXPIRATION_MSG.format(request_id=request_id, ttl=self._timeout)
+                    exception_event = SpanEvent.from_exception(MlflowTracingException(msg))
+                    root_span.add_event(exception_event)
+                    root_span.end()  # Calling end() triggers span export
+                    _logger.info(msg + " You can find the aborted trace in the MLflow UI.")
+                except Exception as e:
+                    _logger.debug(f"Failed to export an expired trace {request_id}: {e}")
+
+                # NB: root_span.end() should pop the trace from the cache. But we need to
+                # double-check it because it may not happens due to some errors.
+                if request_id in self:
+                    del self[request_id]
+
+    def _get_expired_traces(self) -> list[str]:
+        """
+        Find all expired traces and return their request IDs.
+
+        The linked list is ordered by expiration time, so we can traverse the list from the head
+        and return early whenever we find a trace that has not expired yet.
+        """
+        time = self.timer()
+        curr = self._root.next
+
+        if curr.expires and time < curr.expires:
+            return []
+
+        expired = []
+        while curr is not self._root and not (time < curr.expires):
+            expired.append(curr.key)
+            curr = curr.next
+        return expired
+
+    def clear(self):
+        super().clear()
+        self._expire_checker_stop_event.set()
+        self._expire_checker_thread.join()

mlflow/tracing/utils/token.py

@@ -0,0 +1,19 @@
+import contextvars
+from dataclasses import dataclass
+from typing import Optional
+
+from mlflow.entities import LiveSpan
+
+
+@dataclass
+class SpanWithToken:
+    """
+    A utility container to hold an MLflow span and its corresponding OpenTelemetry token.
+
+    The token is a special object that is generated when setting a span as active within
+    the Open Telemetry span context. This token is required when inactivate the span i.e.
+    detaching the span from the context.
+    """
+
+    span: LiveSpan
+    token: Optional[contextvars.Token] = None

mlflow/tracing/utils/truncation.py

@@ -0,0 +1,124 @@
+import json
+from functools import lru_cache
+from typing import Any, Optional
+
+from mlflow.entities.trace_data import TraceData
+from mlflow.entities.trace_info import TraceInfo
+from mlflow.tracing.constant import (
+    TRACE_REQUEST_RESPONSE_PREVIEW_MAX_LENGTH_DBX,
+    TRACE_REQUEST_RESPONSE_PREVIEW_MAX_LENGTH_OSS,
+)
+from mlflow.tracking._tracking_service.utils import get_tracking_uri
+from mlflow.utils.uri import is_databricks_uri
+
+
+def set_request_response_preview(trace_info: TraceInfo, trace_data: TraceData) -> None:
+    """
+    Set the request and response previews for the trace info.
+    """
+    # If request/response preview is already set by users via `mlflow.update_current_trace`,
+    # we don't override it with the truncated version.
+    if trace_info.request_preview is None:
+        trace_info.request_preview = _get_truncated_preview(trace_data.request, role="user")
+    if trace_info.response_preview is None:
+        trace_info.response_preview = _get_truncated_preview(trace_data.response, role="assistant")
+
+
+def _get_truncated_preview(request_or_response: Optional[str], role: str) -> str:
+    """
+    Truncate the request preview to fit the max length.
+    """
+    if request_or_response is None:
+        return ""
+
+    max_length = _get_max_length()
+
+    if len(request_or_response) <= max_length:
+        return request_or_response
+
+    content = None
+
+    # Parse JSON serialized request/response
+    try:
+        obj = json.loads(request_or_response)
+    except json.JSONDecodeError:
+        obj = None
+
+    if messages := _try_extract_messages(obj):
+        msg = _get_last_message(messages, role=role)
+        content = _get_text_content_from_message(msg)
+
+    content = content or request_or_response
+
+    if len(content) <= max_length:
+        return content
+
+    return content[: max_length - 3] + "..."
+
+
+@lru_cache(maxsize=1)
+def _get_max_length() -> int:
+    tracking_uri = get_tracking_uri()
+    return (
+        TRACE_REQUEST_RESPONSE_PREVIEW_MAX_LENGTH_DBX
+        if is_databricks_uri(tracking_uri)
+        else TRACE_REQUEST_RESPONSE_PREVIEW_MAX_LENGTH_OSS
+    )
+
+
+def _try_extract_messages(obj: dict[str, Any]) -> Optional[list[dict[str, Any]]]:
+    if not isinstance(obj, dict):
+        return None
+
+    # Check if the object contains messages with OpenAI ChatCompletion format
+    if messages := obj.get("messages"):
+        return [item for item in messages if _is_message(item)]
+
+    # Check if the object contains a message in OpenAI ChatCompletion response format (choices)
+    if (choices := obj.get("choices")) and len(choices) > 0:
+        return [choices[0].get("message")]
+
+    # Check if the object contains a message in OpenAI Responses API request format
+    if (input := obj.get("input")) and isinstance(input, list):
+        return [item for item in input if _is_message(item)]
+
+    # Check if the object contains a message in OpenAI Responses API response format
+    if (output := obj.get("output")) and isinstance(output, list):
+        return [item for item in output if _is_message(item)]
+
+    # Handle ResponsesAgent input, which contains OpenAI Responses request in 'request' key
+    if "request" in obj:
+        return _try_extract_messages(obj["request"])
+
+    return None
+
+
+def _is_message(item: Any) -> bool:
+    if not isinstance(item, dict):
+        return False
+    return "role" in item and "content" in item
+
+
+def _get_last_message(messages: list[dict[str, Any]], role: str) -> Optional[dict[str, Any]]:
+    """
+    Return last message with the given role.
+    If the messages don't include a message with the given role, return the last one.
+    """
+    for message in reversed(messages):
+        if message.get("role") == role:
+            return message
+    return messages[-1]
+
+
+def _get_text_content_from_message(message: dict[str, Any]) -> str:
+    content = message.get("content")
+    if isinstance(content, list):
+        # content is a list of content parts
+        for part in content:
+            if isinstance(part, str):
+                return part
+            elif isinstance(part, dict) and part.get("type") in ["text", "output_text"]:
+                return part.get("text")
+    elif isinstance(content, str):
+        return content
+    return ""

mlflow/tracing/utils/warning.py

@@ -0,0 +1,76 @@
+import functools
+import importlib
+import logging
+import warnings
+from typing import Optional
+
+_logger = logging.getLogger(__name__)
+
+
+class LogDemotionFilter(logging.Filter):
+    def __init__(self, module: str, message: str):
+        super().__init__()
+        self.module = module
+        self.message = message
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        if record.name == self.module and self.message in record.getMessage():
+            record.levelno = logging.DEBUG  # Change the log level to DEBUG
+            record.levelname = "DEBUG"
+
+            # Check the log level for the logger is debug or not
+            logger = logging.getLogger(self.module)
+            return logger.isEnabledFor(logging.DEBUG)
+        return True
+
+    def __eq__(self, other):
+        if isinstance(other, LogDemotionFilter):
+            return self.module == other.module and self.message == other.message
+        return False
+
+
+def suppress_warning(module: str, message: str):
+    """
+    Convert the "Failed to detach context" log raised by the OpenTelemetry logger to DEBUG
+    level so that it does not show up in the user's console.
+
+    Args:
+        module: The module name of the logger that raises the warning.
+        message: The (part of) message in the log that needs to be demoted to DEBUG level
+    """
+    try:
+        logger = getattr(importlib.import_module(module), "logger", None)
+        log_filter = LogDemotionFilter(module, message)
+        if logger and not any(f == log_filter for f in logger.filters):
+            logger.addFilter(log_filter)
+    except Exception as e:
+        _logger.debug(f"Failed to suppress the warning for {module}", exc_info=e)
+        raise
+
+
+def request_id_backward_compatible(func):
+    """
+    A decorator to support backward compatibility for the `request_id` parameter,
+    which is deprecated and replaced by the `trace_id` parameter in tracing APIs.
+
+    This decorator will adds `request_id` to the function signature and issue
+    a deprecation warning if `request_id` is used with non-null value.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, request_id: Optional[str] = None, **kwargs):
+        if request_id is not None:
+            warnings.warn(
+                f"The request_id parameter is deprecated from the {func.__name__} API "
+                "and will be removed in a future version. Please use the `trace_id` "
+                "parameter instead.",
+                category=DeprecationWarning,
+                stacklevel=2,
+            )
+
+            if kwargs.get("trace_id") is None:
+                kwargs["trace_id"] = request_id
+
+        return func(*args, **kwargs)
+
+    return wrapper

mlflow/tracking/__init__.py

@@ -0,0 +1,39 @@
+"""
+The ``mlflow.tracking`` module provides a Python CRUD interface to MLflow experiments
+and runs. This is a lower level API that directly translates to MLflow
+`REST API <../rest-api.html>`_ calls.
+For a higher level API for managing an "active run", use the :py:mod:`mlflow` module.
+"""
+
+# Minimum APIs required for core tracing functionality of mlflow-tracing package.
+from mlflow.tracking._tracking_service.utils import (
+    _get_artifact_repo,
+    _get_store,
+    get_tracking_uri,
+    is_tracking_uri_set,
+    set_tracking_uri,
+)
+from mlflow.version import IS_TRACING_SDK_ONLY
+
+__all__ = [
+    "get_tracking_uri",
+    "set_tracking_uri",
+    "is_tracking_uri_set",
+    "_get_artifact_repo",
+    "_get_store",
+]
+
+# Importing the following APIs only if mlflow or mlflow-skinny is installed.
+if not IS_TRACING_SDK_ONLY:
+    from mlflow.tracking._model_registry.utils import (
+        get_registry_uri,
+        set_registry_uri,
+    )
+    from mlflow.tracking._tracking_service.utils import _get_artifact_repo
+    from mlflow.tracking.client import MlflowClient
+
+    __all__ += [
+        "get_registry_uri",
+        "set_registry_uri",
+        "MlflowClient",
+    ]

mlflow/tracking/_model_registry/__init__.py

@@ -0,0 +1 @@
+DEFAULT_AWAIT_MAX_SLEEP_SECONDS = 5 * 60