跳到主要内容

为 MLflow 跟踪做贡献

欢迎阅读 MLflow 跟踪贡献指南！这份循序渐进的资源将帮助您实现更多 GenAI 库集成，以便在 MLflow 中进行跟踪。

提示

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

步骤 1. 设置您的环境​

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

步骤 2. 熟悉 MLflow 跟踪​

首先，对 MLflow 跟踪的功能和工作原理有一个扎实的理解。查阅这些文档以快速掌握：

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

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

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

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

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

问：何时会使用 MLflow 客户端创建跟踪？

答：当您需要精细控制如何开始和结束跟踪时，客户端 API 非常有用。例如，您可以在启动 span 时指定父 span ID。
参考： MLflow 跟踪 API 指南

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

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

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

答：异常记录在 span 的 events 属性中，包括异常类型、消息和堆栈跟踪等详细信息。
参考： MLflow 跟踪 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 回调）以实现更可靠的集成。

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

步骤 4. 编写设计文档​

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

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

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

步骤 5. 开始实施​

设计获批后，即可开始实施：

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

在与 MLflow 跟踪集成时，有几点需要注意：

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

步骤 6. 测试集成​

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

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

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

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

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

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

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

步骤 7. 文档化集成​

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

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

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

步骤 8. 发布🚀​

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

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

联系方式​

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