MLflow 插件
作为机器学习的框架无关工具，MLflow Python API 提供了开发者 API，用于编写与不同 ML 框架和后端集成的插件。
插件提供了一种强大的机制，用于自定义 MLflow Python 客户端的行为并集成第三方工具，使您能够
- 与第三方存储解决方案集成，用于存储实验数据、artifacts 和模型
- 与第三方身份验证提供商集成，例如从特殊文件中读取 HTTP 身份验证凭据
- 使用 MLflow 客户端与其它 REST API 通信，例如您组织现有的实验跟踪 API
- 自动捕获额外元数据作为运行标签，例如与运行相关的 git 仓库
- 添加新的后端来执行 MLflow 项目入口点。
MLflow Python API 支持几种类型的插件
- 跟踪存储（Tracking Store）：覆盖跟踪后端逻辑，例如将日志记录到第三方存储解决方案
- Artifact 仓库（ArtifactRepository）：覆盖 artifact 日志记录逻辑，例如将日志记录到第三方存储解决方案
- 运行上下文提供程序（Run context providers）：指定通过
mlflow.start_run()Fluent API 创建的运行上设置的上下文标签。 - 模型注册表存储（Model Registry Store）：覆盖模型注册表后端逻辑，例如将日志记录到第三方存储解决方案
- MLflow 项目后端（MLflow Project backend）：覆盖本地执行后端，以便在您自己的集群（Databricks, kubernetes 等）上执行项目
- MLflow 模型评估器（MLflow ModelEvaluator）：定义自定义模型评估器，可在
mlflow.evaluate()API 中使用。
使用 MLflow 插件
MLflow 插件是可以使用 PyPI 或 conda 安装的 Python 包。本示例从源安装一个跟踪存储插件，并在示例脚本中使用它。
安装插件
要开始使用，请克隆 MLflow 并安装这个示例插件
git clone https://github.com/mlflow/mlflow
cd mlflow
pip install -e tests/resources/mlflow-test-plugin
使用插件运行代码
这个插件定义了一个自定义的跟踪存储，用于跟踪使用 file-plugin 方案的 URI。插件实现委托给 MLflow 内置的基于文件的运行存储。要使用此插件，您可以运行任何使用 MLflow 的代码，并将跟踪 URI 设置为具有 file-plugin:// 方案的 URI
MLFLOW_TRACKING_URI=file-plugin:$(PWD)/mlruns python examples/quickstart/mlflow_tracking.py
启动 MLflow UI
cd ..
mlflow server --backend-store-uri ./mlflow/mlruns
在 http://localhost:5000 查看结果。您应该会看到一个新创建的运行，其中包含一个名为 "param1" 的参数和一个名为 "foo" 的指标
使用插件进行客户端身份验证
MLflow 提供了 RequestAuthProvider 插件，用于自定义传出 http 请求的身份验证头。
要使用它，请实现 RequestAuthProvider 类并覆盖 get_name 和 get_auth 方法。get_name 应该返回您的身份验证提供程序的名称，而 get_auth 应该返回将添加到 http 请求的身份验证对象。
from mlflow.tracking.request_auth.abstract_request_auth_provider import (
RequestAuthProvider,
)
class DummyAuthProvider(RequestAuthProvider):
def get_name(self):
return "dummy_auth_provider_name"
def get_auth(self):
return DummyAuth()
实现请求身份验证提供程序类后，将其注册到 entry_points 中并安装插件。
setup(
entry_points={
"mlflow.request_auth_provider": "dummy-backend=DummyAuthProvider",
},
)
然后设置环境变量 MLFLOW_TRACKING_AUTH 以启用自定义身份验证的注入。此环境变量的值应与身份验证提供程序的名称匹配。
export MLFLOW_TRACKING_AUTH=dummy_auth_provider_name
编写您自己的 MLflow 插件
定义插件
您将 MLflow 插件定义为一个独立的 Python 包，可以通过 PyPI 或 conda 进行分发安装。请参阅https://github.com/mlflow/mlflow/tree/master/tests/resources/mlflow-test-plugin，其中包含一个实现了所有可用插件类型的示例包。
示例包包含一个 setup.py 文件，声明了许多入口点（entry points）
setup(
name="mflow-test-plugin",
# Require MLflow as a dependency of the plugin, so that plugin users can simply install
# the plugin and then immediately use it with MLflow
install_requires=["mlflow"],
...,
entry_points={
# Define a Tracking Store plugin for tracking URIs with scheme 'file-plugin'
"mlflow.tracking_store": "file-plugin=mlflow_test_plugin.file_store:PluginFileStore",
# Define a ArtifactRepository plugin for artifact URIs with scheme 'file-plugin'
"mlflow.artifact_repository": "file-plugin=mlflow_test_plugin.local_artifact:PluginLocalArtifactRepository",
# Define a RunContextProvider plugin. The entry point name for run context providers
# is not used, and so is set to the string "unused" here
"mlflow.run_context_provider": "unused=mlflow_test_plugin.run_context_provider:PluginRunContextProvider",
# Define a RequestHeaderProvider plugin. The entry point name for request header providers
# is not used, and so is set to the string "unused" here
"mlflow.request_header_provider": "unused=mlflow_test_plugin.request_header_provider:PluginRequestHeaderProvider",
# Define a RequestAuthProvider plugin. The entry point name for request auth providers
# is not used, and so is set to the string "unused" here
"mlflow.request_auth_provider": "unused=mlflow_test_plugin.request_auth_provider:PluginRequestAuthProvider",
# Define a Model Registry Store plugin for tracking URIs with scheme 'file-plugin'
"mlflow.model_registry_store": "file-plugin=mlflow_test_plugin.sqlalchemy_store:PluginRegistrySqlAlchemyStore",
# Define a MLflow Project Backend plugin called 'dummy-backend'
"mlflow.project_backend": "dummy-backend=mlflow_test_plugin.dummy_backend:PluginDummyProjectBackend",
# Define a MLflow model deployment plugin for target 'faketarget'
"mlflow.deployments": "faketarget=mlflow_test_plugin.fake_deployment_plugin",
# Define a MLflow model evaluator with name "dummy_evaluator"
"mlflow.model_evaluator": "dummy_evaluator=mlflow_test_plugin.dummy_evaluator:DummyEvaluator",
},
)
这个 entry_points 字典的每个元素指定一个插件。您可以选择在您的包中实现一种或多种插件类型，而不必实现所有类型。下面描述了每个入口点定义的插件类型及其在 MLflow 中的相应参考实现。在编写您自己的插件时，您可以参考这些实现
| 描述 | 入口点组 | 入口点名称和值 | 参考实现 |
|---|---|---|---|
用于覆盖特定跟踪 URI 方案的跟踪 API（如 mlflow.log_metric, mlflow.start_run）定义的插件。 | mlflow.tracking_store | 入口点值（例如 入口点名称（例如 安装示例插件并设置形如 | FileStore |
用于为指定的 artifact URI 方案（例如您的内部 blob 存储系统使用的方案）定义 artifact 读/写 API（如 mlflow.log_artifact, MlflowClient.download_artifacts）的插件。 | mlflow.artifact_repository | 入口点值（例如 入口点名称（例如 安装示例插件并记录到 artifact URI 形如 | LocalArtifactRepository |
| 用于在运行创建时指定自定义上下文标签的插件，例如标识与运行关联的 git 仓库的标签。 | mlflow.run_context_provider | 入口点名称未使用。入口点值（例如 | GitRunContext, DefaultRunContext |
| 用于指定要附加到传出请求的自定义上下文请求头的插件，例如标识客户端环境的请求头。 | mlflow.request_header_provider | 入口点名称未使用。入口点值（例如 | DatabricksRequestHeaderProvider |
| 用于指定要附加到传出请求的自定义请求身份验证的插件。 | mlflow.request_auth_provider | 入口点名称未使用。入口点值（例如 | 暂无（即将添加） |
用于覆盖模型注册表 API（如 mlflow.register_model）定义的插件。 | mlflow.model_registry_store | 入口点值（例如 入口点名称（例如 安装示例插件并设置形如 | SqlAlchemyStore |
| 用于针对自定义执行后端运行 MLflow 项目的插件（例如，针对您的团队内部集群或作业调度器运行项目）。 | mlflow.project.backend | 入口点值（例如 | 暂无（即将添加） |
| 用于将模型部署到自定义服务工具的插件。 | mlflow.deployments | 入口点名称（例如
| PluginDeploymentClient |
| 用于MLflow 模型评估的插件 | mlflow.model_evaluator | 入口点名称（例如
| DummyEvaluator |
| [实验性] 用于自定义 mlflow server flask 应用配置的插件 mlflow.server.app。 | mlflow.app | 入口点 | app |
测试您的插件
我们建议测试您的插件，以确保它遵循 MLflow 预期的契约。例如，跟踪存储插件应包含测试，验证其 log_metric, log_param 等实现的正确性。另请参阅 MLflow 参考实现的测试作为示例
分发您的插件
假设您的插件结构与示例插件类似，您可以通过 PyPI 分发它。
恭喜，您现在已经编写并分发了您自己的 MLflow 插件！
社区插件
SQL Server 插件
mlflow-dbstore 插件允许 MLflow 使用关系型数据库作为 artifact 存储。目前，它只在 SQL Server 作为 artifact 存储时进行过测试。
您可以通过以下方式安装包含 SQL Server 插件的 MLflow：
pip install mlflow[sqlserver]
然后正常使用 MLflow。SQL Server artifact 存储支持将自动提供。
该插件实现了所有 MLflow artifact 存储 API。要使用 SQL Server 作为 artifact 存储，必须提供数据库 URI，如下例所示
db_uri = "mssql+pyodbc://username:password@host:port/database?driver=ODBC+Driver+17+for+SQL+Server"
client.create_experiment(exp_name, artifact_location=db_uri)
mlflow.set_experiment(exp_name)
mlflow.onnx.log_model(onnx, "model")
第一次将 artifact 记录到 artifact 存储中时，插件会自动在数据库 URI 指定的数据库中创建名为 artifacts 的表，并将 artifact 作为 BLOB 存储在那里。后续记录的 artifact 也存储在同一个表中。
在上面提供的示例中，log_model 操作在数据库表中创建三个条目，用于存储 ONNX 模型、MLmodel 文件以及与模型关联的 conda.yaml 文件。
阿里云 (Alibaba Cloud) OSS 插件
aliyunstoreplugin 插件允许 MLflow 使用阿里云 OSS 存储作为 artifact 存储。
pip install mlflow[aliyun-oss]
然后正常使用 MLflow。阿里云 OSS artifact 存储支持将自动提供。
该插件实现了所有 MLflow artifact 存储 API。它期望在 MLFLOW_OSS_ENDPOINT_URL, MLFLOW_OSS_KEY_ID 和 MLFLOW_OSS_KEY_SECRET 环境变量中获取阿里云存储访问凭据，因此您必须在客户端应用和 MLflow 跟踪服务器上都设置这些变量。要使用阿里云 OSS 作为 artifact 存储，必须提供形如 oss://<bucket>/<path> 的 OSS URI，如下例所示
import mlflow
import mlflow.pyfunc
class Mod(mlflow.pyfunc.PythonModel):
def predict(self, context, model_input, params=None):
return 7
exp_name = "myexp"
mlflow.create_experiment(exp_name, artifact_location="oss://mlflow-test/")
mlflow.set_experiment(exp_name)
mlflow.pyfunc.log_model("model_test", python_model=Mod())
在上面提供的示例中，log_model 操作在 OSS 存储 oss://mlflow-test/$RUN_ID/artifacts/model_test/ 中创建三个条目，用于存储 MLmodel 文件以及与模型关联的 conda.yaml 文件。
XetHub 插件
xethub 插件允许 MLflow 使用 XetHub 存储作为 artifact 存储。
pip install mlflow[xethub]
然后正常使用 MLflow。XetHub artifact 存储支持将自动提供。
该插件实现了所有 MLflow artifact 存储 API。它期望通过 xet login CLI 命令或 XET_USER_EMAIL、XET_USER_NAME 和 XET_USER_TOKEN 环境变量获取 XetHub 访问凭据，因此您必须在客户端应用和 MLflow 跟踪服务器上都使用 XetHub 进行身份验证。要使用 XetHub 作为 artifact 存储，必须提供形如 xet://<username>/<repo>/<branch> 的 XetHub URI，如下例所示
import mlflow
import mlflow.pyfunc
class Mod(mlflow.pyfunc.PythonModel):
def predict(self, context, model_input, params=None):
return 7
exp_name = "myexp"
mlflow.create_experiment(
exp_name, artifact_location="xet://<your_username>/mlflow-test/main"
)
mlflow.set_experiment(exp_name)
mlflow.pyfunc.log_model("model_test", python_model=Mod())
在上面提供的示例中，log_model 操作在 OSS 存储 xet://mlflow-test/$RUN_ID/artifacts/model_test/ 中创建三个条目，用于存储 MLmodel 文件以及与模型关联的 conda.yaml 文件。
部署插件
以下已知插件支持使用 MLflow 的模型部署 API 将模型部署到自定义服务工具。请参阅各个插件页面获取安装说明，并参阅Python API 文档和CLI 文档获取使用说明和示例。
- mlflow-redisai
- mlflow-torchserve
- mlflow-algorithmia
- mlflow-ray-serve
- mlflow-azureml
- oci-mlflow 利用 Oracle 云基础设施 (OCI) 模型部署服务部署 MLflow 模型。
- mlflow-jfrog-plugin 通过将 artifact 无缝存储到 JFrog Artifactory 中您偏好的仓库，优化您的 artifact 治理。
模型评估插件
以下已知插件支持使用 MLflow 的mlflow.evaluate() API 使用自定义验证工具评估模型
- mlflow-giskard：在模型投入生产之前，检测从表格到 LLM 的 ML 模型中的隐藏漏洞。预测诸如性能偏差、不稳健性、过度自信、不足自信、伦理偏差、数据泄露、随机性、虚假相关性等问题。通过广泛的测试（自定义或领域特定的测试套件）进行模型比较。
- mlflow-trubrics：使用 Trubrics 验证 ML 模型
项目后端插件
以下已知插件支持针对自定义执行后端运行MLflow 项目。
- mlflow-yarn 在 Hadoop/YARN 上运行 mlflow
- oci-mlflow 在 Oracle 云基础设施 (OCI) 上运行 mlflow 项目
跟踪存储插件
以下已知插件支持针对自定义数据库运行MLflow 跟踪存储。
- mlflow-elasticsearchstore 使用 Elasticsearch 运行 MLflow 跟踪存储
有关此插件的更多信息，请参阅https://github.com/criteo/mlflow-elasticsearchstore/issues。该库可在 PyPI 上找到：https://pypi.ac.cn/project/mlflow-elasticsearchstore/
Artifact 仓库插件
- oci-mlflow 利用 Oracle 云基础设施 (OCI) 对象存储服务存储 MLflow 模型 artifacts。
JFrog Artifactory MLflow 插件
- mlflow-jfrog-plugin 通过将 artifact 无缝存储到 JFrog Artifactory 中您偏好的仓库，优化您的 artifact 治理。
概览
JFrog MLflow 插件通过将 MLflow 的默认 artifacts 位置替换为 JFrog Artifactory 来扩展 MLflow 功能。一旦 MLflow 实验 artifacts 在 JFrog Artifactory 中可用，它们就成为公司发布生命周期的重要组成部分，与任何其他 artifact 一样，并且也受到 JFrog 平台提供的所有安全工具的保护。
特性
- 实验 artifacts 的日志/保存操作针对 JFrog Artifactory 执行
- 使用 MLflow UI 和 API 以及 JFrog UI 和 API 查看和下载实验 artifacts 的操作均针对 JFrog Artifactory 执行
- 实验 Artifacts 的删除遵循实验生命周期（自动或通过 mlflow gc）
- 允许通过实验创建命令更改特定实验 artifacts 的目标位置（通过更改 artifact_location）
安装
使用 pip 安装插件，安装应在 mlflow 跟踪服务器上进行。另外，插件可以安装在任何想要更改特定 artifactory 仓库的默认 artifacts 位置的客户端上
pip install mlflow[jfrog]
或
pip install mlflow-jfrog-plugin
使用 ARTIFACTORY_AUTH_TOKEN 环境变量设置 JFrog Artifactory 身份验证令牌：出于安全原因，最好使用具有所需最低权限的令牌，而不是管理员令牌
export ARTIFACTORY_AUTH_TOKEN=<your artifactory token goes here>
安装插件并设置令牌后，您可以将 JFrog artifactory 仓库作为目标 artifacts 位置启动 mlflow 跟踪服务器。有关其他 mlflow server 选项，请参阅 mlflow 文档
mlflow server --host <mlflow tracking server host> --port <mlflow tracking server port> --artifacts-destination artifactory://<JFrog artifactory URL>/artifactory/<repository[/optional base path]>
为了允许将大型 artifacts 上传到 JFrog artifactory，建议在启动 mlflow server 时增加上传超时设置：--gunicorn-opts '--timeout <timeout in seconds>'
使用方法
MLflow 模型记录代码示例
import mlflow
from mlflow import MlflowClient
from transformers import pipeline
mlflow.set_tracking_uri("<your mlflow tracking server uri>")
mlflow.create_experiment("<your_exp_name>")
classifier = pipeline(
"sentiment-analysis", model="michellejieli/emotion_text_classifier"
)
with mlflow.start_run():
mlflow.transformers.log_model(
transformers_model=classifier, artifact_path=MODEL_NAME
)
mlflow.end_run()
配置
其他可选设置（在 mlflow 跟踪服务器启动前设置）：要使用无 SSL 的 artifactory URL，请将 ARTIFACTORY_NO_SSL 设置为 true。默认为 false
export ARTIFACTORY_NO_SSL=true
要启用 JFrog 操作的调试日志记录，请将 ARTIFACTORY_DEBUG 设置为 true。默认为 false
export ARTIFACTORY_DEBUG=true
为防止 MLflow 垃圾回收从 artifactory 中删除任何 artifacts，请将 ARTIFACTORY_ARTIFACTS_DELETE_SKIP 设置为 true。默认为 false。请注意，此设置可能会导致显著的存储使用量，并且可能需要设置 JFrog 文件保留策略。
export ARTIFACTORY_ARTIFACTS_DELETE_SKIP=true