为 MLflow 追踪做贡献

欢迎来到 MLflow 跟踪贡献指南！这个逐步资源将帮助您在 MLflow 中实现额外的 GenAI 库集成以进行跟踪。

小技巧

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

步骤 1. 设置你的环境

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

第二步：熟悉 MLflow 追踪

首先，深入理解 MLflow Tracing 的功能及其工作原理。查看这些文档以快速上手：

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

Q. What is the difference between a Trace and a Span?

A. Trace is the main object holding multiple Spans, with each Span capturing different parts of an operation. A Trace has metadata (TraceInfo) and a list of Spans (TraceData). Reference: Tracing Schema Guide

Q. What is the easiest way to create a span for a function call?

A. Use the @mlflow.trace decorator to capture inputs, outputs, and execution duration automatically. Reference: MLflow Tracing API Guide

Q. When would you use the MLflow Client for creating traces?

A. The Client API is useful when you want fine-grained control over how you start and end a trace. For example, you can specify a parent span ID when starting a span. Reference: MLflow Tracing API Guide

Q. How do you log input data to a span?

A. You can log input data with the span.set_inputs() method for a span object returned by the ``mlflow.start_span`` context manager or Client APIs. Reference: Tracing Schema Guide

Q. Where is exception information stored in a Trace?

A. Exceptions are recorded in the events attribute of the span, including details such as exception type, message, and stack trace. References: MLflow Tracing API Guide

步骤 3. 理解集成库

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

🧠 LLM 提供者或包装器

像 OpenAI、Anthropic、Ollama 和 LiteLLM 这样的库专注于提供对大型语言模型（LLMs）的访问。这些库通常有简单的客户端SDK，因此，我们通常简单地使用临时补丁来追踪这些API。

对于这种类型的库，首先列出要检测的核心API。例如，在OpenAI自动追踪中，我们修补了 ChatCompletion、Completion 和 Embedding 类的 create() 方法。参考 OpenAI自动追踪实现 作为示例。

⚙️ 编排框架

诸如 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 Tracing 时，有一些需要注意的陷阱：

  • 错误处理：确保捕获异常并记录到具有类型、消息和堆栈跟踪的跨度中。

  • 流输出：对于流（迭代器），挂钩到迭代器以组装并将完整输出记录到范围中。直接记录迭代器对象不仅无益，还可能导致意外行为，例如在序列化期间耗尽迭代器。

  • 序列化：MLflow 通过自定义的 TraceJsonEncoder 实现将跟踪信息序列化为 JSON，该实现支持常见对象和 Pydantic 模型。如果你的库使用自定义对象，建议扩展序列化器，因为不支持的类型会被字符串化，可能会丢失有用细节。

  • 时间戳处理：在使用库提供的时间戳时，验证单位和时区。MLflow 需要时间戳为 自 UNIX 纪元以来的纳秒；错误的时间戳将扰乱跨度持续时间。

步骤 6. 测试集成

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

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

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

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

◻︎ 像异步和流式调用这样的边缘情况按预期工作。

除了本地测试外，还有一些 Databricks 服务与 MLflow Tracing 集成。请联系 MLflow 维护者以获取如何测试这些集成的指导。

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

步骤 7. 文档集成

文档是发布的先决条件。按照以下步骤完成文档：

  1. 在 主要追踪文档 中添加集成库图标和示例。

  2. 如果库已经存在于现有的 MLflow 模型风格中，请在风格文档中添加一个追踪部分（示例页面）。

  3. 添加一个笔记本教程以演示集成 (示例笔记本)

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

步骤 8. 发布🚀

恭喜！现在您已经完成了向 MLflow 添加新的跟踪集成的旅程。发布说明将包含您的名字，并且我们将撰写一篇社交媒体帖子或/和一篇博客文章来突出您的贡献。

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

联系

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