star

venv/lib/python3.9/site-packages/mlflow/dspy/callback.py

@@ -0,0 +1,403 @@
+import logging
+import threading
+from collections import defaultdict
+from functools import wraps
+from typing import Any, Optional, Union
+
+import dspy
+from dspy.utils.callback import BaseCallback
+
+import mlflow
+from mlflow.dspy.save import FLAVOR_NAME
+from mlflow.dspy.util import log_dspy_module_params, save_dspy_module_state
+from mlflow.entities import SpanStatusCode, SpanType
+from mlflow.entities.run_status import RunStatus
+from mlflow.entities.span_event import SpanEvent
+from mlflow.exceptions import MlflowException
+from mlflow.pyfunc.context import get_prediction_context, maybe_set_prediction_context
+from mlflow.tracing.provider import detach_span_from_context, set_span_in_context
+from mlflow.tracing.utils import (
+    end_client_span_or_trace,
+    set_span_chat_messages,
+    start_client_span_or_trace,
+)
+from mlflow.tracing.utils.token import SpanWithToken
+from mlflow.utils.autologging_utils import (
+    get_autologging_config,
+)
+
+_logger = logging.getLogger(__name__)
+_lock = threading.Lock()
+
+
+def skip_if_trace_disabled(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        if get_autologging_config(FLAVOR_NAME, "log_traces"):
+            func(*args, **kwargs)
+
+    return wrapper
+
+
+class MlflowCallback(BaseCallback):
+    """Callback for generating MLflow traces for DSPy components"""
+
+    def __init__(self, dependencies_schema: Optional[dict[str, Any]] = None):
+        self._client = mlflow.MlflowClient()
+        self._dependencies_schema = dependencies_schema
+        # call_id: (LiveSpan, OTel token)
+        self._call_id_to_span: dict[str, SpanWithToken] = {}
+
+        ###### state management for optimization process ######
+        # The current callback logic assumes there is no optimization running in parallel.
+        # The state management may not work when multiple optimizations are running in parallel.
+        # optimizer_stack_level is used to determine if the callback is called within compile
+        # we cannot use boolean flag because the callback can be nested
+        self.optimizer_stack_level = 0
+        # call_id: (key, step)
+        self._call_id_to_metric_key: dict[str, tuple[str, int]] = {}
+        self._evaluation_counter = defaultdict(int)
+
+    def set_dependencies_schema(self, dependencies_schema: dict[str, Any]):
+        if self._dependencies_schema:
+            raise MlflowException(
+                "Dependencies schema should be set only once to the callback.",
+                error_code=MlflowException.INVALID_PARAMETER_VALUE,
+            )
+        self._dependencies_schema = dependencies_schema
+
+    @skip_if_trace_disabled
+    def on_module_start(self, call_id: str, instance: Any, inputs: dict[str, Any]):
+        span_type = self._get_span_type_for_module(instance)
+        attributes = self._get_span_attribute_for_module(instance)
+
+        # The __call__ method of dspy.Module has a signature of (self, *args, **kwargs),
+        # while all built-in modules only accepts keyword arguments. To avoid recording
+        # empty "args" key in the inputs, we remove it if it's empty.
+        if "args" in inputs and not inputs["args"]:
+            inputs.pop("args")
+
+        self._start_span(
+            call_id,
+            name=f"{instance.__class__.__name__}.forward",
+            span_type=span_type,
+            inputs=self._unpack_kwargs(inputs),
+            attributes=attributes,
+        )
+
+    @skip_if_trace_disabled
+    def on_module_end(
+        self, call_id: str, outputs: Optional[Any], exception: Optional[Exception] = None
+    ):
+        # NB: DSPy's Prediction object is a customized dictionary-like object, but its repr
+        # is not easy to read on UI. Therefore, we unpack it to a dictionary.
+        # https://github.com/stanfordnlp/dspy/blob/6fe693528323c9c10c82d90cb26711a985e18b29/dspy/primitives/prediction.py#L21-L28
+        if isinstance(outputs, dspy.Prediction):
+            outputs = outputs.toDict()
+
+        self._end_span(call_id, outputs, exception)
+
+    @skip_if_trace_disabled
+    def on_lm_start(self, call_id: str, instance: Any, inputs: dict[str, Any]):
+        span_type = (
+            SpanType.CHAT_MODEL if getattr(instance, "model_type", None) == "chat" else SpanType.LLM
+        )
+
+        attributes = {
+            **instance.kwargs,
+            "model": instance.model,
+            "model_type": instance.model_type,
+            "cache": instance.cache,
+        }
+
+        inputs = self._unpack_kwargs(inputs)
+
+        span = self._start_span(
+            call_id,
+            name=f"{instance.__class__.__name__}.__call__",
+            span_type=span_type,
+            inputs=inputs,
+            attributes=attributes,
+        )
+
+        if messages := self._extract_messages_from_lm_inputs(inputs):
+            try:
+                set_span_chat_messages(span, messages)
+            except Exception as e:
+                _logger.debug(f"Failed to set input messages for {span}. Error: {e}")
+
+    @skip_if_trace_disabled
+    def on_lm_end(
+        self, call_id: str, outputs: Optional[Any], exception: Optional[Exception] = None
+    ):
+        st = self._call_id_to_span.get(call_id)
+        try:
+            output_msg = self._extract_messages_from_lm_outputs(outputs)
+            set_span_chat_messages(st.span, output_msg, append=True)
+        except Exception as e:
+            _logger.debug(f"Failed to set output messages for {call_id}. Error: {e}")
+
+        self._end_span(call_id, outputs, exception)
+
+    def _extract_messages_from_lm_inputs(self, inputs: dict[str, Any]) -> list[dict[str, str]]:
+        # LM input is either a list of messages or a prompt string
+        # https://github.com/stanfordnlp/dspy/blob/ac5bf56bb1ed7261d9637168563328c1dfeb27af/dspy/clients/lm.py#L92
+        # TODO: Extract tool definition once https://github.com/stanfordnlp/dspy/pull/2023 is merged
+        return inputs.get("messages") or [{"role": "user", "content": inputs.get("prompt")}]
+
+    def _extract_messages_from_lm_outputs(
+        self, outputs: list[Union[str, dict[str, Any]]]
+    ) -> list[dict[str, str]]:
+        # LM output is either a string or a dictionary of text and logprobs
+        # https://github.com/stanfordnlp/dspy/blob/ac5bf56bb1ed7261d9637168563328c1dfeb27af/dspy/clients/lm.py#L105-L114
+        # TODO: Extract tool calls once https://github.com/stanfordnlp/dspy/pull/2023 is merged
+        return [
+            {"role": "assistant", "content": o.get("text") if isinstance(o, dict) else o}
+            for o in outputs
+        ]
+
+    @skip_if_trace_disabled
+    def on_adapter_format_start(self, call_id: str, instance: Any, inputs: dict[str, Any]):
+        self._start_span(
+            call_id,
+            name=f"{instance.__class__.__name__}.format",
+            span_type=SpanType.PARSER,
+            inputs=self._unpack_kwargs(inputs),
+            attributes={},
+        )
+
+    @skip_if_trace_disabled
+    def on_adapter_format_end(
+        self, call_id: str, outputs: Optional[Any], exception: Optional[Exception] = None
+    ):
+        self._end_span(call_id, outputs, exception)
+
+    @skip_if_trace_disabled
+    def on_adapter_parse_start(self, call_id: str, instance: Any, inputs: dict[str, Any]):
+        self._start_span(
+            call_id,
+            name=f"{instance.__class__.__name__}.parse",
+            span_type=SpanType.PARSER,
+            inputs=self._unpack_kwargs(inputs),
+            attributes={},
+        )
+
+    @skip_if_trace_disabled
+    def on_adapter_parse_end(
+        self, call_id: str, outputs: Optional[Any], exception: Optional[Exception] = None
+    ):
+        self._end_span(call_id, outputs, exception)
+
+    @skip_if_trace_disabled
+    def on_tool_start(self, call_id: str, instance: Any, inputs: dict[str, Any]):
+        # DSPy uses the special "finish" tool to signal the end of the agent.
+        if instance.name == "finish":
+            return
+
+        inputs = self._unpack_kwargs(inputs)
+        # Tools are always called with keyword arguments only.
+        inputs.pop("args", None)
+
+        self._start_span(
+            call_id,
+            name=f"Tool.{instance.name}",
+            span_type=SpanType.TOOL,
+            inputs=inputs,
+            attributes={
+                "name": instance.name,
+                "description": instance.desc,
+                "args": instance.args,
+            },
+        )
+
+    @skip_if_trace_disabled
+    def on_tool_end(
+        self, call_id: str, outputs: Optional[Any], exception: Optional[Exception] = None
+    ):
+        if call_id in self._call_id_to_span:
+            self._end_span(call_id, outputs, exception)
+
+    def on_evaluate_start(self, call_id: str, instance: Any, inputs: dict[str, Any]):
+        """
+        Callback handler at the beginning of evaluation call. Available with DSPy>=2.6.9.
+        This callback starts a nested run for each evaluation call inside optimization.
+        If called outside optimization and no active run exists, it creates a new run.
+        """
+        if not get_autologging_config(FLAVOR_NAME, "log_evals"):
+            return
+
+        key = "eval"
+        if callback_metadata := inputs.get("callback_metadata"):
+            if "metric_key" in callback_metadata:
+                key = callback_metadata["metric_key"]
+        if self.optimizer_stack_level > 0:
+            with _lock:
+                # we may want to include optimizer_stack_level in the key
+                # to handle nested optimization
+                step = self._evaluation_counter[key]
+                self._evaluation_counter[key] += 1
+            self._call_id_to_metric_key[call_id] = (key, step)
+            mlflow.start_run(run_name=f"{key}_{step}", nested=True)
+        else:
+            mlflow.start_run(run_name=key, nested=True)
+        if program := inputs.get("program"):
+            save_dspy_module_state(program, "model.json")
+            log_dspy_module_params(program)
+
+    def on_evaluate_end(
+        self,
+        call_id: str,
+        outputs: Any,
+        exception: Optional[Exception] = None,
+    ):
+        """
+        Callback handler at the end of evaluation call. Available with DSPy>=2.6.9.
+        This callback logs the evaluation score to the individual run
+        and add eval metric to the parent run if called inside optimization.
+        """
+        if not get_autologging_config(FLAVOR_NAME, "log_evals"):
+            return
+        if exception:
+            mlflow.end_run(status=RunStatus.to_string(RunStatus.FAILED))
+            return
+        score = None
+        if isinstance(outputs, float):
+            score = outputs
+        elif isinstance(outputs, tuple):
+            score = outputs[0]
+        elif isinstance(outputs, dspy.Prediction):
+            score = float(outputs)
+            try:
+                mlflow.log_table(self._generate_result_table(outputs.results), "result_table.json")
+            except Exception:
+                _logger.debug("Failed to log result table.", exc_info=True)
+        if score is not None:
+            mlflow.log_metric("eval", score)
+
+        mlflow.end_run()
+        # Log the evaluation score to the parent run if called inside optimization
+        if self.optimizer_stack_level > 0 and mlflow.active_run() is not None:
+            if call_id not in self._call_id_to_metric_key:
+                return
+            key, step = self._call_id_to_metric_key.pop(call_id)
+            if score is not None:
+                mlflow.log_metric(
+                    key,
+                    score,
+                    step=step,
+                )
+
+    def reset(self):
+        self._call_id_to_metric_key: dict[str, tuple[str, int]] = {}
+        self._evaluation_counter = defaultdict(int)
+
+    def _start_span(
+        self,
+        call_id: str,
+        name: str,
+        span_type: SpanType,
+        inputs: dict[str, Any],
+        attributes: dict[str, Any],
+    ):
+        prediction_context = get_prediction_context()
+        if prediction_context and self._dependencies_schema:
+            prediction_context.update(**self._dependencies_schema)
+
+        with maybe_set_prediction_context(prediction_context):
+            span = start_client_span_or_trace(
+                self._client,
+                name=name,
+                span_type=span_type,
+                parent_span=mlflow.get_current_active_span(),
+                inputs=inputs,
+                attributes=attributes,
+            )
+
+        token = set_span_in_context(span)
+        self._call_id_to_span[call_id] = SpanWithToken(span, token)
+
+        return span
+
+    def _end_span(
+        self,
+        call_id: str,
+        outputs: Optional[Any],
+        exception: Optional[Exception] = None,
+    ):
+        st = self._call_id_to_span.pop(call_id, None)
+
+        if not st.span:
+            _logger.warning(f"Failed to end a span. Span not found for call_id: {call_id}")
+            return
+
+        status = SpanStatusCode.OK if exception is None else SpanStatusCode.ERROR
+
+        if exception:
+            st.span.add_event(SpanEvent.from_exception(exception))
+
+        try:
+            end_client_span_or_trace(
+                client=self._client,
+                span=st.span,
+                outputs=outputs,
+                status=status,
+            )
+        finally:
+            detach_span_from_context(st.token)
+
+    def _get_span_type_for_module(self, instance):
+        if isinstance(instance, dspy.Retrieve):
+            return SpanType.RETRIEVER
+        elif isinstance(instance, dspy.ReAct):
+            return SpanType.AGENT
+        elif isinstance(instance, dspy.Predict):
+            return SpanType.LLM
+        elif isinstance(instance, dspy.Adapter):
+            return SpanType.PARSER
+        else:
+            return SpanType.CHAIN
+
+    def _get_span_attribute_for_module(self, instance):
+        if isinstance(instance, dspy.Predict):
+            return {"signature": instance.signature.signature}
+        elif isinstance(instance, dspy.ChainOfThought):
+            if hasattr(instance, "signature"):
+                signature = instance.signature.signature
+            else:
+                signature = instance.predict.signature.signature
+
+            attributes = {"signature": signature}
+            if hasattr(instance, "extended_signature"):
+                attributes["extended_signature"] = instance.extended_signature.signature
+            return attributes
+        return {}
+
+    def _unpack_kwargs(self, inputs: dict[str, Any]) -> dict[str, Any]:
+        """Unpacks the kwargs from the inputs dictionary"""
+        # NB: Not using pop() to avoid modifying the original inputs dictionary
+        kwargs = inputs.get("kwargs", {})
+        inputs_wo_kwargs = {k: v for k, v in inputs.items() if k != "kwargs"}
+        return {**inputs_wo_kwargs, **kwargs}
+
+    def _generate_result_table(
+        self, outputs: list[tuple[dspy.Example, dspy.Prediction, Any]]
+    ) -> dict[str, list[Any]]:
+        result = {"score": []}
+        for i, (example, prediction, score) in enumerate(outputs):
+            for k, v in example.items():
+                if f"example_{k}" not in result:
+                    result[f"example_{k}"] = [None] * i
+                result[f"example_{k}"].append(v)
+
+            for k, v in prediction.items():
+                if f"pred_{k}" not in result:
+                    result[f"pred_{k}"] = [None] * i
+                result[f"pred_{k}"].append(v)
+
+            result["score"].append(score)
+
+            for k, v in result.items():
+                if len(v) != i + 1:
+                    result[k].append(None)
+
+        return result