为 MLflow Tracing 做出贡献

欢迎来到 MLflow Tracing 贡献指南！本分步资源将帮助您为 MLflow 实现额外的 GenAI 库集成，以进行追踪。

提示

如果在过程中有任何疑问，请尝试右下角的“Ask AI”功能。它可以提供参考文档以及 MLflow 常见问题的快速解答。

第 1 步. 设置您的环境

请按照 CONTRIBUTING.md 设置开发环境。设置完成后，请运行单元测试，命令为 pytest tests/tracing，并确保所有测试通过，以验证环境是否已准备好进行追踪开发。

第 2 步. 熟悉 MLflow Tracing

首先，请深入了解 MLflow Tracing 的功能及其工作原理。请查阅以下文档以快速掌握：

• 追踪概念 - 了解什么是追踪以及它对 MLflow 用户有何特定优势。
• MLflow Tracing API 指南 - 关于自动仪器化和手动创建追踪的 API 的实用指南。

📝 快速测验：在继续下一步之前，让我们通过几个问题来检验您的理解。如果您不确定答案，请重新查阅文档进行快速回顾。

问：Trace 和 Span 有什么区别？

答：Trace 是包含多个 Span 的主对象，每个 Span 捕获操作的不同部分。Trace 包含元数据 (TraceInfo) 和一个 Span 列表 (TraceData)。
参考： 追踪概念

问：创建函数调用 Span 的最简单方法是什么？

答：使用 @mlflow.trace 装饰器即可自动捕获输入、输出和执行时长。
参考： MLflow Tracing API 指南

问：如何将输入数据记录到 Span 中？

答：您可以使用 span.set_inputs() 方法为由 mlflow.start_span 上下文管理器或 Client API 返回的 Span 对象记录输入数据。
参考： 追踪概念

问：异常信息存储在 Trace 的哪个位置？

答：异常记录在 Span 的 events 属性中，包括异常类型、消息和堆栈跟踪等详细信息。
参考： MLflow Tracing API 指南

第 3 步. 了解集成库

从追踪的角度来看，GenAI 库可以分为两类：

🧠 LLM 提供商或包装器

像 OpenAI、Anthropic、Ollama 和 LiteLLM 这样的库专注于提供对 LLM 的访问。这些库通常有简单的客户端 SDK，因此我们通常使用临时修补 (ad-hoc patching) 来追踪这些 API。

对于此类库，请首先列出需要进行仪器化的核心 API。例如，在 Anthropic 自动追踪中，我们修补了 Messages 类中的 create() 方法。如果库有多个 API（例如，嵌入、转录），并且您不确定应该支持哪些，请咨询维护者。请参阅 Anthropic 自动追踪实现作为示例。

⚙️ 编排框架

像 LangChain、LlamaIndex 和 DSPy 这样的库提供更高级的流程，将 LLM、嵌入、检索器和工具集成到复杂的应用程序中。由于这些库需要来自多个组件的追踪数据，我们不希望依赖临时修补。因此，这些库的自动追踪通常会利用可用的回调（例如 LangChain Callbacks）来实现更可靠的集成。

对于此类库，请首先检查该库是否提供了任何可用的回调机制。如果没有，请考虑向该库提交功能请求，让项目维护者添加此功能，并为请求提供全面的理由。拥有回调机制还可以通过提供灵活性并允许与其他许多工具集成，从而使库的其他用户受益。如果由于某种原因该库无法提供回调，请咨询 MLflow 维护者。我们不太可能采用依赖临时修补的设计，但我们可以讨论其他可能的解决方案。

第 4 步. 编写设计文档

请使用 设计模板起草集成计划的设计文档。以下是一些重要的考虑因素：

• 集成方法：描述您将使用回调、API 钩子还是修补。如果存在多种方法，请将其列为选项并解释您的选择。
• 可维护性：LLM 框架的演变速度很快，因此请尽量避免依赖内部方法。优先使用公共 API，例如回调。
• 标准化：确保与其他 MLflow 集成的**一致性**，以提高可用性和下游任务的兼容性。例如，检索器 Span 应遵循 检索器模式以便 UI 兼容。

请简要概述库的核心功能和用例，为审阅者提供背景信息。草稿准备好后，与 MLflow 维护者分享您的设计，如果时间允许，请创建一个概念验证 (proof of concept) 以尽早突出潜在的挑战。

第 5 步. 开始实现

Python

设计获批后，开始实现。

1. 创建新模块：如果库尚未与 MLflow 集成，请在 mlflow/ 下创建一个新目录（例如，mlflow/llama_index）。添加一个 __init__.py 文件来初始化模块。
2. 开发追踪钩子：实现您选择的追踪方法（修补、回调或装饰器）。如果您选择修补方法，请使用 safe_patch 函数以确保稳定的修补（请参阅 示例）。
3. 定义 mlflow.xxx.autolog() 函数：此函数将是集成的**主入口点**，调用它时将启用追踪（例如，mlflow.llama_index.autolog()）。
4. 编写测试：涵盖边缘情况，例如异步调用、自定义数据类型和流式输出（如果库支持）。

注意

集成 MLflow Tracing 时有几个需要注意的陷阱：

• 错误处理：确保异常被捕获并记录到 Span 中，包含类型、消息和堆栈跟踪。
• 流式输出：对于流式（迭代器），请钩入迭代器以组装完整的输出并将其记录到 Span 中。直接记录迭代器对象不仅无用，而且可能导致意外行为，例如在序列化过程中耗尽迭代器。
• 序列化：MLflow 通过自定义 TraceJsonEncoder 实现将 Trace 序列化为 JSON，该实现支持常见对象和 Pydantic 模型。如果您的库使用自定义对象，请考虑扩展序列化器，因为不受支持的类型会被字符串化，并可能丢失有用的详细信息。
• 时间戳处理：在使用库提供的时间戳时，请验证单位和时区。MLflow 需要**自 UNIX 纪元以来的纳秒**作为时间戳；不正确的时间戳会扰乱 Span 持续时间。

TypeScript

MLflow TypeScript 追踪 SDK 位于 libs/typescript npm 工作区下。在实现新集成时：

1. 创建工作区包：使用 OpenAI 集成作为模板，使用 package.json、tsconfig.json 和可选的文档来搭建 libs/typescript/integrations/<provider> 结构。
2. 公开仪器化辅助函数：在 src/ 中实现追踪入口点，例如，一个 traced<Client>() 包装器或一个 register() 函数，用于将集成连接到 mlflow-tracing。
3. 文档和测试：在 tests/ 下添加 README 片段或内联文档，以及单元/集成测试，以便审阅者能够验证新的仪器化。
4. 更新工作区脚本：如果新包需要自定义构建或测试步骤，请扩展根目录的 package.json 中的 build:integrations 或 test:integrations 命令。

编码完成后，在本地构建并测试工作区，以确保 dist/ 中的工件是最新的。

cd libs/typescript
npm install
npm run build
npm run test

注意

• Trace 形状：使 Span 名称、属性和元数据与 Python 集成保持一致，以便下游 UI 组件能正确解释 Trace。
• 流式和异步工作负载：在记录输出之前，缓冲流式响应或异步迭代器，以便调用者可以继续使用它们。
• 可序列化负载：在将复杂的 SDK 对象传递给 MLflow 客户端之前，将其转换为 JSON 可序列化结构。

第 6 步. 测试集成

实现完成后，请在笔记本中运行端到端测试以验证功能。请确保：

◻︎ Trace 正确出现在 MLflow 实验中。

◻︎ Trace 在 MLflow UI 中正确渲染。

◻︎ MLflow Trace 创建过程中发生的错误不应中断库的原始执行。

◻︎ 边缘情况，例如异步和流式调用，按预期工作。

除了本地测试之外，还有几个与 MLflow Tracing 集成的 Databricks 服务。请咨询 MLflow 维护者以获取有关如何测试这些集成的指导。

当您确信实现已正确运行时，请提交一个 PR，并将测试结果粘贴在 PR 描述中。

第 7 步. 文档集成

文档是发布的前提。请按照以下步骤完成文档：

1. 在 主要 Tracing 文档 中添加集成库的图标和示例。
2. 如果库已存在于现有的 MLflow 模型风味中，请在风味文档中添加一个 Tracing 部分。
3. 添加一个笔记本教程来演示集成。

文档源文件位于 docs/ 文件夹中。请参阅 编写文档以获取有关如何构建和预览文档的更多详细信息。

第 8 步. 发布🚀

恭喜！您已经完成了为 MLflow 添加新的追踪集成的旅程。发布说明将包含您的姓名，我们将撰写 SNS 或/和博客文章来突出您的贡献。

非常感谢您为改进 MLflow Tracing 所做的贡献，我们期待与您再次合作！😊

联系方式

如果您有任何疑问或需要帮助，请随时与维护者联系（POC：@B-Step62，@BenWilson2）以获取进一步指导。