为 MLflow Tracing 贡献

欢迎阅读 MLflow Tracing 贡献指南！本分步指南将协助您为 MLflow 实现额外的 GenAI 库集成以进行跟踪。

提示

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

步骤 1. 设置您的环境

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

步骤 2. 熟悉 MLflow Tracing

首先，扎实理解 MLflow Tracing 的作用和工作原理。查阅这些文档以快速掌握情况

• 跟踪概念 - 了解什么是跟踪以及它对 MLflow 用户的具体好处。
• MLflow 跟踪 API 指南 - 关于自动检测和用于手动创建跟踪的 API 的实用指南。

📝 快速测验：在进入下一步之前，让我们用几个问题挑战一下您的理解。如果您不确定答案，请快速回顾文档。

问：Trace 和 Span 有什么区别？

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

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

答：使用 @mlflow.trace 装饰器自动捕获输入、输出和执行持续时间。
参考： MLflow 跟踪 API 指南

问：如何向 Span 记录输入数据？

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

问：异常信息存储在 Trace 的哪里？

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

步骤 3. 理解集成库

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

🧠 LLM 提供商或包装器

像 OpenAI、Anthropic、Ollama 和 LiteLLM 这样的库专注于提供对 LLM 的访问。这些库通常具有简单的客户端 SDK，因此我们通常只需使用即席补丁来跟踪这些 API。

对于此类库，首先列出需要检测的核心 API。例如，在 Anthropic 自动跟踪中，我们对 Messages 类的 create() 方法进行补丁。如果该库有多个 API（例如嵌入、转录）而您不确定应该支持哪些，请咨询维护人员。请参考 Anthropic 自动跟踪实现 作为一个示例。

⚙️ 编排框架

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

对于此类库，首先检查该库是否提供任何可供您使用的回调机制。如果没有，请考虑向该库提交功能请求，以便项目维护者添加此功能，并为请求提供充分的理由。拥有回调机制也有益于该库的其他用户，因为它提供了灵活性并允许与其他许多工具集成。如果由于某种特定原因该库无法提供回调，请咨询 MLflow 维护人员。我们很可能不会采用依赖即席补丁的设计，但如果有任何替代方法，我们可以讨论。

步骤 4. 撰写设计文档

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

• 集成方法：描述您将使用回调、API 钩子还是打补丁。如果有多种方法，请将它们列为选项并解释您的选择。
• 可维护性：LLM 框架发展迅速，因此应尽可能避免依赖内部方法。首选公共 API，例如回调。
• 标准化：确保与其他 MLflow 集成的连贯性，以实现可用性和下游任务。例如，检索 Span 应遵循 检索器模式 (Retriever Schema) 以实现 UI 兼容性。

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

步骤 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 实现将跟踪序列化为 JSON，该实现支持常见对象和 Pydantic 模型。如果您的库使用自定义对象，请考虑扩展序列化器，因为不受支持的类型会被字符串化，可能会丢失有用的详细信息。
• 时间戳处理：使用库提供的时间戳时，请验证单位和时区。MLflow 要求时间戳以自 UNIX 纪元以来的纳秒为单位；错误的时间戳将扰乱 Span 持续时间。

TypeScript

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

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

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

bash

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

注意

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

步骤 6. 测试集成

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

◻︎ 跟踪正确显示在 MLflow 实验中。

◻︎ 跟踪在 MLflow UI 中正确渲染。

◻︎ 来自 MLflow 跟踪创建的错误不应中断库的原始执行。

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

除了本地测试外，还有一些与 MLflow Tracing 集成的 Databricks 服务。请咨询 MLflow 维护人员，了解如何测试这些集成。

当您确信实现正确工作时，请打开一个 PR，并将测试结果粘贴到 PR 描述中。

步骤 7. 文档化集成

文档是发布的先决条件。请遵循以下步骤完成文档编制

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

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

步骤 8. 发布🚀

恭喜！现在您已经完成了向 MLflow 添加新的跟踪集成的全部历程。发布说明中将刊登您的名字，我们将撰写 SNS 或/和博客文章来突出您的贡献。

非常感谢您帮助改进 MLflow Tracing，我们期待与您再次合作！😊

联系方式

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