在 AWS 中国区把 MLflow 模型端到端部署在 SageMaker

机器学习项目里有一个反复出现的断层：模型在 notebook 里跑得好好的，一旦要"上线"提供推理服务，事情就复杂起来了。实验怎么管理、模型怎么版本化、训练好的工件放哪、推理服务又该怎么部署——每一环都有自己的工具和坑。MLflow + SageMaker 是解决这一断层的一套经典组合：MLflow 管实验与工件，SageMaker 提供托管的推理端点。以下是通用的流程图示

本文的目的，是把这条链路从训练到推理完整跑一遍，用到的项目是一个房价预测模型（Kaggle House Prices 数据集），用 ElasticNet 训练，最终部署成一个 SageMaker 实时推理端点。

源码文件地址：https://gitee.com/zhaojiew-realm/files/blob/master/housing-price.7z

下图是测试环境涉及到的整体架构，三个环节各司其职：

通常来说，模型的训练需要在 SageMaker Notebook 上，测试方便期间我们在本地进行训练。此外，MLflow 跟踪服务器也起在本地。只有真正需要 AWS 托管能力的两样东西留在云上——S3 存模型工件、SageMaker 端点做推理，外加 ECR 存推理镜像。

初始化环境

MLflow 的 mlflow.projects.run() 支持 conda / virtualenv / local 三种 env_manager。我们用 uv 建好环境并装齐依赖，然后让 MLflow 用 env_manager="local" 直接复用当前环境，而不是让 MLflow 自己去建环境。

uv venv --python 3.10 .venv
source .venv/bin/activate
# mlflow 3.14.0、scikit-learn 1.7.2、xgboost 3.2.0。
uv pip install --index-url https://mirrors.aliyun.com/pypi/simple/ \
    mlflow scikit-learn xgboost pandas numpy boto3

对应地，run.py 里要改两处：

mlflow.set_tracking_uri("http://127.0.0.1:5000")   

mlflow.projects.run(
    uri=".",
    entry_point="Training",
    experiment_name="ElasticNet",
    env_manager="local"
)

训练与实验跟踪

数据预处理（data.py）对数值列做 KNN 插值、类别列用众数填充，再用 OneHotEncoder(handle_unknown='ignore') 编码。

import pandas as pd
from sklearn.impute import KNNImputer
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OneHotEncoder

train = pd.read_csv('data/train.csv')
test = pd.read_csv('data/test.csv')

X = train.drop('SalePrice', axis=1)
y = train['SalePrice']
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)

numeric_cols = X_train.select_dtypes(include=['int64', 'float64']).columns
non_numeric_cols = X_train.select_dtypes(exclude=['int64', 'float64']).columns
imputer = KNNImputer()
X_train[numeric_cols] = imputer.fit_transform(X_train[numeric_cols])
X_val[numeric_cols] = imputer.transform(X_val[numeric_cols])
test[numeric_cols] = imputer.transform(test[numeric_cols])

for column in non_numeric_cols:
    X_train[column].fillna(X_train[column].mode()[0], inplace=True)
    X_val[column].fillna(X_val[column].mode()[0], inplace=True)
    test[column].fillna(test[column].mode()[0], inplace=True)

ohe = OneHotEncoder(drop='first', handle_unknown='ignore')
X_train = ohe.fit_transform(X_train)
X_val = ohe.transform(X_val)
test = ohe.transform(test)

训练脚本 train.py 遍历 ElasticNet 的 18 组超参数组合（alpha × l1_ratio × fit_intercept = 3×3×2），每组开一个 MLflow run，记录参数、RMSE/MAPE/R2 指标和模型工件：

import mlflow
from data import X_train, X_val, y_train, y_val
from sklearn.linear_model import ElasticNet
from sklearn.model_selection import ParameterGrid
from params import elasticnet_param_grid
from utils import eval_metrics

for params in ParameterGrid(elasticnet_param_grid):
    with mlflow.start_run():
        lr = ElasticNet(**params)
        lr.fit(X_train, y_train)
        y_pred = lr.predict(X_val)
        metrics = eval_metrics(y_val, y_pred)
        mlflow.log_params(params)
        mlflow.log_metrics(metrics)
        mlflow.sklearn.log_model(
            lr, "ElasticNet",
            input_example=X_train,
            code_paths=['train.py', 'data.py', 'params.py', 'utils.py'],
        )

其中 eval_metrics（utils.py）把三个回归指标打包成字典：

import numpy as np
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_percentage_error

def eval_metrics(y_true, y_pred):
    rmse = np.sqrt(mean_squared_error(y_true, y_pred))
    mape = mean_absolute_percentage_error(y_true, y_pred)
    r2 = r2_score(y_true, y_pred)
    return {"RMSE": rmse, "MAPE": mape, "R2": r2}

后台启动 python run.py，等它跑完 18 个 run。用 MLflow API 挑出最佳模型：

总 run 数: 18 | FINISHED: 18
=== 最佳模型 (R2 最高) ===
R2: 0.8226 | RMSE: 36889.91 | MAPE: 0.1429
params: alpha=0.1 l1_ratio=0.8 fit_intercept=True

在mlflow ui界面查看具体的训练运行

以及具体训练的指标，可以得知本次训练的source代码，使用的数据集和超参数等信息，做到心里有数。

MLflow 3.x 的模型存储结构

要部署最佳模型，得先找到它的工件在哪。MLflow 3.x 不再按 run_id 组织模型，而是按 model_id。工件在

mlartifacts/1/models/m-a1b42d5a6ab940b9b6d593b5fbd8486f/artifacts/

想拿到 run 对应的 model_id，得用 search_logged_models API：

lms = client.search_logged_models(
    experiment_ids=['1'],
    filter_string=f"source_run_id='{run_id}'")
# → model_id: m-a1b42d5a6ab940b9b6d593b5fbd8486f

看一眼这个模型的 MLmodel 文件，有几个信息后面部署要用到：它是 skops 序列化的 sklearn 模型（不是 pickle），签名是 tensor [-1, 7940] 输入、[-1] 输出。

构建推理镜像

先把最佳模型上传到 S3（SageMaker 部署要求模型可从 S3 访问）。

aws s3api create-bucket --bucket mlflow-project-artifacts-<accountID> \
  --region cn-north-1 \
  --create-bucket-configuration LocationConstraint=cn-north-1

aws s3 cp \
  mlartifacts/1/models/m-a1b42d5a6ab940b9b6d593b5fbd8486f/artifacts/ \
  s3://mlflow-project-artifacts-<accountID>/1/cf6908fac6404d099eab94260c24dcc6/artifacts/ElasticNet/ \
  --recursive --region cn-north-1

在构建之前，首先改一些配置来加速。

mlflow/models/docker_utils.py的核心是一个模板字符串 _DOCKERFILE_TEMPLATE。需要在 FROM 之后加两行 ENV，让镜像内所有 pip 都走阿里源：

FROM {base_image}

# 中国区：让镜像内所有 pip install（构建时+运行时）走阿里源
ENV PIP_INDEX_URL=https://mirrors.aliyun.com/pypi/simple/
ENV PIP_TRUSTED_HOST=mirrors.aliyun.com

之后构建卡在 git clone阶段，因为MLflow 默认的 base image 是 ubuntu:22.04，用 virtualenv env_manager 时会走 SETUP_PYENV 路径，从 github.com clone pyenv。

[4/11] RUN git clone ... https://github.com/pyenv/pyenv.git /root/.pyenv
fatal: unable to access 'https://github.com/pyenv/pyenv.git/':
       GnuTLS recv error (-110): The TLS connection was non-properly terminated.

MLflow 中如果 base image 是 python:*-slim，generate_dockerfile 会走一条不装 pyenv的路径。因为此时Python 已经内置，只 apt install nginx 即可。于是改 mlflow/sagemaker/cli.py，把硬编码的 UBUNTU_BASE_IMAGE 换成 python:3.10-slim：

base_image=mlflow.models.docker_utils.PYTHON_SLIM_BASE_IMAGE.format(version="3.10"),

顺手把 python-slim 分支的 apt 源换成阿里云，并补上 xgboost 运行时需要的 libgomp1 和编译工具：

setup_python_venv_steps = (
    "RUN sed -i 's|deb.debian.org|mirrors.aliyun.com|g; "
    "s|security.debian.org|mirrors.aliyun.com|g' "
    "/etc/apt/sources.list /etc/apt/sources.list.d/* 2>/dev/null; "
    "apt-get -y update && apt-get install -y --no-install-recommends "
    "nginx build-essential libgomp1"
)

构建成功后，在 push 阶段出现报错

2026/... INFO mlflow.sagemaker: Logging in to ECR registry:
         <accountID>.dkr.ecr.cn-north-1.amazonaws.com
Error response from daemon: Get "https://<accountID>.dkr.ecr.cn-north-1.amazonaws.com/v2/":
       dial tcp: lookup ... no such host

这是由于MLflow 的 push_image_to_ecr 硬编码了 .amazonaws.com，漏了中国区的 .cn 后缀。

_full_template = "{account}.dkr.ecr.{region}.amazonaws.com/{image}:{version}"
registry = f"{account}.dkr.ecr.{region}.amazonaws.com"

改 MLflow 源码在 push 里做 region 判断可行，但要动好几处。所以这里直接手动将构建好的镜像推送到 ecr：

ECR="<accountID>.dkr.ecr.cn-north-1.amazonaws.com.cn"
aws ecr get-login-password --region cn-north-1 \
  | docker login --username AWS --password-stdin $ECR
docker tag mlflow-pyfunc:latest "$ECR/mlflow-pyfunc:3.14.0"
docker push "$ECR/mlflow-pyfunc:3.14.0"

部署端点

在部署阶段运行 deploy.py，因为 boto3 client 带了 region_name=cn-north-1，会自动解析到中国区端点，日志里全是 .amazonaws.com.cn 和 arn:aws-cn:...。端点资源创建成功，状态进入 Creating。

from mlflow.deployments import get_deploy_client

endpoint_name = "prod-endpoint"
model_uri = "s3://mlflow-project-artifacts-<accountID>/1/cf6908fac6404d099eab94260c24dcc6/artifacts/ElasticNet"

config = {
    "execution_role_arn": "arn:aws-cn:iam::<accountID>:role/AmazonSageMaker-ExecutionRole-all",
    "bucket_name": "mlflow-project-artifacts-<accountID>",
    "image_url": "<accountID>.dkr.ecr.cn-north-1.amazonaws.com.cn/mlflow-pyfunc:3.14.0",
    "region_name": "cn-north-1",
    "archive": False,
    "instance_type": "ml.m5.xlarge",
    "instance_count": 1,
    "synchronous": True,
}

client = get_deploy_client("sagemaker")
client.create_deployment(
    name=endpoint_name,
    model_uri=model_uri,
    flavor="python_function",
    config=config,
)

创建的sagemaker model如下

然后就一直 Creating……14 分钟过去还是 Creating。正常的端点 5-10 分钟就该 InService 了。这时候不能干等，得去看容器日志。SageMaker 的容器日志在 CloudWatch，发现了如下报错：

mlflow.exceptions.MlflowException: Could not find the pyenv binary.
  File ".../container/__init__.py", in _serve_pyfunc
  File ".../container/__init__.py", in _install_pyfunc_deps
  File ".../container/__init__.py", in _install_model_dependencies_to_env
  File ".../utils/virtualenv.py", in _get_or_create_virtualenv
  File ".../utils/virtualenv.py", in _validate_pyenv_is_available

又是 pyenv。 但这次是在运行时。这正是换 base image 埋下的副作用——为了让构建通过，我们把 ubuntu（装 pyenv）换成了 python-slim（没有 pyenv）。构建完成，可 MLflow 的 sagemaker 容器在启动时会用 virtualenv env_manager 为模型创建一个独立环境，这需要 pyenv。slim 镜像里没有，容器起不来，反复崩溃重试，端点卡在 Creating 出不来。

那么为什么容器要在运行时建环境？检查 mlflow/models/container/__init__.py 的 _serve_pyfunc，关键逻辑是这样的：

def _serve_pyfunc(model, env_manager):
    disable_env_creation = MLFLOW_DISABLE_ENV_CREATION.get()
    if pyfunc.ENV in conf:
        # 只有 disable_env_creation 为 False 时，才在容器启动时建环境装依赖
        if not disable_env_creation:
            _install_pyfunc_deps(MODEL_PATH, install_mlflow=True, env_manager=env_manager)
        ...
    # 否则直接用镜像里的系统 Python 跑 scoring server

而 mlflow/sagemaker/cli.py 构建镜像时传的是 disable_env_creation_at_runtime=False——也就是故意让容器在运行时才装模型依赖。这在标准场景是合理的（模型运行时才从 S3 加载，依赖也运行时才知道），但它依赖 pyenv 来建 virtualenv。

因此需要设置 disable_env_creation_at_runtime=True + env_manager='local'，并在构建时把模型依赖预装进镜像。这样 _serve_pyfunc 会跳过 _install_pyfunc_deps，直接用镜像里的系统 Python 跑推理服务，彻底不需要 pyenv。

模型的 requirements.txt 精确列出了它需要什么：

mlflow==3.14.0
numpy==2.2.6
pandas==2.3.3
pyarrow==24.0.0
scikit-learn==1.7.2
scipy==1.15.3
skops==0.14.0

于是改 cli.py 两处：把运行时的最小依赖安装步骤，换成构建时安装这套精确依赖；同时把 disable_env_creation_at_runtime 改成 True：

setup_container = (
    "# 把模型依赖焊进镜像（运行时禁用环境创建）\n"
    "RUN pip install --no-cache-dir mlflow==3.14.0 scikit-learn==1.7.2 "
    "skops==0.14.0 numpy==2.2.6 pandas==2.3.3 scipy==1.15.3 pyarrow==24.0.0 "
    "gunicorn[gevent]"
)
# generate_dockerfile(..., disable_env_creation_at_runtime=True, ...)

用 --env-manager local 重新构建、手动推送到 ECR。删掉那个卡死的失败端点，用新镜像重新部署。这次，2 分半钟就 InService：

The deployment operation completed successfully with message:
"The SageMaker endpoint was created successfully."

推理测试

MLflow pyfunc 端点接受的是 {'inputs': ...}。

import json
import boto3
from data import test

smrt = boto3.client('sagemaker-runtime', region_name='cn-north-1')
test_data_json = json.dumps({'inputs': test[:20].toarray().tolist()})
prediction = smrt.invoke_endpoint(
    EndpointName="prod-endpoint",
    Body=test_data_json,
    ContentType='application/json')
print(prediction['Body'].read().decode("ascii"))

跑起来，端点返回了 20 条房价预测：

{"predictions": [99125.59, 147858.70, 182158.72, 192778.37, 201997.38,
                 171229.13, 164781.89, 164718.58, 182586.14, 128734.96,
                 187770.69, 93003.38, 94895.99, 145722.78, 106453.83,
                 339276.71, 263255.26, 306766.60, 306053.89, 410739.32]}

数值落在 9.9 万到 41 万美元区间，合理。完整链路——本地训练 → MLflow 跟踪 → 构建推送镜像 → 部署端点 → 推理——全部打通。

结语

从"notebook 里能跑"到"云端能推理"，中间隔着的从来不是一个 deploy.py 那么简单。尤其当环境偏离教程的隐含假设时那些藏在框架里的假设会一个接一个地暴露出来。这次实战真正的收获，不是"最终能跑通"，而是看清了这条链路每一层在做什么。这些原理一旦内化，下次无论是换个区域、换个模型格式，还是把这套搬到别的受限环境，不再是对着报错碰运气，而是能一眼看出问题。

posted @ 2026-07-09 10:22 zhaojie10 阅读(2) 评论(0) 收藏举报

刷新页面返回顶部

zhaojiew10学习记录