管理 MLflow 模型中的依赖项

MLflow 模型 是一种标准格式,它将机器学习模型与其依赖项和其他元数据打包在一起。通过构建包含其依赖项的模型,可以实现跨各种平台和工具的可重复性和可移植性。

当你使用 MLflow Tracking APIs 创建一个 MLflow 模型时,例如 mlflow.pytorch.log_model(),MLflow 会自动推断你所使用的模型风格的所需依赖,并将它们记录为模型元数据的一部分。然后,当你为预测服务模型时,MLflow 会自动将这些依赖安装到环境中。因此,你通常不需要担心在 MLflow 模型中管理依赖。

然而,在某些情况下,您可能需要添加或修改一些依赖项。本页提供了关于 MLflow 如何管理依赖项的高层次描述,以及如何根据您的使用情况自定义依赖项的指导。

小技巧

提高MLflow依赖推断准确性的一个技巧是在保存模型时添加一个 input_example 。这使得MLflow在保存模型之前可以执行一次模型预测,从而捕获预测过程中使用的依赖项。有关此参数的更多详细用法,请参阅 模型输入示例

MLflow 如何记录模型依赖

一个 MLflow 模型保存在指定目录中,具有以下结构:

my_model/
├── MLmodel
├── model.pkl
├── conda.yaml
├── python_env.yaml
└── requirements.txt

模型依赖性由以下文件定义(对于其他文件,请参阅讨论 存储格式 的部分中提供的指导):

  • python_env.yaml - 这个文件包含了使用 virtualenv 恢复模型环境所需的信息(1)Python 版本(2)构建工具,如 pip、setuptools 和 wheel(3)模型的 pip 需求(对 requirements.txt 的引用)

  • requirements.txt - 定义了运行模型所需的 pip 依赖项集合。

  • conda.yaml - 定义了运行模型所需的 conda 环境。当你指定 conda 作为恢复模型环境的环境管理器时,会使用此文件。

请注意,**不建议手动编辑这些文件**以添加或删除依赖项。它们由MLflow自动生成,您手动进行的任何更改在您再次保存模型时将被覆盖。相反,您应该使用以下部分中描述的推荐方法之一。

示例

以下展示了使用 mlflow.sklearn.log_model 记录模型时,MLflow 生成的环境文件示例:

  • python_env.yaml

    python: 3.9.8
    build_dependencies:
    - pip==23.3.2
    - setuptools==69.0.3
    - wheel==0.42.0
    dependencies:
    - -r requirements.txt
    
  • requirements.txt

    mlflow==2.9.2
    scikit-learn==1.3.2
    cloudpickle==3.0.0
    
  • conda.yaml

    name: mlflow-env
    channels:
      - conda-forge
    dependencies:
    - python=3.9.8
    - pip
    - pip:
      - mlflow==2.9.2
      - scikit-learn==1.3.2
      - cloudpickle==3.0.0
    

向 MLflow 模型添加额外依赖项

MLflow 推断模型风格库所需的依赖项,但您的模型可能依赖于其他库,例如数据预处理。在这种情况下,您可以通过在记录模型时指定 extra_pip_requirements 参数来向模型添加额外的依赖项。例如,

import mlflow


class CustomModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        # your model depends on pandas
        import pandas as pd

        ...
        return prediction


# Log the model
with mlflow.start_run() as run:
    mlflow.pyfunc.log_model(
        python_model=CustomModel(),
        artifact_path="model",
        extra_pip_requirements=["pandas==2.0.3"],
        input_example=input_data,
    )

额外的依赖项将按如下方式添加到 requirements.txt 中(同样适用于 conda.yaml):

mlflow==2.9.2
cloudpickle==3.0.0
pandas==2.0.3  # added

在这种情况下,MLflow 在为预测服务模型时,除了推断的依赖项外,还会安装 Pandas 2.0.3。

备注

一旦你记录了带有依赖项的模型,建议在沙盒环境中进行测试,以避免在将模型部署到生产环境时出现任何依赖问题。自 MLflow 2.10.0 起,你可以使用 mlflow.models.predict() API 在虚拟环境中快速测试你的模型。更多详情请参阅 验证预测环境

自行定义所有依赖项

或者,您也可以从头定义所有依赖项,而不是添加额外的依赖项。为此,请在记录模型时指定 pip_requirements。例如,

import mlflow

# Log the model
with mlflow.start_run() as run:
    mlflow.sklearn.log_model(
        sk_model=model,
        artifact_path="model",
        pip_requirements=[
            "mlflow-skinny==2.9.2",
            "cloudpickle==2.5.8",
            "scikit-learn==1.3.1",
        ],
    )

手动定义的依赖项将覆盖 MLflow 从模型风格库中检测到的默认依赖项:

mlflow-skinny==2.9.2
cloudpickle==2.5.8
scikit-learn==1.3.1

警告

在声明与训练期间使用的依赖项不同的依赖项时,请务必小心,因为这可能很危险且容易导致意外行为。确保一致性的最安全方法是依赖 MLflow 推断的默认依赖项。

备注

一旦你记录了带有依赖项的模型,建议在沙盒环境中进行测试,以避免在将模型部署到生产环境时出现任何依赖问题。自 MLflow 2.10.0 起,你可以使用 mlflow.models.predict() API 在虚拟环境中快速测试你的模型。更多详情请参阅 验证预测环境

使用 MLflow 模型保存额外代码依赖 - 自动推理

备注

自动代码依赖推断是 MLflow 2.13.0 中引入的一个功能,并被标记为实验性。基础实现可能会在没有事先通知的情况下进行修改、改进和调整,以解决潜在问题和边缘情况。

备注

目前仅支持对 Python 函数模型进行自动代码依赖推断。未来 MLflow 的版本中将增加对其他命名模型风格的支持。

在 MLflow 2.13.0 版本中,引入了一种包含自定义依赖代码的新方法,该方法扩展了在保存或记录模型时声明 code_paths 的现有功能。这一新功能利用导入依赖分析,通过检查 Python 模型定义中的引用模块来自动推断模型所需的代码依赖。

为了使用这个新功能,你可以在记录时简单地将参数 infer_code_paths (默认 False)设置为 True。在使用这种依赖推断方法时,你不需要通过声明 code_paths 目录位置来明确地定义文件位置,就像在MLflow 2.13.0之前你必须做的那样。

使用此功能的示例如下所示,其中我们正在记录一个包含外部依赖的模型。在第一部分中,我们定义了一个名为 custom_code 的外部模块,该模块存在于与我们的模型定义不同的位置。

custom_code.py
from typing import List

iris_types = ["setosa", "versicolor", "viginica"]


def map_iris_types(predictions: int) -> List[str]:
    return [iris_types[pred] for pred in predictions]

定义了 custom_code.py 模块后,它就可以在我们的 Python 模型中使用了:

model.py
from typing import Any, Dict, List, Optional

from custom_code import map_iris_types  # import the external reference

import mlflow


class FlowerMapping(mlflow.pyfunc.PythonModel):
    """Custom model with an external dependency"""

    def predict(
        self, context, model_input, params: Optional[Dict[str, Any]] = None
    ) -> List[str]:
        predictions = [pred % 3 for pred in model_input]

        # Call the external function
        return map_iris_types(predictions)


with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        artifact_path="flowers",
        python_model=FlowerMapping(),
        infer_code_paths=True,  # Enabling automatic code dependency inference
    )

infer_code_paths 设置为 True 时,将分析 map_iris_types 的依赖关系,其源声明将被检测为源自 custom_code.py 模块,并且 custom_code.py 中的代码引用将与模型工件一起存储。请注意,通过使用 code_paths 参数(在下一节中讨论)来定义外部代码依赖关系是不需要的。

小技巧

只有当前工作目录内的模块是可访问的。依赖推断不会跨模块边界工作,或者如果你的自定义代码定义在完全不同的库中。如果你的代码库结构使得公共模块完全位于你的模型日志记录代码执行路径之外,那么需要使用原始的 code_paths 选项来记录这些依赖关系,因为 infer_code_paths 依赖推断不会捕捉到这些需求。

infer_code_paths 的限制

警告

在使用 infer_code_paths 进行依赖推断之前,请确保您的依赖代码模块中没有硬编码的敏感数据(例如,密码、访问令牌或密钥)。代码推断不会混淆敏感信息,并且会捕获并记录(保存)模块,无论其内容是什么。

在使用 infer_code_paths 时需要注意的一个重要方面是避免在代码的主入口点内定义依赖关系。当一个 Python 代码文件作为 __main__ 模块加载时,它不能被推断为代码路径文件。这意味着如果你直接运行你的脚本(例如,使用 python script.py),该脚本中定义的函数和类将成为 __main__ 模块的一部分,而不能被其他模块轻松访问。

如果你的模型依赖于这些类或函数,这可能会带来问题,因为它们不是标准模块命名空间的一部分,因此不容易序列化。要处理这种情况,你应该使用 cloudpickle 来序列化你的模型实例。cloudpickle 是 Python 的 pickle 模块的扩展版本,可以序列化更广泛的 Python 对象,包括在 __main__ 模块中定义的函数和类。

为何这很重要
  • 代码路径推断:MLflow 使用代码路径来理解和记录与您的模型相关的代码。当脚本作为 __main__ 执行时,代码路径无法推断,这使得MLflow实验的跟踪和可重复性变得复杂。

  • 序列化:像 pickle 这样的标准序列化方法可能无法处理 __main__ 模块对象,导致在尝试保存和加载模型时出现问题。cloudpickle 通过启用这些对象的序列化提供了一个解决方案,确保您的模型可以正确地保存和恢复。

最佳实践
  • 避免在 __main__ 模块中定义关键函数和类。相反,将它们放在可以按需导入的单独模块文件中。

  • 如果你必须在 __main__ 模块中定义函数和类,请使用 cloudpickle 来序列化你的模型,以确保所有依赖项都被正确处理。

使用 MLflow 模型保存额外代码 - 手动声明

MLflow 还支持将您的自定义 Python 代码保存为模型的依赖项。当您想要部署与模型预测所需的定制模块时,这特别有用。为此,请在记录模型时指定 code_paths。例如,如果您的项目中有以下文件结构:

my_project/
├── utils.py
└── train.py
train.py
import mlflow


class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        from utils import my_func

        x = my_func(model_input)
        # .. your prediction logic
        return prediction


# Log the model
with mlflow.start_run() as run:
    mlflow.pyfunc.log_model(
        python_model=MyModel(),
        artifact_path="model",
        input_example=input_data,
        code_paths=["utils.py"],
    )

然后 MLflow 会将 utils.py 保存在模型目录下的 code/ 目录中:

model/
├── MLmodel
├── ...
└── code/
    └── utils.py

当 MLflow 加载模型以进行服务时,code 目录将被添加到系统路径中,以便您可以在模型代码中使用模块,例如 from utils import my_func。您还可以将目录路径指定为 code_paths,以在目录下保存多个文件:

code_paths 选项的注意事项

在使用 code_paths 选项时,请注意指定的文件或目录 必须与您的模型脚本在同一目录中 。如果指定的文件或目录位于父目录或子目录中,例如 my_project/src/utils.py,模型服务将会失败并抛出 ModuleNotFoundError。例如,假设您的项目中有以下文件结构

my_project/
|── train.py
└── src/
    └──  utils.py

那么,以下模型代码 起作用:

class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        from src.utils import my_func

        # .. your prediction logic
        return prediction


with mlflow.start_run() as run:
    mlflow.pyfunc.log_model(
        python_model=MyModel(),
        artifact_path="model",
        input_example=input_data,
        code_paths=[
            "src/utils.py"
        ],  # the file will be saved at code/utils.py not code/src/utils.py
    )

# => Model serving will fail with ModuleNotFoundError: No module named 'src'

此限制是由于 MLflow 如何保存和加载指定的文件和目录。当它在 code/ 目标中复制指定的文件或目录时,它**不会**保留它们原本所在的相对路径。例如,在上面的例子中,MLflow 会将 utils.py 复制到 code/utils.py,而不是 code/src/utils.py。因此,必须将其导入为 from utils import my_func,而不是 from src.utils import my_func。然而,这可能不太理想,因为导入路径与原始训练脚本不同。

要解决这个问题,code_paths 应该指定父目录,在本例中是 code_paths=["src"]。这样,MLflow 会将整个 src/ 目录复制到 code/ 下,您的模型代码将能够导入 src.utils

class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        from src.utils import my_func

        # .. your prediction logic
        return prediction


with mlflow.start_run() as run:
    mlflow.pyfunc.log_model(
        python_model=model,
        artifact_path="model",
        input_example=input_data,
        code_paths=["src"],  # the whole /src directory will be saved at code/src
    )

警告

出于同样的原因,code_paths 选项无法处理 code_paths=["../src"] 的相对导入。

在使用 code_paths 加载具有相同模块名称但不同实现的多模型时的限制

code_paths 选项的当前实现有一个限制,即它不支持在同一个 Python 进程中加载依赖于具有相同名称但不同实现的模块的多个模型,如下例所示:

import importlib
import sys
import tempfile
from pathlib import Path

import mlflow

with tempfile.TemporaryDirectory() as tmpdir:
    tmpdir = Path(tmpdir)
    my_model_path = tmpdir / "my_model.py"
    code_template = """
import mlflow

class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        return [{n}] * len(model_input)
"""

    my_model_path.write_text(code_template.format(n=1))

    sys.path.insert(0, str(tmpdir))
    import my_model

    # model 1
    model1 = my_model.MyModel()
    assert model1.predict(context=None, model_input=[0]) == [1]

    with mlflow.start_run():
        info1 = mlflow.pyfunc.log_model(
            artifact_path="model",
            python_model=model1,
            code_paths=[my_model_path],
        )

    # model 2
    my_model_path.write_text(code_template.format(n=2))
    importlib.reload(my_model)
    model2 = my_model.MyModel()
    assert model2.predict(context=None, model_input=[0]) == [2]

    with mlflow.start_run():
        info2 = mlflow.pyfunc.log_model(
            artifact_path="model",
            python_model=model2,
            code_paths=[my_model_path],
        )

# To simulate a fresh Python process, remove the `my_model` module from the cache
sys.modules.pop("my_model")

# Now we have two models that depend on modules with the same name but different implementations.
# Let's load them and check the prediction results.
pred = mlflow.pyfunc.load_model(info1.model_uri).predict([0])
assert pred == [1], pred  # passes

# As the `my_model` module was loaded and cached in the previous `load_model` call,
# the next `load_model` call will reuse it and return the wrong prediction result.
assert "my_model" in sys.modules
pred = mlflow.pyfunc.load_model(info2.model_uri).predict([0])
assert pred == [2], pred  # doesn't pass, `pred` is [1]

为了绕过这个限制,你可以在加载模型之前从缓存中移除模块。例如:

model1 = mlflow.pyfunc.load_model(info1.model_uri)
sys.modules.pop("my_model")
model2 = mlflow.pyfunc.load_model(info2.model_uri)

另一种解决方法是使用不同的模块名来区分不同的实现。例如:

mlflow.pyfunc.log_model(
    artifact_path="model1",
    python_model=model1,
    code_paths=["my_model1.py"],
)

mlflow.pyfunc.log_model(
    artifact_path="model",
    python_model=model2,
    code_paths=["my_model2.py"],
)

预测环境的验证

在部署之前验证您的模型是确保生产就绪的关键步骤。MLflow 提供了几种在本地测试模型的方法,无论是在虚拟环境中还是在 Docker 容器中。如果在验证过程中发现任何依赖性问题,请按照 如何在服务我的模型时修复依赖错误? 中的指导进行操作。

在虚拟环境中测试离线预测

你可以通过 Python 或 CLI 使用 MLflow Models 的 predict API 来对你的模型进行测试预测。这将根据模型 URI 加载你的模型,创建一个包含模型依赖项(在 MLflow Model 中定义)的虚拟环境,并使用该模型进行离线预测。请参考 mlflow.models.predict()CLI 参考 以获取 predict API 的更多详细用法。

备注

自 MLflow 2.10.0 起,Python API 可用。如果您使用的是较旧版本,请使用 CLI 选项。

import mlflow

mlflow.models.predict(
    model_uri="runs:/<run_id>/model",
    input_data=<input_data>,
)

使用 mlflow.models.predict() API 对于快速测试您的模型和推理环境非常方便。然而,它可能不是服务的一个完美模拟,因为它不会启动在线推理服务器。也就是说,这是一个测试您的预测输入是否正确格式化的好方法。

格式化取决于您的已记录模型的 predict() 方法所支持的类型。如果模型是使用签名记录的,输入数据应该可以从 MLflow UI 中查看,或者通过 mlflow.models.get_model_info() 查看,该方法具有 signature 字段。

更一般地,MLflow 有能力支持各种特定风格(flavor-specific)的输入类型,例如 tensorflow 张量。MLflow 还支持不特定于某种风格的类型,例如 pandas DataFrame、numpy ndarray、python Dict、python List、scipy.sparse 矩阵和 spark 数据框。

使用虚拟环境测试在线推理端点

如果你想通过实际运行在线推理服务器来测试你的模型,你可以使用 MLflow 的 serve API。这将创建一个包含你的模型和依赖项的虚拟环境,类似于 predict API,但会启动推理服务器并暴露 REST 端点。然后你可以发送测试请求并验证响应。请参考 CLI 参考 以获取 serve API 的更多详细用法。

mlflow models serve -m runs:/<run_id>/model -p <port>
# In another terminal
curl -X POST -H "Content-Type: application/json" \
    --data '{"inputs": [[1, 2], [3, 4]]}' \
    http://localhost:<port>/invocations

虽然这是一种在部署前测试模型的可靠方法,但一个需要注意的是,虚拟环境不会吸收你的机器与生产环境之间的操作系统级别的差异。例如,如果你使用MacOS作为本地开发机器,但你的部署目标是在Linux上运行,你可能会遇到一些在虚拟环境中无法重现的问题。

在这种情况下,您可以使用 Docker 容器来测试您的模型。虽然它不像虚拟机那样提供完整的操作系统级隔离,例如我们不能在 Linux 机器上运行 Windows 容器,但 Docker 涵盖了一些流行的测试场景,例如运行不同版本的 Linux 或在 Mac 或 Windows 上模拟 Linux 环境。

使用 Docker 容器测试在线推理端点

MLflow 的 build-docker API 用于 CLI 和 Python,能够为您的模型构建一个基于 Ubuntu 的 Docker 镜像。该镜像将包含您的模型和依赖项,以及用于启动推理服务器的入口点。与 serve API 类似,您可以发送测试请求并验证响应。有关 build-docker API 的更多详细用法,请参阅 CLI 参考

mlflow models build-docker -m runs:/<run_id>/model -n <image_name>
docker run -p <port>:8080 <image_name>
# In another terminal
curl -X POST -H "Content-Type: application/json" \
    --data '{"inputs": [[1, 2], [3, 4]]}' \
    http://localhost:<port>/invocations

故障排除

如何修复在服务我的模型时遇到的依赖错误

在模型部署过程中遇到的最常见问题之一是依赖问题。当记录或保存模型时,MLflow 会尝试推断模型依赖关系,并将它们保存为 MLflow 模型元数据的一部分。然而,这可能并不总是完整的,并且可能会遗漏某些依赖项,例如某些库的 [extras] 依赖项。这可能会在服务模型时导致错误,例如“ModuleNotFoundError”或“ImportError”。以下是一些可以帮助诊断和修复缺失依赖错误的方法。

提示

为了减少依赖错误的可能性,您可以在保存模型时添加 input_example。这使得 MLflow 能够在保存模型之前执行模型预测,从而捕获预测过程中使用的依赖项。有关此参数的更多详细用法,请参阅 模型输入示例

1. 检查缺失的依赖项

缺失的依赖项列在错误信息中。例如,如果你看到以下错误信息:

ModuleNotFoundError: No module named 'cv2'

2. 尝试使用 predict API 添加依赖项

既然你知道了缺失的依赖项,你可以创建一个带有正确依赖项的新模型版本。然而,为了尝试新的依赖项而创建新模型可能会有些繁琐,特别是因为你可能需要多次迭代才能找到正确的解决方案。相反,你可以使用 mlflow.models.predict() API 来测试你的更改,而无需在解决安装错误时反复重新记录模型。

为此,使用 pip-requirements-override 选项来指定 pip 依赖项,例如 opencv-python==4.8.0

import mlflow

mlflow.models.predict(
    model_uri="runs:/<run_id>/<model_path>",
    input_data=<input_data>,
    pip_requirements_override=["opencv-python==4.8.0"],
)

指定的依赖项将被安装到虚拟环境中,除了(或代替)模型元数据中定义的依赖项。由于这不会改变模型,您可以快速且安全地迭代以找到正确的依赖项。

请注意,在Python实现中,input_data 参数的函数接受一个由模型 predict() 函数支持的Python对象。一些例子可能包括特定类型的输入,如tensorflow张量,或更通用的类型,如pandas DataFrame、numpy ndarray、python字典或python列表。在使用CLI时,我们无法传递Python对象,而是传递包含输入负载的CSV或JSON文件的路径。

备注

自 MLflow 2.10.0 起,可以使用 pip-requirements-override 选项。

3. 更新模型元数据

一旦你找到正确的依赖项,你可以用这些正确的依赖项创建一个新模型。为此,在记录模型时指定 extra_pip_requirements 选项。

import mlflow

mlflow.pyfunc.log_model(
    artifact_path="model",
    python_model=python_model,
    extra_pip_requirements=["opencv-python==4.8.0"],
    input_example=input_data,
)

请注意,您也可以利用 CLI 就地更新模型依赖项,从而避免重新记录模型。

mlflow models update-pip-requirements -m runs:/<run_id>/<model_path> add "opencv-python==4.8.0"

如何迁移Anaconda依赖以适应许可证变更

Anaconda Inc. 更新了他们针对 anaconda.org 频道的 服务条款。根据新的服务条款,如果您依赖 Anaconda 的打包和分发,您可能需要一个商业许可证。更多信息请参见 Anaconda 商业版常见问题。您使用任何 Anaconda 频道的行为都受其服务条款的约束。

v1.18 之前记录的 MLflow 模型默认使用 conda defaults 频道 (https://repo.anaconda.com/pkgs/) 作为依赖项。由于此许可证变更,MLflow 已停止对使用 MLflow v1.18 及以上版本记录的模型使用 defaults 频道。现在默认记录的频道是 conda-forge,它指向社区管理的 https://conda-forge.org/

如果你在MLflow v1.18之前记录了一个模型,并且没有从模型的conda环境中排除``defaults``通道,那么该模型可能会有一个你可能不希望的``defaults``通道依赖。要手动确认一个模型是否具有此依赖关系,你可以检查与记录模型一起打包的``conda.yaml``文件中的``channel``值。例如,一个具有``defaults``通道依赖的模型的``conda.yaml``可能看起来像这样:

name: mlflow-env
channels:
- defaults
dependencies:
- python=3.8.8
- pip
- pip:
    - mlflow==2.3
    - scikit-learn==0.23.2
    - cloudpickle==1.6.0

如果你想更改模型环境中使用的通道,可以通过使用新的 conda.yaml 重新注册模型到模型注册表。你可以通过在 log_model()conda_env 参数中指定通道来实现这一点。

有关 log_model() API 的更多信息,请参阅您正在使用的模型风格的 MLflow 文档,例如 mlflow.sklearn.log_model()