71 lines
2.4 KiB
Python
71 lines
2.4 KiB
Python
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
|