为何mlflow.langchain.autolog()无法捕获自定义LangGraph StateGraph追踪信息？

问题背景 #

我正尝试在项目中集成MLflow autologging与基于LangGraph的智能体。根据MLflow文档，启用以下代码应该能自动捕获LangGraph工作流的追踪信息：

import mlflow
mlflow.set_tracking_uri("http://127.0.0.1:5000")
mlflow.set_experiment("thingworx_agentic_v2")
mlflow.langchain.autolog()

问题现象 #

实际应用中，我用StateGraph构建自定义图：

from langgraph.graph import StateGraph, END
from app.core.graph_builder import build_agent_graph

graph = build_agent_graph(checkpointer=memory_saver)
graph.stream({...})

即使在构建图前已启用mlflow.langchain.autolog()，MLflow中虽显示运行记录，但包含0个参数、0个指标，且无追踪日志。

疑问 #

• mlflow.langchain.autolog()是否仅支持create_react_agent及其他预构建LangGraph智能体？
• 若是，是否有办法将自定义StateGraph工作流接入MLflow追踪，使追踪信息显示在UI中？
• 或者我需要手动为每个节点添加Span埋点？

版本信息 #

• mlflow: 2.12.0
• langchain-core: 0.3.72
• langgraph: 0.6.3
• Python: 3.11

我的预期是mlflow.langchain.autolog()能与LangGraph的StateGraph兼容。

解答 #

1. autolog的支持范围
   mlflow.langchain.autolog()在v2.12.0版本中对LangGraph的支持更偏向预构建智能体（比如create_react_agent生成的图），对自定义StateGraph的自动追踪支持有限。原因是预构建智能体有标准化的结构和钩子，MLflow自动日志能识别这些标准化节点；而自定义StateGraph的节点逻辑、流转规则都是自定义的，自动日志无法自动识别所有执行细节。

2. 自定义StateGraph接入MLflow追踪的方案
   有两种可行方式：

• 利用LangGraph回调系统结合MLflow：LangGraph支持在图初始化时传入回调函数，可自定义回调捕获节点的输入输出、状态变化，再手动调用MLflow API记录信息。示例代码：

from langchain_core.callbacks.base import BaseCallbackHandler
import mlflow

class MLflowLangGraphCallback(BaseCallbackHandler):
    def on_chain_start(self, serialized, inputs, **kwargs):
        mlflow.log_params(inputs)
        self.current_span = mlflow.start_run(run_name=serialized.get("name", "node_run"), nested=True)
    
    def on_chain_end(self, outputs, **kwargs):
        # 可根据实际需求记录执行时间等指标
        mlflow.log_params(outputs)
        mlflow.end_run()

# 构建图时传入回调
graph = build_agent_graph(checkpointer=memory_saver, callbacks=[MLflowLangGraphCallback()])

• 手动嵌入MLflow日志代码：在每个自定义节点的逻辑中直接添加MLflow埋点，记录参数、指标和日志：

def custom_node(state):
    with mlflow.start_run(nested=True):
        mlflow.log_param("node_input", state["input"])
        # 节点业务逻辑
        result = ...
        mlflow.log_param("node_output", result)
        mlflow.log_metric("node_processed_length", len(result))
        return {"output": result}

3. 手动添加Span埋点的必要性
   如果回调方案无法满足需求，手动添加Span（利用MLflow的嵌套Run）是可行方案。每个节点的执行作为嵌套MLflow Run，能在UI中展示层级追踪信息，清晰呈现整个工作流的执行路径。

内容的提问来源于stack exchange，提问作者Rupesh Akotkar