MLflow 插件

作为一个与框架无关的机器学习工具，MLflow Python API 提供了开发者API，用于编写与不同机器学习框架和后端集成的插件。

插件提供了一种强大的机制，用于定制 MLflow Python 客户端的行为并集成第三方工具，使您能够：

集成第三方存储解决方案以处理实验数据、工件和模型
集成第三方身份验证提供者，例如从特殊文件中读取HTTP身份验证凭据
使用 MLflow 客户端与其他 REST API 通信，例如您组织的现有实验跟踪 API。
自动捕获额外的元数据作为运行标签，例如与运行相关的git仓库
添加新的后端以执行 MLflow 项目入口点。

MLflow Python API 支持多种类型的插件：

跟踪存储：覆盖跟踪后端逻辑，例如，记录到第三方存储解决方案
ArtifactRepository: 覆盖工件日志记录逻辑，例如记录到第三方存储解决方案
运行上下文提供者：指定通过 mlflow.start_run() 流式 API 创建的运行中要设置的上下文标签。
模型注册存储：覆盖模型注册后端逻辑，例如，记录到第三方存储解决方案
MLflow 项目后端：覆盖本地执行后端以在您自己的集群（Databricks、Kubernetes 等）上执行项目。
MLflow ModelEvaluator: 定义自定义模型评估器，可以在 mlflow.evaluate() API 中使用。

使用 MLflow 插件

MLflow 插件是你可以使用 PyPI 或 conda 安装的 Python 包。这个示例从源代码安装了一个 Tracking Store 插件，并在一个示例脚本中使用它。

安装插件

要开始使用，请克隆 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，它声明了若干入口点：

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	入口点值（例如 `mlflow_test_plugin.local_store:PluginFileStore`）指定了一个 mlflow.tracking.store.AbstractStore 的自定义子类（例如，`mlflow_test_plugin` 模块中的 PluginFileStore 类）。入口点名称（例如 `file-plugin`）是与自定义 AbstractStore 实现相关联的跟踪 URI 方案。安装示例插件并设置形式为 `file-plugin://<path>` 的跟踪 URI 的用户将使用在 `PluginFileStore` 中定义的自定义 AbstractStore 实现。完整的跟踪 URI 被传递给 `PluginFileStore` 构造函数。	FileStore
用于定义工件读/写API的插件，如 `mlflow.log_artifact` ， `MlflowClient.download_artifacts` ，用于指定的工件URI方案（例如，您内部使用的blob存储系统所使用的方案）。	mlflow.artifact_repository	入口点值（例如 `mlflow_test_plugin.local_artifact:PluginLocalArtifactRepository`）指定了一个 mlflow.store.artifact.artifact_repo.ArtifactRepository 的自定义子类（例如，`mlflow_test_plugin` 模块中的 PluginLocalArtifactRepository 类）。入口点名称（例如 `file-plugin`）是与自定义 ArtifactRepository 实现相关联的工件 URI 方案。安装了示例插件并登录到其工件URI为 `file-plugin://<path>` 格式的运行的用户，将使用在 `PluginLocalArtifactRepository` 中定义的自定义 ArtifactRepository 实现。完整的工件URI被传递给 `PluginLocalArtifactRepository` 构造函数。	LocalArtifactRepository
用于在创建运行时指定自定义上下文标签的插件，例如，标识与运行关联的 git 仓库的标签。	mlflow.run_context_provider	入口点名称未使用。入口点值（例如 `mlflow_test_plugin.run_context_provider:PluginRunContextProvider`）指定了一个 mlflow.tracking.context.abstract_context.RunContextProvider 的自定义子类（例如，`mlflow_test_plugin` 模块中的 PluginRunContextProvider 类）以进行注册。	GitRunContext, DefaultRunContext
用于指定自定义上下文请求头的插件，以便附加到外发请求中，例如标识客户端环境的头信息。	mlflow.request_header_provider	入口点名称未使用。入口点值（例如 `mlflow_test_plugin.request_header_provider:PluginRequestHeaderProvider`）指定了一个 mlflow.tracking.request_header.abstract_request_header_provider.RequestHeaderProvider 的自定义子类（例如，`mlflow_test_plugin` 模块中的 PluginRequestHeaderProvider 类）以进行注册。	DatabricksRequestHeaderProvider 是一个请求头提供者，用于在Databricks环境中跟踪MLflow请求。
用于指定自定义请求认证以附加到传出请求的插件。	mlflow.request_auth_provider	入口点名称未使用。入口点值（例如 `mlflow_test_plugin.request_auth_provider:PluginRequestAuthProvider`）指定了一个 mlflow.tracking.request_auth.abstract_request_auth_provider.RequestAuthProvider 的自定义子类（例如，`mlflow_test_plugin` 模块中的 PluginRequestAuthProvider 类）以进行注册。	N/A (即将添加)
用于覆盖模型注册API定义的插件，例如 `mlflow.register_model`。	mlflow.model_registry_store	入口点值（例如 `mlflow_test_plugin.sqlalchemy_store:PluginRegistrySqlAlchemyStore`）指定了一个 mlflow.tracking.model_registry.AbstractStore 的自定义子类（例如，`mlflow_test_plugin` 模块中的 PluginRegistrySqlAlchemyStore 类）入口点名称（例如 `file-plugin`）是与自定义 AbstractStore 实现相关联的跟踪 URI 方案。安装示例插件并设置形式为 `file-plugin://<path>` 的跟踪 URI 的用户将使用在 `PluginFileStore` 中定义的自定义 AbstractStore 实现。完整的跟踪 URI 被传递给 `PluginFileStore` 构造函数。	SqlAlchemyStore
用于针对自定义执行后端运行 MLflow 项目的插件（例如，针对团队内部集群或作业调度程序运行项目）。	mlflow.project.backend	入口点值（例如 `mlflow_test_plugin.dummy_backend:PluginDummyProjectBackend`）指定了一个 `mlflow.project.backend.AbstractBackend` 的自定义子类。	N/A (即将添加)
用于将模型部署到自定义服务工具的插件。	mlflow.deployments	入口点名称（例如 `redisai`）是目标名称。入口点值（例如 `mlflow_test_plugin.fake_deployment_plugin`）指定一个模块，该模块定义：1) 恰好一个 mlflow.deployments.BaseDeploymentClient 的子类（例如，PluginDeploymentClient 类）。MLflow 的 `mlflow.deployments.get_deploy_client` API 直接向用户返回该子类的实例，因此鼓励您在插件实现中编写清晰的用户面向的方法和类文档字符串。2) `run_local` 和 `target_help` 函数，不包括 `target` 参数，如此处所示。	PluginDeploymentClient.
用于 MLflow 模型评估的插件	mlflow.model_evaluator	入口点名称（例如 `dummy_evaluator`）是评估器的名称，用于 `mlflow.evaluate` API 的 `evaluators` 参数中。入口点值（例如 `dummy_evaluator:DummyEvaluator`）必须引用 `mlflow.models.evaluation.ModelEvaluator` 的子类；该子类必须实现两个方法：1) `can_evaluate`：接受仅限关键字的参数 `model_type` 和 `evaluator_config`。如果评估器可以评估具有指定评估器配置的指定模型类型，则返回 `True`。否则返回 `False`。2) `evaluate`：计算并记录指标和工件，返回评估结果作为 `mlflow.models.EvaluationResult` 的实例。接受以下参数：model``（一个 pyfunc 模型实例），``model_type``（与 :py:func:`mlflow.evaluate()` 中的 ``model_type 参数相同），dataset``（一个包含特征和标签（可选）的 ``mlflow.data.evaluation_dataset._EvaluationDataset 实例，用于模型评估），``run_id``（记录结果的 MLflow Run 的 ID），以及 ``evaluator_config``（评估器的附加配置字典）。	DummyEvaluator.
[实验性] 自定义 mlflow 服务器 flask 应用配置的插件 mlflow.server.app。	mlflow.app	入口点 `<app_name>=<object_reference>` （例如 `custom_app=mlflow_test_plugin.app:app`）指定了一个自定义的Flask应用。这对于实现身份验证/授权的请求钩子、自定义日志记录和自定义Flask配置非常有用。插件必须导入 mlflow.server.app （例如 `from mlflow.server import app`），并且可以向应用添加自定义配置、中间件等。插件应避免更改现有的应用路由、处理程序和环境变量，以避免意外行为。安装了示例插件的用户将拥有一个自定义的Flask应用。要运行自定义的Flask应用，请使用 `mlflow server --app-name <app_name>`。	app.

测试您的插件

我们建议测试您的插件，以确保它遵循MLflow预期的契约。例如，一个跟踪存储插件应包含测试，验证其 log_metric、log_param 等实现的正确性。另请参阅MLflow参考实现的测试作为示例：

分发您的插件

假设你已经按照示例插件的结构组织了你的插件，你可以通过 PyPI 分发它。

恭喜，您现在已经编写并分发了您自己的 MLflow 插件！

该插件实现了所有 MLflow 工件存储 API。它期望在 MLFLOW_OSS_ENDPOINT_URL、MLFLOW_OSS_KEY_ID 和 MLFLOW_OSS_KEY_SECRET 环境变量中提供阿里云存储访问凭证，因此您必须在客户端应用程序和 MLflow 跟踪服务器上设置这些变量。要使用阿里云 OSS 作为工件存储，必须提供形式为 oss://<bucket>/<path> 的 OSS URI，如下例所示：

import mlflow
import mlflow.pyfunc


class Mod(mlflow.pyfunc.PythonModel):
    def predict(self, ctx, inp, 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 存储作为工件存储。

pip install mlflow[xethub]

然后像往常一样使用 MLflow。XetHub 工件存储支持将自动提供。

该插件实现了所有 MLflow 工件存储 API。它通过 xet login CLI 命令或 XET_USER_EMAIL、XET_USER_NAME 和 XET_USER_TOKEN 环境变量期望 XetHub 访问凭证，因此您必须使用 XetHub 对您的客户端应用程序和 MLflow 跟踪服务器进行身份验证。要使用 XetHub 作为工件存储，必须提供形式为 xet://<username>/<repo>/<branch> 的 XetHub URI，如下例所示：

import mlflow
import mlflow.pyfunc


class Mod(mlflow.pyfunc.PythonModel):
    def predict(self, ctx, inp, 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 是一个与 RedisAI 集成的 MLflow 插件。
mlflow-torchserve 的中文翻译
mlflow-algorithmia 是一个与 Algorithmia 集成的 MLflow 插件。
mlflow-ray-serve 是一个项目。
mlflow-azureml 的翻译
oci-mlflow 利用 Oracle 云基础设施 (OCI) 模型部署服务来部署 MLflow 模型。
mlflow-jfrog-plugin 通过无缝地将您的工件存储在JFrog Artifactory中的首选仓库中，优化您的工件治理。

模型评估插件

以下已知插件通过使用 MLflow 的 mlflow.evaluate() API 提供了对使用自定义验证工具评估模型的支持：

mlflow-giskard: 在投入生产之前，检测从表格到LLMs的ML模型中的隐藏漏洞。预见诸如性能偏差、不稳健性、过度自信、自信不足、伦理偏差、数据泄露、随机性、虚假相关等问题。通过自定义或领域特定的测试套件，使用广泛的测试进行模型比较。
mlflow-trubrics: 使用 Trubrics 验证 ML 模型

项目后端插件

以下已知的插件提供对在自定义执行后端上运行 MLflow 项目的支持。

mlflow-yarn 在 Hadoop/YARN 上运行 mlflow
oci-mlflow 在 Oracle 云基础设施 (OCI) 上运行 mlflow 项目

跟踪存储插件

以下已知的插件提供了对运行 MLflow Tracking Store 的支持，以针对自定义数据库。

mlflow-elasticsearchstore 使用 Elasticsearch 运行 MLflow 跟踪存储

关于此插件的更多信息，请参阅 <https://github.com/criteo/mlflow-elasticsearchstore/issues>。该库在 PyPI 上可获取，地址为：<https://pypi.org/project/mlflow-elasticsearchstore/>

工件仓库插件

oci-mlflow 利用 Oracle 云基础设施 (OCI) 对象存储服务来存储 MLflow 模型工件。

JFrog Artifactory MLflow 插件

mlflow-jfrog-plugin 通过无缝地将您的工件存储在JFrog Artifactory中的首选仓库中，优化您的工件治理。

概述

JFrog MLflow 插件通过将 MLflow 的默认工件位置替换为 JFrog Artifactory 来扩展 MLflow 功能。一旦 MLflow 实验工件在 JFrog Artifactory 中可用，它们就成为公司发布生命周期中不可或缺的一部分，就像其他任何工件一样，并且还受到 JFrog 平台提供的所有安全工具的保护。

特性

实验工件日志/保存是针对JFrog Artifactory进行的
使用 MLflow UI 和 API 以及 JFrog UI 和 API 对 JFrog Artifactory 进行实验工件的查看和下载。
实验工件的删除遵循实验生命周期（自动或通过 mlflow gc）
通过实验创建命令（通过更改 artifact_location）允许更改特定实验的工件目标位置。

安装

使用 pip 安装插件，安装应在 mlflow 跟踪服务器上完成。可选地，插件可以安装在任何希望更改特定制品库默认制品位置的客户端上。

pip install mlflow[jfrog]

或

pip install mlflow-jfrog-plugin

设置 JFrog Artifactory 认证令牌，使用 ARTIFACTORY_AUTH_TOKEN 环境变量：出于安全考虑，最好使用具有最低权限要求的令牌，而不是管理员令牌。

export ARTIFACTORY_AUTH_TOKEN=<your artifactory token goes here>

一旦插件安装并设置好令牌，您的 mlflow 跟踪服务器就可以启动，并将 JFrog 制品库作为目标制品目的地。使用 mlflow 文档获取更多 mlflow 服务器选项。

mlflow server --host <mlflow tracking server host> --port <mlflow tracking server port> --artifacts-destination artifactory://<JFrog artifactory URL>/artifactory/<repository[/optional base path]>

为了允许将大型工件上传到 JFrog 工件库，建议在启动 mlflow 服务器时增加上传超时设置：–gunicorn-opts ‘–timeout <超时时间（秒）>’

用法

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 中移除的工件，将 ARTIFACTORY_ARTIFACTS_DELETE_SKIP 设置为 true。默认值为 false。请注意，此设置可能会导致显著的存储使用，并可能需要 JFrog 文件保留设置。

export ARTIFACTORY_ARTIFACTS_DELETE_SKIP=true