MLflow Fluent API 完全指南:Tracing 可观测性与 Logged Model 版本管理(mlflow 模块 Python API 详解)
【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow
mlflow模块是 MLflow 面向 Python 开发者暴露的高层 "fluent"(流式)API 入口,它把实验跟踪、模型管理、Tracing 可观测性等能力收敛为一行行直观的函数调用。本文以官方 API 参考文档 docs/api_reference/source/python_api/mlflow.rst 为骨架,系统讲解其中两大核心 API 族群——MLflow Tracing APIs(trace / span / assessment)与MLflow Logged Model APIs(模型初始化、定稿、检索与标签管理),并结合仓库源码深入剖析每个函数的签名、参数语义与底层调用链。读完本文,你将能熟练地为自己的函数一键加装可观测性、把远端服务的 trace 合并进本地调用链、为 LLM 应用记录人工反馈,并用 Logged Model 体系对 Agent / LLM 应用做版本化追踪。
一、mlflow 模块定位:一切高层的入口
1.1 从一条最简单的示例看 fluent API
mlflow模块的设计哲学是:用最少的样板代码完成实验记录。官方模块文档(见 mlflow/init.py 顶部 docstring)给出了最经典的用法:
import mlflow mlflow.start_run() mlflow.log_param("my", "param") mlflow.log_metric("score", 100) mlflow.end_run()也可以使用上下文管理器语法,with块结束时自动结束 run:
with mlflow.start_run() as run: mlflow.log_param("my", "param") mlflow.log_metric("score", 100)需要特别注意两点:
- 线程安全:模块文档明确声明,fluent tracking API 目前不是线程安全的,任何并发调用者必须自行实现互斥;
- 高低层分工:如果需要对 tracking 行为做更精细、更低层的控制,应当改用
:py:mod:mlflow.client`` 模块(即MlflowClient),fluent API 则面向大多数常规场景。
1.2 模块内部结构:lazy load 与版本自检
从 mlflow/init.py 的源码可以看到这个入口模块的工程细节:
- 版本自检:导入时先调用
mlflow.mismatch._check_version_mismatch(),并在异常时静默忽略,防止客户端与服务器版本不一致导致崩溃; - LazyLoader 机制:
anthropic、langchain、openai、sklearn、transformers等数十个 flavor 子模块全部通过LazyLoader延迟加载(见 mlflow/init.py),避免import mlflow时拉入大量重型依赖; __all__白名单:模块末尾显式定义了__all__列表(mlflow/init.py),涵盖 tracing 核心 API、assessment API 以及最少量的 tracking API,方便 IDE 补全与静态检查。
1.3 API 参考文档的组织方式
关联文档mlflow.rst使用 Sphinx 的automodule+autofunction指令组织内容:
automodule:: mlflow会为模块内除:exclude-members:之外的全部公开成员生成文档;- 被排除的成员(如
trace、start_span、search_traces、initialize_logged_model等)随后在正文中以autofunction的形式单独成节,最终形成两个专章:MLflow Tracing APIs与MLflow Logged Model APIs——这正是本文的核心主题。
顺带说明:load_prompt、register_prompt、search_prompts、set_prompt_alias、delete_prompt_alias等 Prompt Registry API 同样在排除列表中,但未在本文件内展开(它们已迁移至mlflow.genai命名空间,见 mlflow/init.py 中的注释),本文不做展开。
二、MLflow Tracing APIs:为 LLM 与 Agent 应用加装可观测性
MLflow Tracing 是一套 OpenTelemetry 兼容的 LLM/Agent 可观测性方案,能够捕获一次请求中每个中间步骤的输入、输出与元数据,帮助快速定位 bug 与异常行为(参见总览文档 docs/docs/genai/tracing/index.mdx)。mlflow模块为此提供了一组高层 tracing fluent API,全部实现在 mlflow/tracing/fluent.py。
2.1mlflow.trace:装饰器式埋点
mlflow.trace是最简单、最常用的埋点方式。它会为被装饰的函数自动创建一个 span,并自动捕获函数的输入输出;方法上的self参数不会被捕获。函数内任何异常都会将 span 状态置为ERROR,异常信息与堆栈会记录到 span 的attributes字段(见 mlflow/tracing/fluent.py)。
import mlflow @mlflow.trace def my_function(x, y): return x + y上面的写法等价于用mlflow.start_span上下文管理器手写:
with mlflow.start_span("my_function") as span: x, y = 1, 2 span.set_inputs({"x": x, "y": y}) result = my_function(x, y) span.set_outputs({"output": result})@mlflow.trace装饰器支持的函数类型(源码 docstring 明确列出):
| 函数类型 | 支持情况 |
|---|---|
| 同步函数 Sync | ✅ |
| 异步函数 Async | ✅(>= 2.16.0) |
| 生成器 Generator | ✅(>= 2.20.2) |
| 异步生成器 Async Generator | ✅(>= 2.20.2) |
| 类方法 ClassMethod | ✅(>= 3.0.0) |
| 静态方法 StaticMethod | ✅(>= 3.0.0) |
mlflow.trace的完整签名(mlflow/tracing/fluent.py):
mlflow.trace( func=None, # 被装饰函数;作为装饰器使用时必须省略 name=None, # span 名称,缺省时使用函数名 span_type=SpanType.UNKNOWN, # span 类型,可为字符串或 SpanType 枚举 attributes=None, # 附加到 span 的属性字典 output_reducer=None, # 输出归约函数,用于压缩输出 trace_destination=None, # trace 的写入目标 sampling_ratio_override=None, # 采样率覆盖 log_level=None, # SpanLogLevel 或 "INFO"/"DEBUG" 等名称 links=None, # Link 对象列表 description=None, # 人类可读的 span 描述 )函数包装(Function Wrapping):除了作为装饰器,mlflow.trace还可以直接包装外部库中的函数而不修改其定义:
import math import mlflow mlflow.trace(math.factorial)(5) # 直接包装调用 traced_pow = mlflow.trace(math.pow) # 或先包装再调用 traced_pow(2, 10)这在为第三方库函数加装追踪时非常有用。详见手动埋点指南 docs/docs/genai/tracing/app-instrumentation/manual-tracing.mdx。
2.2mlflow.start_span:上下文管理器式埋点
当需要在一个代码块内手动管理输入输出、自定义属性时,使用mlflow.start_span上下文管理器(mlflow/tracing/fluent.py):
with mlflow.start_span("my_span") as span: span.set_inputs({"x": x, "y": y}) z = x + y span.set_outputs(z) span.set_attribute("key", "value")关键语义:
- 在上下文管理器内部新建的 span 会被自动挂为子 span,父子关系自动维护;
- 在顶层作用域(不在任何 span 上下文内)使用时,该 span 是根 span,整个 trace 会在根 span 结束时被写入;
- 上下文管理器退出时 span 自动结束;内部抛出的异常会将 span 置为
ERROR并记录异常信息; - 上下文管理器默认不跨线程传播 span 上下文,多线程场景需要手动传播(参考手动埋点文档中的 Multi-Threading 章节)。
start_span参数(mlflow/tracing/fluent.py):
| 参数 | 默认值 | 说明 |
|---|---|---|
name | "span" | span 名称 |
span_type | SpanType.UNKNOWN | 字符串或SpanType枚举 |
attributes | None | 附加属性字典 |
trace_destination | None | 写入目标(如 MLflow Experiment);仅对根 span 生效,非根 span 设置会被忽略并告警 |
log_level | None | 严重级别(SpanLogLevel或名称),缺省在结束时按 span 类型解析 |
run_id | None | 关联的 MLflow run ID;仅在创建根 span 时生效,优先级高于mlflow.start_run()设置的活跃 run |
links | None | 关联的Link对象列表 |
description | None | span 的人类可读描述 |
2.3mlflow.start_span_no_context:低层手动控制
当需要完全掌控 span 生命周期与父子关系时,使用start_span_no_context(mlflow/tracing/fluent.py)。该 API 创建的 span不会挂到全局 tracing 上下文,因此get_current_active_span等 API 无法感知它;span 必须手动调用end()结束:
root_span = mlflow.start_span_no_context("my_trace") # 创建子 span,手动指定父 span child_span = mlflow.start_span_no_context( "child_span", parent_span=root_span, inputs={"query": "..."}, ) # ... 业务逻辑 ... child_span.end() root_span.end()相比start_span,它额外暴露了parent_span、inputs、tags、metadata、experiment_id、start_time_ns等参数。源码 docstring 建议:只要start_span能满足需求就优先使用它,因为样板代码更少、更不易出错。
2.4 查询 Trace:mlflow.get_trace与mlflow.search_traces
mlflow.get_trace(trace_id, silent=False, flush=False)(mlflow/tracing/fluent.py):按 trace ID 获取单个 trace,返回mlflow.entities.Trace或None。内部先查内存缓冲,再查 tracking store;silent=True时找不到不告警,flush=True时若未找到会先冲刷待写入的异步 trace 再重试(对脚本/测试中异步日志尚未落盘的情况很关键):
with mlflow.start_span(name="span") as span: span.set_attribute("key", "value") trace = mlflow.get_trace(span.trace_id) print(trace)mlflow.search_traces(...)(mlflow/tracing/fluent.py):按条件批量检索 trace,返回pandas.DataFrame或Trace列表。其完整参数如下:
| 参数 | 说明 |
|---|---|
experiment_ids | 限定搜索范围的实验 ID 列表 |
filter_string | 搜索过滤字符串 |
max_results | 最大返回条数,None表示返回全部匹配项 |
order_by | 排序子句列表 |
extract_fields | 已弃用(3.6.0),仅return_type="pandas"时生效,格式为"span_name.[inputs\|outputs].field_name" |
run_id | 关联的 run ID;活跃 run 下创建的 trace 可用 run_id 过滤 |
return_type | "pandas"(默认,已安装 pandas 时)或"list" |
model_id | 只检索与该模型 ID 关联的 trace |
sql_warehouse_id | 已弃用,改用环境变量MLFLOW_TRACING_SQL_WAREHOUSE_ID,仅 Databricks 场景 |
include_spans | True时返回 trace 含 span;False只返回元数据(trace ID、起止时间等),默认True |
locations | 指定搜索位置列表:实验 ID,或 Databricks UC 表格式<catalog>.<schema>[.<table_prefix>];缺省搜索当前活跃实验 |
flush | True时先冲刷待写入的异步 trace 再搜索 |
使用示例:
# 以 DataFrame 形式返回当前实验的所有 trace mlflow.search_traces() # 按 model_id 过滤并以 Trace 对象列表返回 traces = mlflow.search_traces(model_id=active_model_id, return_type="list", flush=True)源码 docstring 给出了一条实用建议:如果预期结果集很大,应直接使用MlflowClient.search_traces分页拉取,因为本函数会一次性把全部结果载入内存。
2.5 活跃上下文查询:get_current_active_span与get_last_active_trace_id
mlflow.get_current_active_span()(mlflow/tracing/fluent.py):返回全局上下文中的当前活跃 span(LiveSpan),没有则返回None。注意:只对@mlflow.trace或with mlflow.start_span创建的 span 生效;start_span_no_context创建的 span 不在全局上下文中,不会被返回。
@mlflow.trace def f(): span = mlflow.get_current_active_span() span.set_attribute("key", "value") return 0mlflow.get_last_active_trace_id(thread_local=False)(mlflow/tracing/fluent.py):返回最近一次活跃的 trace ID;thread_local=True时按线程隔离取值。这与get_active_trace_id()(线程安全地返回当前活跃 trace ID)一起,常用于在 span 之外关联 trace。
2.6mlflow.add_trace:合并远端服务的 trace
mlflow.add_trace(trace, target=None)(mlflow/tracing/fluent.py)可以把一个已经完成的 trace 对象(Trace实例或 dict)合并进当前活跃的本地 trace。典型场景是调用了一个由 MLflow Tracing 埋点的远程服务,远程 trace 随响应返回,本地通过该函数将远端 trace 并入当前调用链,从而看到完整的端到端链路:
@mlflow.trace(name="predict") def predict(input): resp = requests.get("https://your-service-endpoint", ...) trace_json = resp.json().get("trace") # 将远端 trace 合并到当前活跃 trace 的 "predict" span 之下 mlflow.add_trace(trace_json)也可以指定目标 span:
with mlflow.start_span(name="predict") as span: resp = requests.get("https://your-service-endpoint", ...) mlflow.add_trace(resp.json().get("trace"), target=span)需要注意:传入的 trace 必须是已完成状态(不能再被修改),且 span 必须按"父 span 在子 span 之前"的顺序排列,否则会抛异常。
2.7 Assessment APIs:为 trace 记录期望与反馈
Assessment(评估标注)是 Tracing 体系的重要一环:把期望答案(expectation,ground truth)与质量反馈(feedback)附加到 trace 上,供后续评估与监控使用。全部实现位于 mlflow/tracing/assessment.py。
mlflow.log_assessment(trace_id, assessment)(mlflow/tracing/assessment.py):通用入口,接受Expectation或Feedback对象:
- Expectation:某个操作的期望值标签,例如聊天机器人的期望答案;
- Feedback:对操作质量的反向标注,可来自人工评审、启发式打分或 LLM-as-a-Judge。
from mlflow.entities import Feedback feedback = Feedback( name="faithfulness", value=0.9, rationale="The model is faithful to the input.", metadata={"model": "gpt-4o-mini"}, ) mlflow.log_assessment(trace_id="1234", assessment=feedback)记录带来源信息的人工期望值(未指定source时默认来源为类型HUMAN、ID 为"default"):
from mlflow.entities import AssessmentSource, AssessmentSourceType, Expectation source = AssessmentSource(source_type=AssessmentSourceType.HUMAN, source_id="john@example.com") expectation = Expectation(name="expected_answer", value=42, source=source) mlflow.log_assessment(trace_id="1234", assessment=expectation)期望值可以是任意 JSON 可序列化的值,例如完整的 LLM 消息(含期望的 tool calls)。
mlflow.log_expectation(*, trace_id, name, value, source=None, metadata=None, span_id=None)(mlflow/tracing/assessment.py):只接受关键字参数的便捷入口,用于记录期望值:
mlflow.log_expectation( trace_id="tr-1234567890abcdef", name="expected_answer", value="The capital of France is Paris.", source=AssessmentSource(source_type=AssessmentSourceType.HUMAN, source_id="annotator@company.com"), metadata={"question_type": "factual", "difficulty": "easy"}, )span_id参数可将期望关联到 trace 内的特定 span。
mlflow.log_feedback(*, trace_id, name=DEFAULT_FEEDBACK_NAME, value=None, source=None, error=None, rationale=None, metadata=None, span_id=None)(mlflow/tracing/assessment.py):记录质量反馈,name缺省为"feedback",支持error参数传入异常以标注失败反馈。
mlflow.update_assessment(trace_id, assessment_id, assessment)(mlflow/tracing/assessment.py):更新已存在的期望或反馈:
response = mlflow.log_assessment( trace_id="1234", assessment=Expectation(name="expected_answer", value=42), ) mlflow.update_assessment( trace_id="1234", assessment_id=response.assessment_id, assessment=Expectation(name="expected_answer", value=43), )mlflow.delete_assessment(trace_id, assessment_id)(mlflow/tracing/assessment.py):删除 trace 上的某个评估标注。
以上 tracing API 的行为在 tests/tracing/test_fluent.py(覆盖trace装饰器、start_span上下文管理器、get_trace、search_traces等)与 tests/tracing/test_assessment.py 中有大量测试用例佐证,例如test_trace、test_start_span_context_manager、test_search_traces等。
三、MLflow Logged Model APIs:版本化追踪 Agent / LLM 应用
Logged Model 是 MLflow 3 引入的模型实体概念:一个包含模型元数据、状态与关联 trace/评估的容器,用于对"未按 MLflow Model 格式打包的"模型、应用或 Agent 做版本化追踪与性能数据管理。本组 API 主要实现在 mlflow/tracking/fluent.py。
3.1 生命周期管理:initialize / finalize / create_external_model
mlflow.initialize_logged_model(name=None, source_run_id=None, tags=None, params=None, model_type=None, experiment_id=None)(mlflow/tracking/fluent.py):初始化一个状态为PENDING、无任何 artifact 的 LoggedModel。随后必须向模型添加 artifact 并调用finalize_logged_model将其置为READY,例如通过mlflow.pyfunc.log_model()等 flavor 方法。name缺省时随机生成;source_run_id缺省时使用活跃 run 的 ID。
mlflow.finalize_logged_model(model_id, status)(mlflow/tracking/fluent.py):更新模型终态,status取值"READY"/"FAILED"或LoggedModelStatus枚举:
from mlflow.entities import LoggedModelStatus model = mlflow.initialize_logged_model(name="model") logged_model = mlflow.finalize_logged_model(model_id=model.model_id, status=LoggedModelStatus.READY) assert logged_model.status == LoggedModelStatus.READYmlflow.create_external_model(name=None, source_run_id=None, tags=None, params=None, model_type=None, experiment_id=None)(mlflow/tracking/fluent.py):直接创建一个状态为READY的 LoggedModel,其 artifact存储在 MLflow 之外。非常适合跟踪未采用 MLflow Model 格式打包的模型、应用或生成式 AI Agent 的参数与性能数据。源码实现会为其打上内部标记MLFLOW_MODEL_IS_EXTERNAL="true"(mlflow/tracking/fluent.py)。model_type是用户自定义字符串,例如设置model_type="agent"便于日后按类型搜索比较同类模型。
3.2 检索:get_logged_model / last_logged_model / search_logged_models
mlflow.get_logged_model(model_id)(mlflow/tracking/fluent.py):按 ID 获取 LoggedModel,返回mlflow.entities.LoggedModel。mlflow.last_logged_model()(mlflow/tracking/fluent.py):返回当前会话中最近一次记录的 LoggedModel(实现上借助进程内_last_logged_model_id变量,见 mlflow/tracking/fluent.py);从未记录过则返回None。mlflow.search_logged_models(experiment_ids=None, filter_string=None, datasets=None, max_results=None, order_by=None, output_format="pandas")(mlflow/tracking/fluent.py):按条件搜索,output_format可选"pandas"(返回 DataFrame)或"list"(返回LoggedModel列表)。
典型组合用法:
model_info = mlflow.pyfunc.log_model(name="model", python_model=DummyModel()) logged_model = mlflow.get_logged_model(model_id=model_info.model_id) last_model = mlflow.last_logged_model() assert last_model.model_id == model_info.model_id3.3 标签与参数:set_logged_model_tags / delete_logged_model_tag / log_model_params
mlflow.set_logged_model_tags(model_id, tags)(mlflow/tracking/fluent.py):为 LoggedModel 批量设置标签:
model_info = mlflow.pyfunc.log_model(name="model", python_model=DummyModel()) mlflow.set_logged_model_tags(model_info.model_id, {"key": "value"}) model = mlflow.get_logged_model(model_info.model_id) assert model.tags["key"] == "value"mlflow.delete_logged_model_tag(model_id, key)(mlflow/tracking/fluent.py):删除指定标签键。mlflow.log_model_params(params, model_id=None)(mlflow/tracking/fluent.py):为 LoggedModel 记录参数;model_id缺省时使用当前活跃模型 ID(get_active_model_id())。
3.4 活跃模型:set_active_model / clear_active_model
mlflow.set_active_model(*, name=None, model_id=None)(mlflow/tracking/fluent.py):设置活跃模型,其生命周期内产生的 trace 会自动与模型关联。返回值是ActiveModel对象,可当上下文管理器使用;否则需手动调用mlflow.clear_active_model()更新。
name:按名称设置;若该名称的 LoggedModel 不存在,会在当前实验下自动创建;存在多个同名模型时取最新的一个并告警;model_id:按 ID 设置;ID 不存在则抛异常。
# 按名称设置(不存在则自动创建) mlflow.set_active_model(name="my_model") # 按 model_id 设置 model = mlflow.create_external_model(name="test_model") mlflow.set_active_model(model_id=model.model_id) # 上下文管理器用法 with mlflow.set_active_model(name="new_model"): print(mlflow.get_active_model_id()) # 活跃模型生命周期内产生的 trace 自动关联 mlflow.set_active_model(name="my_model") @mlflow.trace def predict(model_input): return model_input predict("abc") traces = mlflow.search_traces(model_id=mlflow.get_active_model_id(), return_type="list", flush=True) assert len(traces) == 1mlflow.clear_active_model()(mlflow/tracking/fluent.py):清除当前活跃模型,使后续 trace 不再关联任何模型。
3.5 实战串联:一个完整的版本化应用示例
官方快速入门 docs/docs/genai/version-tracking/quickstart.mdx 演示了将 Prompt 版本化、LangChain 应用追踪与 Logged Model 串联的完整流程(需 MLflow 3.0+):
import mlflow # 1. 注册带版本历史的 prompt 模板 system_prompt = mlflow.genai.register_prompt( name="chatbot_prompt", template="You are a chatbot that can answer questions about IT. Answer this question: {{question}}", commit_message="Initial version of chatbot", ) # 2. 构建 LangChain 链 from langchain.schema.output_parser import StrOutputParser from langchain_core.prompts import ChatPromptTemplate from langchain_openai import ChatOpenAI prompt = ChatPromptTemplate.from_template(system_prompt.to_single_brace_format()) chain = prompt | ChatOpenAI(temperature=0.7) | StrOutputParser() # 3. 设置活跃模型并开启 autolog,trace 自动关联到模型 mlflow.set_active_model(name="langchain_model") mlflow.langchain.autolog() outputs = [chain.invoke(q) for q in questions] active_model_id = mlflow.get_active_model_id() mlflow.search_traces(model_id=active_model_id) # 验证 trace 已关联以上 Logged Model 相关 API 均有对应的测试用例,见 tests/tracking/fluent/test_fluent.py 中的test_initialize_logged_model_active_run、test_log_model_params、test_finalized_logged_model、test_search_logged_models、test_set_active_model等。
四、与 MlflowClient 的分工:何时用 fluent,何时用 client
关联文档在automodule指令的:exclude-members:列表中排除了MlflowClient,因为它属于低层编程接口,单独文档化。二者的分工原则:
- fluent API(本文所述):面向"当前进程内、当前活跃上下文"的高频操作,代码简短、自动维护活跃 run / span / model 上下文,适合交互式开发与常规脚本;
MlflowClient:面向精确控制,例如为 trace 检索分页、按 ID 精确操作实体、跨上下文管理(见 mlflow/client.py 与模块入口 mlflow/init.py 中的导入)。
五、小结
围绕 docs/api_reference/source/python_api/mlflow.rst 这份 API 参考文档,本文完整覆盖了mlflow模块的两大 API 族群:
- Tracing APIs:
trace(装饰器/包装)、start_span、start_span_no_context(手动控制)、get_trace/search_traces(查询)、get_current_active_span/get_last_active_trace_id(上下文感知)、add_trace(远端链路合并),以及log_assessment/log_expectation/log_feedback/update_assessment/delete_assessment一组评估标注 API; - Logged Model APIs:
initialize_logged_model/finalize_logged_model/create_external_model(生命周期)、get_logged_model/last_logged_model/search_logged_models(检索)、set_logged_model_tags/delete_logged_model_tag/log_model_params(元数据)、set_active_model/clear_active_model(trace 关联)。
所有 API 均可在源码 mlflow/tracing/fluent.py、mlflow/tracing/assessment.py、mlflow/tracking/fluent.py 中查看完整 docstring 与实现,并可通过对应测试目录tests/tracing与tests/tracking/fluent验证行为。上手时建议从@mlflow.trace+mlflow.set_active_model的组合起步,即可在数行代码内获得完整的 LLM 应用可观测性与版本化追踪能力。
【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考