MLOps with Databricks

Chapter 8. MLflow for GenAI

MLflow for GenAI can be thought of as a distinct layer on top of the MLflow concepts we’ve seen so far. While familiar building blocks still apply, such as experiment tracking for logging runs and registering models, the application of these features shifts significantly in the context of generative AI.

In GenAI workflows, we still log and register models, but we rely on the PyFunc flavor with a different abstraction. Instead of extending PythonModel, GenAI applications typically use classes such as ResponsesAgent, which are designed to represent agents rather than traditional predictive models.

Evaluation is also fundamentally different. Rather than focusing on standard ML metrics, GenAI evaluation centers on behavioral qualities and agent decisions. To support this, MLflow introduces traces to provide fine-grained observability into everything that happens inside an agent, including tool calls, intermediate steps, and execution flow in real time.

In the following sections, we’ll explore MLflow’s GenAI flavors and the fundamentals of MLflow tracing. From there, we’ll define a custom agent with tracing enabled and progressively build toward evaluation, logging, and registering the agent.

MLflow Flavors

In Chapter 3, we explored logging and registering ML models using MLflow, and we discussed different built-in flavors like scikit-learn (sklearn), LightGBM, and others. In the same way, MLflow provides specialized flavors for agentic frameworks, including ...