跳到主要内容

为 MLflow Tracing 贡献

欢迎来到 MLflow Tracing 贡献指南！这份循序渐进的资源将帮助您实现在 MLflow 中对更多 GenAI 库集成进行 Tracing。

提示

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

步骤 1. 设置您的环境​

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

步骤 2. 熟悉 MLflow Tracing​

首先，深入了解 MLflow Tracing 的作用和工作原理。查看以下文档以快速掌握

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

问：Trace 和 Span 有什么区别？

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

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

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

问：何时使用 MLflow Client 创建 trace？

答：当您想要对 trace 的开始和结束方式进行细粒度控制时，Client API 非常有用。例如，您可以在开始 span 时指定父 span ID。
参考： MLflow Tracing API 指南

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

答：您可以使用 span.set_inputs() 方法将输入数据记录到由 mlflow.start_span 上下文管理器或 Client API 返回的 span 对象中。
参考： Tracing Schema 指南

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

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

步骤 3. 了解集成库​

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

🧠 LLM 提供商或包装器

OpenAI、Anthropic、Ollama 和 LiteLLM 等库主要提供对 LLM 的访问。这些库通常具有简单的客户端 SDK，因此我们通常直接使用 ad-hoc patching 来 trace 这些 API。

对于这类库，首先列出需要进行插桩的核心 API。例如，在 Anthropic 自动 tracing 中，我们 patch 了 Messages 类的 create() 方法。如果库有多个 API（例如，embeddings、transcription），并且您不确定支持哪些 API，请咨询维护者。请参考 Anthropic 自动 tracing 实现 作为示例。

⚙️ 编排框架

LangChain、LlamaIndex 和 DSPy 等库提供更高级的工作流程，将 LLM、embeddings、retrievers 和工具集成到复杂应用中。由于这些库需要来自多个组件的 trace 数据，我们不希望依赖 ad-hoc patching。因此，对这些库进行自动 tracing 通常会利用可用的回调机制（例如，LangChain Callbacks）来实现更可靠的集成。

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

步骤 4. 撰写设计文档​

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

  • 集成方法：描述您将使用回调、API 钩子还是 patching。如果存在多种方法，请将它们列为选项并解释您的选择。
  • 可维护性：LLM 框架发展迅速，因此尽可能避免依赖内部方法。首选公共 API，如回调。
  • 标准化：确保与其他 MLflow 集成保持一致性，以提高可用性和下游任务效率。例如，retrieval span 应遵循Retriever Schema，以确保 UI 兼容性。

包含库核心功能和用例的简要概述，以便为评审者提供上下文。草稿准备好后，与 MLflow 维护者分享您的设计，如果时间允许，创建一个概念验证（proof of concept）来尽早突出潜在挑战。

步骤 5. 开始实现​

设计批准后，开始实现

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

在与 MLflow Tracing 集成时，有一些需要注意的陷阱

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

步骤 6. 测试集成​

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

◻︎ Trace 在 MLflow Experiment 中正确显示。

◻︎ Trace 在 MLflow UI 中正确呈现。

◻︎ MLflow trace 创建产生的错误不应中断库的原始执行。

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

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

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

步骤 7. 撰写集成文档​

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

  1. 在主 Tracing 文档中添加集成库的图标和示例。
  2. 如果该库已存在于现有的 MLflow 模型 flavor 中，请在该 flavor 文档中添加 Tracing 部分（示例页面）。
  3. 添加一个 notebook 教程来演示集成（示例 notebook）

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

步骤 8. 发布🚀​

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

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

联系方式​

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