跳到主要内容

为 MLflow Tracing 贡献

欢迎来到 MLflow Tracing 贡献指南！本循序渐进的资源将协助您为 MLflow 实现额外的 GenAI 库集成以进行 tracing。

提示

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

第 1 步：设置您的环境​

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

第 2 步：熟悉 MLflow Tracing​

首先，深入了解 MLflow Tracing 的功能及其工作原理。请参阅以下文档以快速了解：

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

问：Trace 和 Span 有什么区别？

答：Trace 是包含多个 Span 的主对象，每个 Span 捕获操作的不同部分。Trace 包含元数据（TraceInfo）和 Span 列表（TraceData）。
参考： Tracing Concepts

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

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

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

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

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

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

第 3 步：理解集成库​

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

🧠 LLM Providers or Wrappers (LLM 提供商或包装器)

像 OpenAI、Anthropic、Ollama 和 LiteLLM 这样的库专注于提供对 LLM 的访问。这些库通常具有简单的客户端 SDK，因此我们经常使用 ad-hoc patching 来 trace 这些 API。

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

⚙️ Orchestration Frameworks (编排框架)

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

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

第 4 步：编写设计文档​

使用 design template 起草您的集成计划设计文档。以下是一些重要的注意事项：

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

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

第 5 步：开始实现​

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

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

集成 MLflow Tracing 时有几个需要注意的地方：

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

第 6 步：测试集成​

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

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

◻︎ Traces 在 MLflow UI 中正确渲染。

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

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

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

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

第 7 步：文档集成​

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

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

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

第 8 步：发布🚀​

恭喜！您已经完成了为 MLflow 添加新的 tracing 集成的旅程。发布说明将刊登您的名字，我们将撰写 SNS 或/和博客文章来表彰您的贡献。

非常感谢您为改进 MLflow Tracing 所做的贡献，我们期待与您再次合作！😊

联系方式​

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