跳到主要内容

从代码创建模型

注意

从代码创建模型功能在 MLflow 2.12.2 及更高版本中提供。如果您使用的版本低于支持此功能的版本，则需要使用自定义 Python 模型文档中概述的传统序列化方法。

注意

从代码创建模型仅适用于 LangChain、LlamaIndex，以及使用 pyfunc 编写的自定义 Python Agent 或生成式 AI 应用。对于其他用例（例如使用 xgboost 的经典机器学习），如果您直接使用 ML 库，建议使用相应的模型 Flavor 内的保存和日志记录功能。

从代码创建模型功能全面改进了定义、存储和加载自定义模型以及不依赖于序列化模型权重（例如 LangChain 和 LlamaIndex）的特定 Flavor 实现的过程。如果您正在编写自定义 Python 模型或生成式 AI Agent/应用，则应使用从代码创建模型。

从代码创建模型与这些模型的传统序列化之间的关键区别在于模型在序列化时的表示方式。

在传统方法中，序列化是使用 cloudpickle（自定义 pyfunc 和 LangChain）或自定义序列化器（对于 LlamaIndex）对模型对象执行的，自定义序列化器对底层包的所有功能的覆盖不完整。对于自定义 pyfunc，使用 cloudpickle 序列化对象实例会创建一个二进制文件，用于在加载时重构对象。

在从代码创建模型中，对于支持的模型类型，会保存一个包含自定义 pyfunc 或 Flavor 接口定义的简单脚本（例如，对于 LangChain，我们可以在脚本中直接定义一个 LCEL Chain 并将其标记为模型）。

使用从代码创建模型处理自定义 pyfunc 和支持库实现的最大优势在于减少了在实现过程中可能出现的重复试错调试。下文所示的工作流程说明了这两种方法在处理自定义模型解决方案时的比较。

Models from code comparison with legacy serialization

与传统序列化的区别​

在自定义模型的传统模式中，调用 log_model 时会提交您的子类 mlflow.pyfunc.PythonModel 的一个实例。通过对象引用调用时，MLflow 将利用 cloudpickle 尝试序列化您的对象。

在 LangChain 的原生 Flavor 序列化中，使用 cloudpickle 存储对象引用。然而，由于外部状态引用或在 API 中使用 lambda 函数，LangChain 中所有可用的对象类型中只有一部分可以进行序列化。另一方面，LlamaIndex 在 Flavor 的原生实现中使用了自定义序列化器，由于需要过度复杂的实现来支持库中的边缘案例特性，该序列化器无法覆盖库的所有可能用法。

在从代码创建模型中，您只需传递包含模型定义的脚本的路径引用，而不是传递自定义模型实例的对象引用。使用此模式时，MLflow 将在执行环境中简单地执行此脚本（以及运行主脚本之前的任何 code_paths 依赖项），并实例化您在调用 mlflow.models.set_model() 中定义的任何对象，将该对象指定为推理目标。

在此过程中，没有任何对 pickle 或 cloudpickle 等序列化库的依赖，消除了这些序列化包的广泛限制，例如：

  • 可移植性和兼容性：在与序列化对象所使用的 Python 版本不同的 Python 版本中加载 pickle 或 cloudpickle 文件不能保证兼容性。
  • 复杂对象序列化：文件句柄、socket、外部连接、动态引用、lambda 函数和系统资源无法进行 pickle 序列化。
  • 可读性：Pickle 和 CloudPickle 都将其序列化对象存储为二进制格式，人类无法阅读。
  • 性能：对象序列化和依赖检查可能非常慢，特别是对于具有许多代码引用依赖项的复杂实现。

使用从代码创建模型的核心要求​

使用从代码创建模型功能时需要注意一些重要概念，因为通过脚本记录模型时会执行一些可能不会立即显而易见的操作。

  • 导入：从代码创建模型不捕获非 pip 可安装包的外部引用，就像传统的 cloudpickle 实现一样。如果您有外部引用（参见下面的示例），则必须通过 code_paths 参数定义这些依赖项。
  • 日志记录期间执行：为了验证您正在记录的脚本文件是否有效，代码将在写入磁盘之前执行，这与其他模型日志记录方法完全相同。
  • 依赖推理：如果在您定义的模型脚本顶部导入的包可以从 PyPI 安装，则无论您是否在模型执行逻辑中使用它们，它们都将被推断为依赖项。
提示

如果您定义了在脚本中从未使用过的 import 语句，这些语句仍将包含在依赖列表中。建议在编写实现时使用能够确定未使用 import 语句的 linter，这样就不会包含不相关的包依赖项。

警告

从代码记录模型时，请确保您的代码不包含任何敏感信息，例如 API 密钥、密码或其他机密数据。代码将以纯文本形式存储在 MLflow 模型 Artifact 中，任何有权访问该 Artifact 的人都可以查看代码。

在 Jupyter Notebook 中使用从代码创建模型​

Jupyter (IPython Notebooks) 是处理 AI 应用和建模的非常便捷的方式。它们的一个小限制在于其基于单元格的执行模型。由于其定义和运行方式的性质，从代码创建模型功能不直接支持将 Notebook 定义为模型。相反，此功能要求模型定义为 Python 脚本（文件扩展名必须以 '.py' 结尾）。

幸运的是，维护 Jupyter 使用的核心内核（IPython）的人员创建了许多可在 Notebook 中使用的魔术命令，以增强 Notebook 作为 AI 从业者开发环境的可用性。在任何基于 IPython 的 Notebook 环境（Jupyter、Databricks Notebooks 等）中都可以使用的最有用的魔术命令之一是 %%writefile 命令。

当 %%writefile 魔术命令作为 Notebook 单元格的第一行编写时，将捕获单元格的内容（请注意，不是整个 Notebook，仅是当前单元格范围），除了魔术命令本身之外，并将这些内容写入您定义的文件。

例如，在 Notebook 中运行以下内容：

%%writefile "./hello.py"

print("hello!")

将导致在与您的 Notebook 相同的目录中创建一个文件，该文件包含：

print("hello!")
注意

%%writefile 魔术命令可以使用可选的 -a 附加命令。此选项将把单元格内容附加到目标文件以保存单元格内容。不建议使用此选项，因为它可能会在脚本中创建难以调试的覆盖，脚本中可能包含模型定义逻辑的多个副本。建议使用 %%writefile 的默认行为，即每次执行单元格时都覆盖本地文件，以确保单元格内容的状态始终反映在保存的脚本文件中。

使用从代码创建模型的示例​

以下每个示例都将在脚本定义单元格块的顶部展示 %%writefile 魔术命令的用法，以模拟在单个 Notebook 中定义模型代码或其他依赖项。如果您是在 IDE 或文本编辑器中编写实现，请不要将此魔术命令放在脚本顶部。

构建一个简单的从代码创建模型​

在此示例中，我们将定义一个非常基本的模型，当通过 predict() 调用时，它将利用输入的浮点值作为数字 2 的指数。第一个代码块代表一个独立的 Notebook 单元格，它将在与 Notebook 相同的目录中创建一个名为 basic.py 的文件。此文件的内容将是模型定义 BasicModel，以及 import 语句和用于实例化此模型以进行推理的 MLflow 函数 set_model。

# If running in a Jupyter or Databricks notebook cell, uncomment the following line:
# %%writefile "./basic.py"

import pandas as pd
from typing import List, Dict
from mlflow.pyfunc import PythonModel
from mlflow.models import set_model


class BasicModel(PythonModel):
    def exponential(self, numbers):
        return {f"{x}": 2**x for x in numbers}

    def predict(self, context, model_input) -> Dict[str, float]:
        if isinstance(model_input, pd.DataFrame):
            model_input = model_input.to_dict()[0].values()
        return self.exponential(model_input)


# Specify which definition in this script represents the model instance
set_model(BasicModel())

下一节展示了包含日志记录逻辑的另一个单元格。

import mlflow

mlflow.set_experiment("Basic Model From Code")

model_path = "basic.py"

with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        python_model=model_path,  # Define the model as the path to the script that was just saved
        artifact_path="arithmetic_model",
        input_example=[42.0, 24.0],
    )

在 MLflow UI 中查看此存储的模型，我们可以看到第一个单元格中的脚本被记录为运行的一个 Artifact。

The MLflow UI showing the stored model code as a serialized python script

当我们通过 mlflow.pyfunc.load_model() 加载此模型时，此脚本将被执行，并构建 BasicModel 的一个实例，将 predict 方法暴露为我们进行推理的入口点，就像记录自定义模型的另一种传统模式一样。

my_model = mlflow.pyfunc.load_model(model_info.model_uri)
my_model.predict([2.2, 3.1, 4.7])

# or, with a Pandas DataFrame input
my_model.predict(pd.DataFrame([5.0, 6.0, 7.0]))

从代码创建模型的常见问题解答​

在使用从代码创建模型功能进行模型日志记录时，您应该了解几个方面。虽然其行为与使用传统模型序列化类似，但在开发工作流程和代码架构方面，您需要进行一些显著的调整。

依赖管理和要求​

正确管理依赖项和要求对于确保您的模型可以在新环境中加载或部署至关重要。

从保存的脚本加载模型时，为什么会遇到 NameError？​

在定义脚本（如果在 Notebook 中开发，则为单元格）时，确保所有必需的 import 语句都定义在脚本中。未包含导入依赖项不仅会导致名称解析错误，而且依赖要求也将不会包含在模型的 requirements.txt 文件中。

加载模型时遇到 ImportError。​

如果您的模型定义脚本具有不在 PyPI 上的外部依赖项，则在记录或保存模型时必须使用 code_paths 参数包含这些引用。您可能需要在记录模型时手动将这些外部脚本的导入依赖项添加到 extra_pip_requirements 参数中，以确保模型在加载期间能够获得所有必需的依赖项。

为什么我的 requirements.txt 文件中充满了模型未使用的包？​

MLflow 将根据模块级别的 import 语句从从代码创建模型的脚本构建依赖项列表。没有运行检查过程来验证模型的逻辑是否需要所有声明为 import 的内容。强烈建议精简这些脚本中的 import，仅包含模型正常运行所需的最低限度的 import 语句。过度导入大型包会在加载或部署模型时引入安装延迟，并增加部署的推理环境中的内存压力。

使用从代码创建模型进行日志记录​

从定义的 Python 文件记录模型时，您将遇到与提供对象引用的传统模型序列化过程之间的一些细微差异。

我不小心在脚本中包含了 API Key。我该怎么办？​

由于从代码创建模型功能以纯文本形式存储您的脚本定义，在 MLflow UI 的 Artifact 查看器中完全可见，因此包含敏感数据（如访问密钥或其他基于授权的秘密）存在安全风险。如果您在记录模型时不小心将敏感密钥直接定义在脚本中，建议执行以下操作：

  1. 删除包含泄露密钥的 MLflow 运行。您可以通过 UI 或 delete_run API 来执行此操作。
  2. 删除与该运行关联的 Artifact。您可以通过 mlflow gc CLI 命令来执行此操作。
  3. 通过生成新密钥并从源系统管理界面删除泄露的秘密来轮换您的敏感密钥。
  4. 将模型重新记录到新的运行中，确保不要在模型定义脚本中设置敏感密钥。

为什么记录模型时会执行我的模型？​

为了验证定义模型的 Python 文件中的代码是否可执行，MLflow 将实例化在 set_model API 中定义为模型的对象。如果在模型初始化期间进行了外部调用，这些调用将会执行，以确保您的代码在记录之前是可执行的。如果此类调用需要对服务的认证访问，请确保您记录模型的环境已配置适当的认证，以便您的代码可以运行。

其他资源​

有关可增强您对 MLflow "从代码创建模型" 功能理解的其他相关主题，请考虑查阅 MLflow 文档中的以下部分：