跳到主要内容

为 MLflow 跟踪做贡献

欢迎阅读 MLflow 跟踪贡献指南！这份分步资源将帮助您为 MLflow 实现更多 GenAI 库的跟踪集成。

提示

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

第 1 步. 设置您的环境​

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

第 2 步. 熟悉 MLflow 跟踪​

首先，深入了解 MLflow 跟踪的功能和工作原理。查看这些文档以快速掌握：

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

问：跟踪（Trace）和跨度（Span）之间有什么区别？

答：跟踪是包含多个跨度的主要对象，每个跨度捕获操作的不同部分。一个跟踪有元数据（TraceInfo）和一个跨度列表（TraceData）。
参考：跟踪概念

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

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

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

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

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

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

第 3 步. 了解集成库​

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

🧠 LLM 提供商或包装器

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

对于这类库，首先列出需要插桩的核心 API。例如，在 Anthropic 自动跟踪中，我们修补了 Messages 类的 create() 方法。如果库有多个 API（例如，嵌入、转录），而您不确定支持哪些，请咨询维护者。请参考 Anthropic 自动跟踪实现 作为示例。

⚙️ 编排框架

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

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

第 4 步. 编写设计文档​

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

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

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

第 5 步. 开始实施​

设计获得批准后，开始实施：

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

在与 MLflow 跟踪集成时，需要注意以下几个陷阱：

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

第 6 步. 测试集成​

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

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

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

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

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

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

当您确信实施工作正常后，请提交一个 PR，并将测试结果粘贴在 PR 描述中。

第 7 步. 记录集成​

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

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

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

第 8 步. 发布🚀​

恭喜！您现在已经完成了向 MLflow 添加新跟踪集成的旅程。发布说明将包含您的名字，我们将撰写 SNS 或/和博客文章来重点介绍您的贡献。

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

联系方式​

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