跳到主要内容

在 MLflow 模型中管理依赖项

MLflow 模型是一种标准格式,用于打包机器学习模型及其依赖项和其他元数据。构建包含其依赖项的模型可确保在各种平台和工具之间的可复现性和可移植性。

当您使用 MLflow 跟踪 API(例如 mlflow.pytorch.log_model())创建 MLflow 模型时,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
mlflow.pyfunc.log_model(
python_model=CustomModel(),
name="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
mlflow.sklearn.log_model(
sk_model=model,
name="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 模型一起保存额外代码依赖项 - 自动推断

注意

目前仅支持 Python Function Models 的代码依赖项自动推断。未来版本的 MLflow 将支持其他命名的模型风格。

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

要使用此新功能,您只需在记录时将参数 infer_code_paths(默认 False)设置为 True。在使用此依赖项推断方法时,您无需像 MLflow 2.13.0 之前那样通过声明 code_paths 目录位置来显式定义文件位置。

下面显示了使用此功能的一个示例,我们正在记录一个包含外部依赖项的模型。在第一部分中,我们定义了一个名为 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(
name="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(),
name="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 选项用于自定义库

要在记录模型时包含 PyPI 上不可公开的自定义库,可以使用 code_paths 参数。此选项允许您将 .whl 文件或其他依赖项与模型一起上传,确保在服务期间所有必需的库都可用。

警告

以下示例演示了一种为开发目的包含自定义库的快速方法。此方法不建议用于生产环境。对于生产用途,请将库上传到自定义 PyPI 服务器或云存储,以确保可靠和安全的访问。

例如,假设您的项目具有以下文件结构

my_project/
|── train.py
└── custom_package.whl

然后以下代码可以记录您的模型以及自定义包

train.py
import mlflow
from custom_package import my_func


class MyModel(mlflow.pyfunc.PythonModel):
def predict(self, context, model_input):
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(),
name="model",
extra_pip_requirements=["code/custom_package.whl"],
input_example=input_data,
code_paths=["custom_package.whl"],
)

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(),
name="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,
name="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(
name="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(
name="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(
name="model1",
python_model=model1,
code_paths=["my_model1.py"],
)

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

考虑到 code_paths 的此限制,推荐的项目结构如下所示

my_project/
|-- model.py # Defines the custom pyfunc model
|── train.py # Trains and logs the model
|── core/ # Required modules for prediction
| |── preprocessing.py
| └── ...
└── helper/ # Other helper modules used for training, evaluation
|── evaluation.py
└── ...

这样,您可以使用 code_paths=["core"] 记录模型,以包含预测所需的模块,同时排除仅用于开发的辅助模块。

验证预测环境

在部署之前验证您的模型是确保生产就绪的关键一步。MLflow 提供了几种在本地测试模型的方法,无论是在虚拟环境还是 Docker 容器中。如果您在验证期间发现任何依赖项问题,请遵循如何修复模型服务时的依赖项错误?中的指南

使用虚拟环境测试离线预测

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

注意

Python API 自 MLflow 2.10.0 起可用。如果您使用的是旧版本,请使用 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 能够支持各种特定于风格的输入类型,例如 TensorFlow 张量。MLflow 还支持不特定于给定风格的类型,例如 Pandas DataFrame、Numpy ndarray、Python Dict、Python List、SciPy 稀疏矩阵和 Spark DataFrame。

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

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

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]]}' \
https://:<port>/invocations

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

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

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

MLflow build-docker CLI 和 Python API 能够构建基于 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]]}' \
https://:<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 参数,该函数接受一个 Python 对象,该对象受您的模型 predict() 函数支持。一些示例可能包括特定于风格的输入类型(例如 tensorflow 张量)或更通用的类型(例如 pandas DataFrame、numpy ndarray、python Dict 或 python List)。在使用 CLI 时,我们不能传递 Python 对象,而是需要传递包含输入有效负载的 CSV 或 JSON 文件的路径。

注意

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

3. 更新模型元数据

找到正确的依赖项后,您可以使用 mlflow.models.update_model_requirements() API 更新已记录模型的依赖项。

import mlflow

mlflow.models.update_model_requirements(
model_uri="runs:/<run_id>/<model_path>",
operation="add",
requirement_list=["opencv-python==4.8.0"],
)

请注意,您还可以利用 CLI 更新模型要求

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

或者,您可以通过在记录模型时指定 extra_pip_requirements 选项来记录一个具有更新依赖项的新模型。

import mlflow

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

如何为许可证变更迁移 Anaconda 依赖项

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

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

如果您在 MLflow v1.18 之前记录了模型,但未将 defaults 频道从模型的 conda 环境中排除,则该模型可能对 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()