wehub-resource-sync
456 lines
17 KiB
Python
456 lines
17 KiB
Python
import logging
|
|
import threading
|
|
from collections import defaultdict
|
|
from functools import wraps
|
|
from typing import Any
|
|
|
|
import dspy
|
|
from dspy.utils.callback import BaseCallback
|
|
|
|
import mlflow
|
|
from mlflow.dspy.constant import FLAVOR_NAME
|
|
from mlflow.dspy.util import (
|
|
log_dspy_lm_state,
|
|
log_dspy_module_params,
|
|
sanitize_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.tracing.constant import SpanAttributeKey, TokenUsageKey
|
|
from mlflow.tracing.fluent import start_span_no_context
|
|
from mlflow.tracing.provider import detach_span_from_context, set_span_in_context
|
|
from mlflow.tracing.utils import maybe_set_prediction_context
|
|
from mlflow.tracing.utils.token import SpanWithToken
|
|
from mlflow.utils import _get_fully_qualified_class_name
|
|
from mlflow.utils.autologging_utils import (
|
|
get_autologging_config,
|
|
)
|
|
from mlflow.version import IS_TRACING_SDK_ONLY
|
|
|
|
_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
|
|
|
|
|
|
def _convert_signature(val):
|
|
# serialization of dspy.Signature is quite slow, so we should convert it to string
|
|
if isinstance(val, type) and issubclass(val, dspy.Signature):
|
|
return repr(val)
|
|
return val
|
|
|
|
|
|
class MlflowCallback(BaseCallback):
|
|
"""Callback for generating MLflow traces for DSPy components"""
|
|
|
|
def __init__(self, dependencies_schema: dict[str, Any] | None = None):
|
|
self._dependencies_schema = dependencies_schema
|
|
# call_id: (LiveSpan, OTel token)
|
|
self._call_id_to_span: dict[str, SpanWithToken] = {}
|
|
self._call_id_to_module: dict[str, Any] = {}
|
|
|
|
###### 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)
|
|
self._disabled_eval_call_ids = set()
|
|
self._eval_runs_started: set[str] = set()
|
|
|
|
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,
|
|
)
|
|
self._call_id_to_module[call_id] = instance
|
|
|
|
@skip_if_trace_disabled
|
|
def on_module_end(self, call_id: str, outputs: Any | None, exception: Exception | None = None):
|
|
instance = self._call_id_to_module.pop(call_id)
|
|
attributes = {}
|
|
|
|
if _get_fully_qualified_class_name(instance) == "dspy.retrieve.databricks_rm.DatabricksRM":
|
|
from mlflow.entities.document import Document
|
|
|
|
if isinstance(outputs, dspy.Prediction):
|
|
# Convert outputs to MLflow document format to make it compatible with
|
|
# agent evaluation.
|
|
num_docs = len(outputs.doc_ids)
|
|
doc_uris = outputs.doc_uris if outputs.doc_uris is not None else [None] * num_docs
|
|
outputs = [
|
|
Document(
|
|
page_content=doc_content,
|
|
metadata={
|
|
"doc_id": doc_id,
|
|
"doc_uri": doc_uri,
|
|
}
|
|
| extra_column_dict,
|
|
id=doc_id,
|
|
).to_dict()
|
|
for doc_content, doc_id, doc_uri, extra_column_dict in zip(
|
|
outputs.docs,
|
|
outputs.doc_ids,
|
|
doc_uris,
|
|
outputs.extra_columns,
|
|
)
|
|
]
|
|
else:
|
|
# 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):
|
|
usage_by_model = (
|
|
outputs.get_lm_usage() if hasattr(outputs, "get_lm_usage") else None
|
|
)
|
|
outputs = outputs.toDict()
|
|
if usage_by_model:
|
|
usage_data = {
|
|
TokenUsageKey.INPUT_TOKENS: 0,
|
|
TokenUsageKey.OUTPUT_TOKENS: 0,
|
|
TokenUsageKey.TOTAL_TOKENS: 0,
|
|
}
|
|
for usage in usage_by_model.values():
|
|
usage_data[TokenUsageKey.INPUT_TOKENS] += usage.get("prompt_tokens", 0)
|
|
usage_data[TokenUsageKey.OUTPUT_TOKENS] += usage.get("completion_tokens", 0)
|
|
usage_data[TokenUsageKey.TOTAL_TOKENS] += usage.get("total_tokens", 0)
|
|
attributes[SpanAttributeKey.CHAT_USAGE] = usage_data
|
|
# TODO: the span may not contain model name so we cannot calculate cost
|
|
self._end_span(call_id, outputs, exception, attributes)
|
|
|
|
@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
|
|
)
|
|
|
|
filtered_kwargs = sanitize_params(instance.kwargs)
|
|
attributes = {
|
|
**filtered_kwargs,
|
|
"model": instance.model,
|
|
"model_type": instance.model_type,
|
|
"cache": instance.cache,
|
|
SpanAttributeKey.MESSAGE_FORMAT: "dspy",
|
|
SpanAttributeKey.MODEL: instance.model,
|
|
}
|
|
match instance.model.split("/", 1):
|
|
case [provider, _]:
|
|
attributes[SpanAttributeKey.MODEL_PROVIDER] = provider
|
|
|
|
inputs = self._unpack_kwargs(inputs)
|
|
|
|
self._start_span(
|
|
call_id,
|
|
name=f"{instance.__class__.__name__}.__call__",
|
|
span_type=span_type,
|
|
inputs=inputs,
|
|
attributes=attributes,
|
|
)
|
|
|
|
@skip_if_trace_disabled
|
|
def on_lm_end(self, call_id: str, outputs: Any | None, exception: Exception | None = None):
|
|
self._end_span(call_id, outputs, exception)
|
|
|
|
@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: Any | None, exception: Exception | None = 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: Any | None, exception: Exception | None = 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: Any | None, exception: Exception | None = 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 callback_metadata.get("disable_logging"):
|
|
self._disabled_eval_call_ids.add(call_id)
|
|
return
|
|
started_run = False
|
|
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)
|
|
started_run = True
|
|
elif mlflow.active_run() is None:
|
|
mlflow.start_run(run_name=key, nested=True)
|
|
started_run = True
|
|
|
|
if started_run:
|
|
self._eval_runs_started.add(call_id)
|
|
if program := inputs.get("program"):
|
|
save_dspy_module_state(program, "model.json")
|
|
log_dspy_module_params(program)
|
|
|
|
# Log the current DSPy LM state
|
|
log_dspy_lm_state()
|
|
|
|
def on_evaluate_end(
|
|
self,
|
|
call_id: str,
|
|
outputs: Any,
|
|
exception: Exception | None = 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 call_id in self._disabled_eval_call_ids:
|
|
self._disabled_eval_call_ids.discard(call_id)
|
|
return
|
|
run_started = call_id in self._eval_runs_started
|
|
if exception:
|
|
if run_started:
|
|
mlflow.end_run(status=RunStatus.to_string(RunStatus.FAILED))
|
|
self._eval_runs_started.discard(call_id)
|
|
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)
|
|
|
|
if run_started:
|
|
mlflow.end_run()
|
|
self._eval_runs_started.discard(call_id)
|
|
# 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)
|
|
self._eval_runs_started = set()
|
|
|
|
def _start_span(
|
|
self,
|
|
call_id: str,
|
|
name: str,
|
|
span_type: SpanType,
|
|
inputs: dict[str, Any],
|
|
attributes: dict[str, Any],
|
|
):
|
|
if not IS_TRACING_SDK_ONLY:
|
|
from mlflow.pyfunc.context import get_prediction_context
|
|
|
|
prediction_context = get_prediction_context()
|
|
if prediction_context and self._dependencies_schema:
|
|
prediction_context.update(**self._dependencies_schema)
|
|
else:
|
|
prediction_context = None
|
|
|
|
with maybe_set_prediction_context(prediction_context):
|
|
span = start_span_no_context(
|
|
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: Any | None,
|
|
exception: Exception | None = None,
|
|
attributes: dict[str, Any] | None = 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))
|
|
|
|
if attributes:
|
|
st.span.set_attributes(attributes)
|
|
|
|
try:
|
|
st.span.end(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"}
|
|
merged = inputs_wo_kwargs | kwargs
|
|
return {k: _convert_signature(v) for k, v in merged.items()}
|
|
|
|
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
|