在 Spark 作业中部署和运行 MLflow 模型

项目
02/18/2024

本文介绍如何在 Spark 作业中部署和运行 MLflow 模型，以针对大量数据或在数据整理过程中执行推理。

关于此示例

此示例演示如何将 Azure 机器学习中注册的 MLflow 模型部署到托管 Spark 群集（预览）中运行的 Spark 作业、Azure Databricks 或 Azure Synapse Analytics，以便对大量数据执行推理。

该模型基于 UCI 心脏病数据集。数据库包含 76 个属性，但我们使用其中 14 个。该模型尝试预测患者是否存在心脏疾病。它是从 0（不存在）到 1（存在）的整数值。它已使用 XGBBoost 分类器进行训练，所有必需的预处理都打包为 scikit-learn 管道，使此模型成为从原始数据到预测的端到端管道。

本文中的信息基于 azureml-examples 存储库中包含的代码示例。若要在不复制/粘贴文件的情况下在本地运行命令，请克隆存储库，然后将目录更改为 sdk/using-mlflow/deploy。

git clone https://github.com/Azure/azureml-examples --depth 1
cd sdk/python/using-mlflow/deploy

先决条件

在按照本文中的步骤操作之前，请确保满足以下先决条件：

安装 MLflow SDK 包 mlflow 和适用于 MLflow azureml-mlflow 的 Azure 机器学习插件。
```
pip install mlflow azureml-mlflow
```
提示

可以使用 mlflow-skinny 包，它是一个不带 SQL 存储、服务器、UI 或数据科学依赖项的轻型 MLflow 包。对于主要需要用到 MLflow 的跟踪和记录功能但不需要导入整个功能套件（包括部署）的用户，建议使用 mlflow-skinny。
一个 Azure 机器学习工作区。可以按照创建机器学习资源教程创建一个。
- 查看在工作区中执行 MLflow 操作所需的访问权限。
如果要执行远程跟踪（即跟踪在 Azure 机器学习外部运行的试验），请将 MLflow 配置为指向 Azure 机器学习工作区的跟踪 URI。有关如何将 MLflow 连接到工作区的详细信息，请参阅为 Azure 机器学习配置 MLflow。

必须在工作区中注册 MLflow 模型。具体而言，此示例将注册针对糖尿病数据集训练的模型。

连接到工作区

首先，我们连接到要在其中注册模型的 Azure 机器学习工作区。

Azure 机器学习计算
远程计算

已为你配置跟踪。使用 MLflow 时，还会使用你的默认凭据。

配置跟踪 URI

获取工作区的跟踪 URI：
- Azure CLI
- Python
- 工作室
- 手动
适用于：Azure CLI ml 扩展 v2（当前版）
1. 登录并配置工作区：
```
az account set --subscription <subscription>
az configure --defaults workspace=<workspace> group=<resource-group> location=<location> 
```
2. 可以使用 az ml workspace 命令获取跟踪 URI：
```
az ml workspace show --query mlflow_tracking_uri
```
适用于：Python SDK azure-ai-ml v2（当前版本）

可以使用适用于 Python 的 Azure 机器学习 SDK v2 获取 Azure ML MLflow 跟踪 URI。确保在使用的计算中安装了库 azure-ai-ml。以下示例获取与工作区关联的唯一 MLFLow 跟踪 URI。
1. 使用 MLClient 登录到工作区。要执行此操作，更简单方法是使用工作区配置文件：
```
from azure.ai.ml import MLClient
from azure.identity import DefaultAzureCredential

ml_client = MLClient.from_config(credential=DefaultAzureCredential())
```
  提示
  
  可以通过以下方式下载工作区配置文件：
  
  导航到 Azure ML 工作室
  
  单击页面右上角 -> 下载配置文件。
  
  将文件 config.json 保存到正在使用的同一目录中。
2. 或者，也可以使用订阅 ID、资源组名称和工作区名称来获取：
```
from azure.ai.ml import MLClient
from azure.identity import DefaultAzureCredential

#Enter details of your AzureML workspace
subscription_id = '<SUBSCRIPTION_ID>'
resource_group = '<RESOURCE_GROUP>'
workspace_name = '<WORKSPACE_NAME>'

ml_client = MLClient(credential=DefaultAzureCredential(),
                        subscription_id=subscription_id, 
                        resource_group_name=resource_group,
                        workspace_name=workspace_name)
```
  重要
  
  DefaultAzureCredential 将尝试从可用上下文中拉取凭据。如果想要以其他方式指定凭据，例如以交互式方式使用 Web 浏览器，则可以使用 InteractiveBrowserCredential 或 azure.identity 包中提供的任何其他方法。
3. 获取 Azure 机器学习跟踪 URI：
```
mlflow_tracking_uri = ml_client.workspaces.get(ml_client.workspace_name).mlflow_tracking_uri
```
使用 Azure 机器学习门户获取跟踪 URI：
1. 打开 Azure 机器学习工作室并使用你的凭据登录。
2. 在右上角，单击工作区的名称以显示“目录 + 订阅 + 工作区”边栏选项卡。
3. 单击“在 Azure 门户中查看所有属性”。
4. 可以在“概要”部分中找到“MLflow 跟踪 URI”属性。
Azure 机器学习跟踪 URI 可以使用订阅 ID、资源部署区域、资源组名称和工作区名称来生成。以下代码示例显示了如何：

警告

如果在启用了专用链接的工作区中工作，MLflow 终结点还将使用专用链接与 Azure 机器学习通信。因此，跟踪 URI 看起来会有所不同，如此处所示。在这些情况下，需要使用 Azure ML SDK 或 CLI v2 获取跟踪 URI。
```
region = "<LOCATION>"
subscription_id = '<SUBSCRIPTION_ID>'
resource_group = '<RESOURCE_GROUP>'
workspace_name = '<AML_WORKSPACE_NAME>'

mlflow_tracking_uri = f"azureml://{region}.api.azureml.ms/mlflow/v1.0/subscriptions/{subscription_id}/resourceGroups/{resource_group}/providers/Microsoft.MachineLearningServices/workspaces/{workspace_name}"
```
配置跟踪 URI：
- 使用 MLflow SDK
- 使用环境变量
然后，set_tracking_uri() 方法将 MLflow 跟踪 URI 指向该 URI。
```
import mlflow

mlflow.set_tracking_uri(mlflow_tracking_uri)
```
可以在计算中设置 MLflow 环境变量 MLFLOW_TRACKING_URI，从而将与该计算中的 MLflow 的任何交互默认指向 Azure 机器学习。
```
MLFLOW_TRACKING_URI=$(az ml workspace show --query mlflow_tracking_uri | sed 's/"//g') 
```
提示

使用共享环境（如 Azure Databricks 群集、Azure Synapse Analytics 群集或类似环境）时，可以在群集级别设置环境变量 MLFLOW_TRACKING_URI 以自动配置 MLflow 跟踪 URI，以指向群集中运行的基于所有会话的 Azure 机器学习，而不是指向群集中运行的基于每会话的 Azure 机器学习。

配置身份验证

配置跟踪后，还需要配置对关联工作区进行身份验证的方式。默认情况下，适用于 MLflow 的 Azure 机器学习插件将通过打开默认浏览器提示输入凭据来执行交互式身份验证。有关在 Azure 机器学习工作区中为 MLflow 配置身份验证的其他方法，请参阅为 Azure 机器学习配置 MLflow：配置身份验证。

如果交互式作业有用户连接到会话，你可以依赖交互式身份验证，因此无需进一步操作。

警告

在系统提示输入凭据时，交互式浏览器身份验证将阻止代码执行。此选项不适用于无人参与的环境（如训练作业）中的身份验证。建议配置其他身份验证模式。

如果需要无人参与的执行，必须将服务主体配置为与 Azure 机器学习通信。

MLflow SDK
使用环境变量

import os

os.environ["AZURE_TENANT_ID"] = "<AZURE_TENANT_ID>"
os.environ["AZURE_CLIENT_ID"] = "<AZURE_CLIENT_ID>"
os.environ["AZURE_CLIENT_SECRET"] = "<AZURE_CLIENT_SECRET>"

export AZURE_TENANT_ID="<AZURE_TENANT_ID>"
export AZURE_CLIENT_ID="<AZURE_CLIENT_ID>"
export AZURE_CLIENT_SECRET="<AZURE_CLIENT_SECRET>"

提示

使用共享环境时，建议在计算中配置这些环境变量。最佳做法是尽可能在 Azure Key Vault 实例中将其作为机密进行管理。例如，在 Azure Databricks 中，可以在群集配置的环境变量中使用机密，如下所示：AZURE_CLIENT_SECRET={{secrets/<scope-name>/<secret-name>}}。请参阅引用环境变量中的机密，或参阅平台中的类似文档，了解如何在 Azure Databricks 中执行此操作。

注册模型

我们需要在 Azure 机器学习注册表中注册的模型来执行推理。在这种情况下，我们已在存储库中拥有模型的本地副本，因此我们只需要将模型发布到工作区中的注册表。如果打算部署的模型已注册，则可以跳过此步骤。

model_name = 'heart-classifier'
model_local_path = "model"

registered_model = mlflow_client.create_model_version(
    name=model_name, source=f"file://{model_local_path}"
)
version = registered_model.version

或者，如果模型是在运行中记录的，则你可以直接注册它。

提示

若要注册模型，需要知道模型的存储位置。如果使用的是 MLflow 的 autolog 功能，则路径将取决于所使用模型的类型和框架。建议检查作业输出以确定该文件夹的名称。可以查找包含名为 MLModel 的文件的文件夹。如果使用 log_model 手动记录模型，则路径是传递给此类方法的参数。例如，如果使用 mlflow.sklearn.log_model(my_model, "classifier") 记录模型，则存储模型的路径为 classifier。

model_name = 'heart-classifier'

registered_model = mlflow_client.create_model_version(
    name=model_name, source=f"runs://{RUN_ID}/{MODEL_PATH}"
)
version = registered_model.version

注意

路径 MODEL_PATH 是模型在运行中存储的位置。

获取要评分的输入数据

我们需要一些用于运行作业的输入数据。在本示例中，我们将从 Internet 下载示例数据，并将其放置在 Spark 群集使用的共享存储中。

import urllib

urllib.request.urlretrieve("https://azuremlexampledata.blob.core.chinacloudapi.cn/data/heart-disease-uci/data/heart.csv", "/tmp/data")

将数据移动到可供整个群集使用的已装载存储帐户。

dbutils.fs.mv("file:/tmp/data", "dbfs:/")

重要

前面的代码使用 dbutils，这是 Azure Databricks 群集中可用的工具。根据所使用的平台使用适当的工具。

然后，输入数据放置在以下文件夹中：

input_data_path = "dbfs:/data"

在 Spark 群集中运行模型

以下部分介绍如何在 Spark 作业中运行在 Azure 机器学习中注册的 MLflow 模型。

确保在群集中安装以下库：

  - mlflow<3,>=2.1
  - cloudpickle==2.2.0
  - scikit-learn==1.2.0
  - xgboost==1.7.2

我们将使用笔记本来演示如何使用 Azure 机器学习中注册的 MLflow 模型来创建评分例程。创建笔记本并使用 PySpark 作为默认语言。

导入所需的命名空间：

import mlflow
import pyspark.sql.functions as f

配置模型 URI。以下 URI 在其最新版本中引入名为 heart-classifier 的模型。
```
model_uri = "models:/heart-classifier/latest"
```
将模型加载为 UDF 函数。用户定义函数 (UDF) 是由用户定义的函数，用于在用户环境中重复使用自定义逻辑。
```
predict_function = mlflow.pyfunc.spark_udf(spark, model_uri, result_type='double') 
```
提示

使用参数 result_type 控制 predict() 函数返回的类型。
读取要评分的数据：
```
df = spark.read.option("header", "true").option("inferSchema", "true").csv(input_data_path).drop("target")
```
在本示例中，输入数据采用 CSV 格式并放置在 dbfs:/data/ 文件夹中。我们还要删除 target 列，因为此数据集包含要预测的目标变量。在生产方案中，数据不包含此列。
运行函数 predict_function，并将预测结果放在新列中。在本示例中，我们要将预测放在 predictions 列中。
```
df.withColumn("predictions", score_function(*df.columns))
```
提示

predict_function 会接收所需的列作为参数。在本示例中，数据帧的所有列都是模型预期，因此使用了 df.columns。如果模型需要列的子集，可以手动将其引入。如果模型具有签名，则类型需要在输入与预期类型之间兼容。

可以将预测写回到存储：

scored_data_path = "dbfs:/scored-data"
scored_data.to_csv(scored_data_path)

在 Azure 机器学习的独立 Spark 作业中运行模型

Azure 机器学习支持创建独立的 Spark 作业，以及创建可在 Azure 机器学习管道中使用的可重用 Spark 组件。在本例中，我们将部署评分作业（在 Azure 机器学习独立 Spark 作业中运行并运行 MLflow 模型）来执行推理。

注意

若要详细了解 Azure 机器学习中的 Spark 作业，请参阅在 Azure 机器学习中提交 Spark 作业（预览版）。

Spark 作业需要采用参数的 Python 脚本。创建评分脚本：

score.py

import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--model")
parser.add_argument("--input_data")
parser.add_argument("--scored_data")

args = parser.parse_args()
print(args.model)
print(args.input_data)

# Load the model as an UDF function
predict_function = mlflow.pyfunc.spark_udf(spark, args.model, env_manager="conda")

# Read the data you want to score
df = spark.read.option("header", "true").option("inferSchema", "true").csv(input_data).drop("target")

# Run the function `predict_function` and place the predictions on a new column
scored_data = df.withColumn("predictions", score_function(*df.columns))

# Save the predictions
scored_data.to_csv(args.scored_data)

上述脚本采用 --model、--input_data 和 --scored_data 这三个参数。前两个是输入，表示要运行的模型和输入数据，最后一个是输出，这是用于放置预测的输出文件夹。

提示

Python 包安装：前面的评分脚本将 MLflow 模型加载到 UDF 函数中，但它指示参数 env_manager="conda"。设置此参数后，MLflow 将在仅有 UDF 函数运行的隔离环境中还原模型定义中指定的所需包。有关更多详细信息，请参阅 mlflow.pyfunc.spark_udf 文档。

创建作业定义：

mlflow-score-spark-job.yml

$schema: http://azureml/sdk-2-0/SparkJob.json
type: spark

code: ./src
entry:
  file: score.py

conf:
  spark.driver.cores: 1
  spark.driver.memory: 2g
  spark.executor.cores: 2
  spark.executor.memory: 2g
  spark.executor.instances: 2

inputs:
  model:
    type: mlflow_model
    path: azureml:heart-classifier@latest
  input_data:
    type: uri_file
    path: https://azuremlexampledata.blob.core.chinacloudapi.cn/data/heart-disease-uci/data/heart.csv
    mode: direct

outputs:
  scored_data:
    type: uri_folder

args: >-
  --model ${{inputs.model}}
  --input_data ${{inputs.input_data}}
  --scored_data ${{outputs.scored_data}}

identity:
  type: user_identity

resources:
  instance_type: standard_e4s_v3
  runtime_version: "3.2"

提示

若要使用附加的 Synapse Spark 池，请在上面所示的示例 YAML 规范文件中定义 compute 属性，而不是 resources 属性。

上面显示的 YAML 文件可以在带有 --file 参数的 az ml job create 命令中使用，以创建独立的 Spark 作业，如下所示：
```
az ml job create -f mlflow-score-spark-job.yml
```