star

venv/lib/python3.9/site-packages/mlflow/pyfunc/context.py

@@ -0,0 +1,70 @@
+import contextlib
+from contextvars import ContextVar
+from dataclasses import dataclass
+from typing import Optional
+
+# A thread local variable to store the context of the current prediction request.
+# This is particularly used to associate logs/traces with a specific prediction request in the
+# caller side. The context variable is intended to be set by the called before invoking the
+# predict method, using the set_prediction_context context manager.
+_PREDICTION_REQUEST_CTX = ContextVar("mlflow_prediction_request_context", default=None)
+
+
+@dataclass
+class Context:
+    # A unique identifier for the current prediction request.
+    request_id: Optional[str] = None
+    # Whether the current prediction request is as a part of MLflow model evaluation.
+    is_evaluate: bool = False
+    # The schema of the dependencies to be added into the tag of trace info.
+    dependencies_schemas: Optional[dict] = None
+
+    def update(self, **kwargs):
+        for key, value in kwargs.items():
+            if hasattr(self, key):
+                setattr(self, key, value)
+            else:
+                raise AttributeError(f"Context has no attribute named '{key}'")
+
+
+@contextlib.contextmanager
+def set_prediction_context(context: Optional[Context]):
+    """
+    Set the context for the current prediction request. The context will be set as a thread-local
+    variable and will be accessible globally within the same thread.
+
+    Args:
+        context: The context for the current prediction request.
+    """
+    if context and not isinstance(context, Context):
+        raise TypeError(f"Expected context to be an instance of Context, but got: {context}")
+
+    token = _PREDICTION_REQUEST_CTX.set(context)
+    try:
+        yield
+    finally:
+        _PREDICTION_REQUEST_CTX.reset(token)
+
+
+def get_prediction_context() -> Optional[Context]:
+    """
+    Get the context for the current prediction request. The context is thread-local and is set
+    using the set_prediction_context context manager.
+
+    Returns:
+        The context for the current prediction request, or None if no context is set.
+    """
+    return _PREDICTION_REQUEST_CTX.get()
+
+
+@contextlib.contextmanager
+def maybe_set_prediction_context(context: Optional[Context]):
+    """
+    Set the prediction context if the given context
+    is not None. Otherwise no-op.
+    """
+    if context:
+        with set_prediction_context(context):
+            yield
+    else:
+        yield