Mlflow 初探

1. 引言

MLflow是一个由Databricks公司开发并开源，用于管理机器学习生命周期的开源平台。它提供了一组工具，可以帮助数据科学家和工程师跟踪实验、重现结果、部署模型并共享代码。

到目前为止mlflow在github上已经达到了3.6K个fork和15.6K的starts，已经成为机器学习领域中最受欢迎的平台之一。

本文通过查找相关公开资料对mlflow做初步探究，总结其基本概念，使用场景系统架构以及基本使用方法。

2. Mlflow的目标

Mlflow 的目目标在于将机器学习过程中涉及到的步骤与操作都标准化，管理整个机器学习工程的生命周期。

本质上可以将它理解为一套规范流程，技术上它的作用类似于Java开发过程使用到的maven工具： java开发者使用 mvn 命令来管理项目的生命周期，发布制品等；而Mlflow则使用 mlflow 命令对模型训练，调参再训练，然后进入 mlflow UI界面进行对比不同参数的模型效果，发布选中的效果最好的模型等操作，参考下图：

其中，上图中的前几步 DataPreparation，ExploratoryDataAnalysis ，FeatureEngineering 则期望使用开源工具spark处理，这样，整个机器学习的生命周期在databricks公司的产品中闭环。

当前Mlflow已经集成了多种常用的机器学习算法项目，包括 TensorFlow、 xgboost、 sklearn 等，方便用户以极小的工作量快速开展常用机器学习探索工作。同时为了支持方便集成用户自己开发的算法工程，Mlflow定制了一套 project标准规范，一般情况下用户只需要做简单的配置即可将自己的算法工程接入到mlflow。

3. Mlflow的技术架构与实现

在按操作类型主题的维度上划分，Mlflow 项目主要包含5个子模块：

3.1 Mlflow Tracking

MLflow的核心组件之一，用于跟踪机器学习实验。在MLflow Tracking中，每个实验都被称为一个run，每个run可以包含多个参数、指标和模型。这个组件通过API的形式管理实验的参数、代码、结果，并且以Mlflow UI的形式展示，方便用户进行比较。

3.1.1 SDK

Mlflow 封装了一套记录 实验参数，指标以及模型的api，供开发人员使用:

mlflow.log_param() # 记录参数
mlflow.log_metric() # 记录指标
mlflow.log_artifact() # 记录一个本地文件或者文件夹作为一个artifact(制品)
mlflow.log_artifacts()
...

Mlflow Tracking针对常用的编程语言都封装了SDK，目前已经支持的编程语言包括：Python，java，scala，R 。例如以下两段代码分别使用了Python和Java/scala 语言的sdk来记录实验运行过程中的指标：

with mlflow.start_run():
    for epoch in range(0, 3):
        mlflow.log_metric(key="quality", value=2 * epoch, step=epoch)

MlflowClient client = new MlflowClient();
RunInfo run = client.createRun();
for (int epoch = 0; epoch < 3; epoch ++) {
    client.logMetric(run.getRunId(), "quality", 2 * epoch, System.currentTimeMillis(), epoch);
}

• auto_log

目前，Mlflow集成的Scikit-learn、Keras、XGBoost等常用机器学习类库都已经支持 自动记录(auto_log)需要用到的常用实验数据, 这样做的目的就是希望可以让用户快速上手常用的机器学习算法训练，而不需要去编写代码重复定义各种常用指标数据。

3.1.2 Tracking UI

指标保存之后，关键的一步就是将数据可视化，让数据科学家可以观测和评估模型的质量。mlflow 提供了一套Tracking UI组件，可以将各种指标数据按照多种方式可视化，同时这个UI界面具备将 模型制品（model artifact） 打版本，保存到model registry等管理能力。

3.2 Mlflow Models

Mlflow Models模块定义了一组训练、打包、管理、运行模型的api，并提供了命令行操作接口。用户可以通过一条标准的命令即可启动一个训练模型的任务，将生成的模型打包成docker镜像并发布到镜像仓库，将模型以server的模式启动提供 restful api服务等。

# 模型训练
mlflow run --env-manager local examples/sklearn_elasticnet_wine -P alpha=0.5 

# 将模型打包成镜像
mlflow models build-docker --model-uri "runs:/some-run-uuid/my-model" --name "my-image-name" 

# 运行一个模型，指定服务端口
mlflow models serve --env-manager local -h 0.0.0.0 -m my_model

3.3 Mlflow Model Registry

Mlflow Model Registry 是一个集中的模型存储、API和UI，用于全周期地管理模型，可以提供模型血缘、模型版本以及模型的阶段切换。开发者通过该模块管理训练好的模型。

3.4 Mlflow Projects

Mlflow Projects 是一套 开发模型算法的代码规范，mlflow运行一个run(实验)时，通过MLProject文件获取算法工程的名称，环境依赖，入口函数，运行参数等重要信息。一个自定义开发的模型算法模块，包含以下内容：

sklearn_elasticnet_wine
├── conda.yaml        # 可选，通过conda管理python依赖时需要提供
├── MLproject         # mlflow调用该子模块的入口描述文件
├── python_env.yaml   # python环境描述文件
├── train.ipynb       # 工程说明文件(readme)
├── train.py          # 代码文件

入口文件MLProject 文件描述信息如下：

name: tutorial
python_env: python_env.yaml   # 亦可以选择conda.yaml ， docker等环境
entry_points:
  main:
    parameters:
      alpha: {type: float, default: 0.5}
      l1_ratio: {type: float, default: 0.1}
    command: "python train.py {alpha} {l1_ratio}"   # 执行训练代码

以上实例算法模块，可通过命令 mlflow run --env-manager sklearn_elasticnet_wine -P alpha 0.7 来执行。

3.5 Mlflow Recipes

mlflow recipes 的前身是mlflow pipeline，用于机器学习中用到的各种步骤通过recipes编排形成一个可复用且容易拓展的mlflow套件。recipes可以缓存中间运行的结果，使得任务可以从失败节点开始执行，节省算力资源。

3.5.1 核心概念

3.5.1.1 Templates

一个mlflow工作流编排的工程模板目录，里面包含了 配置文件模板recipes.yaml , profile，执行步骤的代码等数据文件。用户开发一个mlflow工作流可以参考templates的目录结构进行修改，官方提供了一个mlflow recipes模板，参考github地址 recipes-regression-template

下面是该模板工程的目录组成：

recipes-regression-template
|-- LICENSE
|-- notebooks               # 工作流编排文件，支持python和java两种api
|   |-- databricks.py
|   `-- jupyter.ipynb
|-- profiles
|   |-- databricks.yaml     # 存放profile实例，运行一个工作流时选择一个profile来填充 recipes.yaml 模板配置文件的参数
|   `-- local.yaml
|-- README.md
|-- recipe.yaml             # 模板配置文件，定义recipes的基本信息： steps，name等
|-- requirements            
|   |-- lint-requirements.txt
|   `-- test-requirements.txt
|-- requirements.txt        # 运行一个recipes需要的 环境依赖信息
|-- steps                   # 工作流步骤对应的代码存放处
|   |-- custom_metrics.py
|   |-- ingest.py
|   |-- split.py
|   |-- train.py
|   `-- transform.py
`-- tests                   # 单元测试
    |-- __init__.py
    |-- train_test.py
    `-- transform_test.py

3.5.1.2 Recipes

编排工作流的核心，提供了python和java两种编程语言的SDK。开发者通过编写代码实现一套工作流程，因此十分容易拓展。编排代码存放到notebooks目录。

# 准备环境依赖：requirements.txt
dbutils.library.restartPython()

from mlflow.recipes import Recipe
# 指定profile，填充recipes模板参数
r = Recipe(profile="databricks")
# 生成dag可视化数据供UI解析
r.inspect()

# 执行步骤 ingest
r.run("ingest")

# 执行步骤 split
r.run("split")

training_data = r.get_artifact("training_data")
training_data.describe()

# 执行步骤 transform
r.run("transform")

# 执行步骤 train
r.run("train")

trained_model = r.get_artifact("model")
print(trained_model)

# 执行步骤，评估模型
r.run("evaluate")

# 执行步骤，注册模型
r.run("register")

3.5.1.3 step

表示工作流程中的一个步骤，即DAG 中的一个节点。

3.5.4 Profiles

工作流配置文件 recipes.yaml 是一个模板，yaml文件可以通过占位符的方式定义参数，参数的值在运行工作流时通过profile注入，profile通过yaml格式编写。

from mlflow.recipes import Recipe
# 指定profile，填充recipes模板参数
r = Recipe(profile="databricks")

3.5.2 可视化

通过调用函数 Recipe.inspect() 来生成DAG数据，供UI解析，如下图所示，深色的transform表示工作流正在运行的进度：

3.6 业务流程及关键组件

从Mlflow自己定义的标准化流程以及技术角度上讨论，则可以将 整个Mlflow项目架构涉及到的关键组件，总结为下图：

4. mlflow的安装和使用

以安装 mlflow-2.7.1 为例，快速体验mlflow：

1. 宿主机安装docker

2. 拉取docker镜像

   docker pull ghcr.io/mlflow/mlflow:v2.7.1 

3. 启动一个mlflow docker 容器, 其中 5001 端口映射预留给mlflow ui使用；映射端口8081预留给模型服务使用

   docker run -d -it --name mlflow_demo -p 5001:5000 -p 8081:8080 ghcr.io/mlflow/mlflow:v2.7.1 /bin/bash

4. github下载 mlflow-2.7.1 版本代码，并拷贝到 mlflow容器内部：

   docker cp mlflow-2.7.1 a4b3d1454fd1:/

5. 进入docker容器，运行一个 算法训练实验

   docker exec -it a4b3d1454fd1 /bin/bash 
   
   mkdir /mlflow_workspace
   cd /mlflow_workspace/
   ## 运行模型训练实验, 使用本地python环境
   export GIT_PYTHON_REFRESH=quiet ; mlflow run --env-manager local /mlflow-2.7.1/examples/sklearn_elasticnet_wine -P alpha=0.7

   运行结束之后，可以看到 mlflow_workspace 下有一个mlruns目录，保存了指标数据以及模型制品。

6. 启动mlflow ui服务

   nohup mlflow ui -h 0.0.0.0 -p 5000 > mlflow_ui.log 2>&1 &

7. 浏览器访问UI

8. 启动模型服务

   mlflow models serve --env-manager local -h 0.0.0.0   --port 8080   -m runs:/ee55b2f4bd844dc99bdff3865ffa3d43/model

   其中 ee55b2f4bd844dc99bdff3865ffa3d43 表示一个runId（实验）， runs:/ 表示从本地中按照实验id查找并加载模型，也可以从s3协议中加载。

9. curl验证模型服务

   curl -d '{"dataframe_split": {"columns": ["x"], "data": [[1], [-1]]}}' -H 'Content-Type: application/json' -X POST localhost:8080/invocations