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

📅 2026/7/10 19:09:05
在 AWS 中国区把 MLflow 模型端到端部署在 SageMaker
机器学习项目里有一个反复出现的断层模型在 notebook 里跑得好好的一旦要上线提供推理服务事情就复杂起来了。实验怎么管理、模型怎么版本化、训练好的工件放哪、推理服务又该怎么部署——每一环都有自己的工具和坑。MLflow SageMaker 是解决这一断层的一套经典组合MLflow 管实验与工件SageMaker 提供托管的推理端点。以下是通用的流程图示本文的目的是把这条链路从训练到推理完整跑一遍用到的项目是一个房价预测模型Kaggle House Prices 数据集用 ElasticNet 训练最终部署成一个 SageMaker 实时推理端点。源码文件地址https://gitee.com/zhaojiew-realm/files/blob/master/housing-price.7z下图是测试环境涉及到的整体架构三个环节各司其职AWS 中国区 (cn-north-1)本地机器 (uv venv, Python 3.10)记录元数据产出工件上传读取模型拉取镜像调用data.py预处理train.py训练 记录deploy.py部署端点test.py推理测试MLflow Server127.0.0.1:5002(sqlite 元数据)本地 mlartifactsS3模型工件ECR推理镜像SageMaker Endpoint(推理)通常来说模型的训练需要在 SageMaker Notebook 上测试方便期间我们在本地进行训练。此外MLflow 跟踪服务器也起在本地。只有真正需要 AWS 托管能力的两样东西留在云上——S3 存模型工件、SageMaker 端点做推理外加 ECR 存推理镜像。初始化环境MLflow 的mlflow.projects.run()支持conda/virtualenv/local三种 env_manager。我们用 uv 建好环境并装齐依赖然后让 MLflow 用env_managerlocal直接复用当前环境而不是让 MLflow 自己去建环境。uv venv--python3.10.venvsource.venv/bin/activate# mlflow 3.14.0、scikit-learn 1.7.2、xgboost 3.2.0。uv pipinstall--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_pointTraining,experiment_nameElasticNet,env_managerlocal)训练与实验跟踪数据预处理data.py对数值列做 KNN 插值、类别列用众数填充再用OneHotEncoder(handle_unknownignore)编码。importpandasaspdfromsklearn.imputeimportKNNImputerfromsklearn.model_selectionimporttrain_test_splitfromsklearn.preprocessingimportOneHotEncoder trainpd.read_csv(data/train.csv)testpd.read_csv(data/test.csv)Xtrain.drop(SalePrice,axis1)ytrain[SalePrice]X_train,X_val,y_train,y_valtrain_test_split(X,y,test_size0.2,random_state42)numeric_colsX_train.select_dtypes(include[int64,float64]).columns non_numeric_colsX_train.select_dtypes(exclude[int64,float64]).columns imputerKNNImputer()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])forcolumninnon_numeric_cols:X_train[column].fillna(X_train[column].mode()[0],inplaceTrue)X_val[column].fillna(X_val[column].mode()[0],inplaceTrue)test[column].fillna(test[column].mode()[0],inplaceTrue)oheOneHotEncoder(dropfirst,handle_unknownignore)X_trainohe.fit_transform(X_train)X_valohe.transform(X_val)testohe.transform(test)训练脚本train.py遍历 ElasticNet 的 18 组超参数组合alpha × l1_ratio × fit_intercept 3×3×2每组开一个 MLflow run记录参数、RMSE/MAPE/R2 指标和模型工件importmlflowfromdataimportX_train,X_val,y_train,y_valfromsklearn.linear_modelimportElasticNetfromsklearn.model_selectionimportParameterGridfromparamsimportelasticnet_param_gridfromutilsimporteval_metricsforparamsinParameterGrid(elasticnet_param_grid):withmlflow.start_run():lrElasticNet(**params)lr.fit(X_train,y_train)y_predlr.predict(X_val)metricseval_metrics(y_val,y_pred)mlflow.log_params(params)mlflow.log_metrics(metrics)mlflow.sklearn.log_model(lr,ElasticNet,input_exampleX_train,code_paths[train.py,data.py,params.py,utils.py],)其中eval_metricsutils.py把三个回归指标打包成字典importnumpyasnpfromsklearn.metricsimportmean_squared_error,r2_score,mean_absolute_percentage_errordefeval_metrics(y_true,y_pred):rmsenp.sqrt(mean_squared_error(y_true,y_pred))mapemean_absolute_percentage_error(y_true,y_pred)r2r2_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: alpha0.1 l1_ratio0.8 fit_interceptTrue在mlflow ui界面查看具体的训练运行以及具体训练的指标可以得知本次训练的source代码使用的数据集和超参数等信息做到心里有数。MLflow 3.x 的模型存储结构要部署最佳模型得先找到它的工件在哪。MLflow 3.x 不再按 run_id 组织模型而是按 model_id。工件在mlartifacts/1/models/m-a1b42d5a6ab940b9b6d593b5fbd8486f/artifacts/想拿到 run 对应的 model_id得用search_logged_modelsAPIlmsclient.search_logged_models(experiment_ids[1],filter_stringfsource_run_id{run_id})# → model_id: m-a1b42d5a6ab940b9b6d593b5fbd8486f看一眼这个模型的MLmodel文件有几个信息后面部署要用到它是skops 序列化的 sklearn 模型不是 pickle签名是tensor [-1, 7940]输入、[-1]输出。构建推理镜像先把最佳模型上传到 S3SageMaker 部署要求模型可从 S3 访问。aws s3api create-bucket--bucketmlflow-project-artifacts-accountID\--regioncn-north-1\--create-bucket-configurationLocationConstraintcn-north-1 aws s3cp\mlartifacts/1/models/m-a1b42d5a6ab940b9b6d593b5fbd8486f/artifacts/\s3://mlflow-project-artifacts-accountID/1/cf6908fac6404d099eab94260c24dcc6/artifacts/ElasticNet/\--recursive--regioncn-north-1在构建之前首先改一些配置来加速。mlflow/models/docker_utils.py的核心是一个模板字符串_DOCKERFILE_TEMPLATE。需要在FROM之后加两行 ENV让镜像内所有 pip 都走阿里源FROM {base_image} # 中国区让镜像内所有 pip install构建时运行时走阿里源 ENV PIP_INDEX_URLhttps://mirrors.aliyun.com/pypi/simple/ ENV PIP_TRUSTED_HOSTmirrors.aliyun.com之后构建卡在git clone阶段因为MLflow 默认的 base image 是ubuntu:22.04用virtualenvenv_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:*-slimgenerate_dockerfile会走一条不装 pyenv的路径。因为此时Python 已经内置只apt install nginx即可。于是改mlflow/sagemaker/cli.py把硬编码的UBUNTU_BASE_IMAGE换成python:3.10-slimbase_imagemlflow.models.docker_utils.PYTHON_SLIM_BASE_IMAGE.format(version3.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}registryf{account}.dkr.ecr.{region}.amazonaws.com改 MLflow 源码在 push 里做 region 判断可行但要动好几处。所以这里直接手动将构建好的镜像推送到 ecrECRaccountID.dkr.ecr.cn-north-1.amazonaws.com.cnaws ecr get-login-password--regioncn-north-1\|dockerlogin--usernameAWS --password-stdin$ECRdockertag mlflow-pyfunc:latest$ECR/mlflow-pyfunc:3.14.0dockerpush$ECR/mlflow-pyfunc:3.14.0部署端点在部署阶段运行deploy.py因为 boto3 client 带了region_namecn-north-1会自动解析到中国区端点日志里全是.amazonaws.com.cn和arn:aws-cn:...。端点资源创建成功状态进入Creating。frommlflow.deploymentsimportget_deploy_client endpoint_nameprod-endpointmodel_uris3://mlflow-project-artifacts-accountID/1/cf6908fac6404d099eab94260c24dcc6/artifacts/ElasticNetconfig{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,}clientget_deploy_client(sagemaker)client.create_deployment(nameendpoint_name,model_urimodel_uri,flavorpython_function,configconfig,)创建的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 容器在启动时会用virtualenvenv_manager 为模型创建一个独立环境这需要 pyenv。slim 镜像里没有容器起不来反复崩溃重试端点卡在Creating出不来。那么为什么容器要在运行时建环境检查mlflow/models/container/__init__.py的_serve_pyfunc关键逻辑是这样的def_serve_pyfunc(model,env_manager):disable_env_creationMLFLOW_DISABLE_ENV_CREATION.get()ifpyfunc.ENVinconf:# 只有 disable_env_creation 为 False 时才在容器启动时建环境装依赖ifnotdisable_env_creation:_install_pyfunc_deps(MODEL_PATH,install_mlflowTrue,env_managerenv_manager)...# 否则直接用镜像里的系统 Python 跑 scoring server而mlflow/sagemaker/cli.py构建镜像时传的是disable_env_creation_at_runtimeFalse——也就是故意让容器在运行时才装模型依赖。这在标准场景是合理的模型运行时才从 S3 加载依赖也运行时才知道但它依赖 pyenv 来建 virtualenv。因此需要设置disable_env_creation_at_runtimeTrueenv_managerlocal并在构建时把模型依赖预装进镜像。这样_serve_pyfunc会跳过_install_pyfunc_deps直接用镜像里的系统 Python 跑推理服务彻底不需要 pyenv。模型的requirements.txt精确列出了它需要什么mlflow3.14.0 numpy2.2.6 pandas2.3.3 pyarrow24.0.0 scikit-learn1.7.2 scipy1.15.3 skops0.14.0于是改cli.py两处把运行时的最小依赖安装步骤换成构建时安装这套精确依赖同时把disable_env_creation_at_runtime改成Truesetup_container(# 把模型依赖焊进镜像运行时禁用环境创建\nRUN pip install --no-cache-dir mlflow3.14.0 scikit-learn1.7.2 skops0.14.0 numpy2.2.6 pandas2.3.3 scipy1.15.3 pyarrow24.0.0 gunicorn[gevent])# generate_dockerfile(..., disable_env_creation_at_runtimeTrue, ...)用--env-manager local重新构建、手动推送到 ECR。删掉那个卡死的失败端点用新镜像重新部署。这次2 分半钟就InServiceThe deployment operation completed successfully with message: The SageMaker endpoint was created successfully.推理测试MLflow pyfunc 端点接受的是{inputs: ...}。importjsonimportboto3fromdataimporttest smrtboto3.client(sagemaker-runtime,region_namecn-north-1)test_data_jsonjson.dumps({inputs:test[:20].toarray().tolist()})predictionsmrt.invoke_endpoint(EndpointNameprod-endpoint,Bodytest_data_json,ContentTypeapplication/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那么简单。尤其当环境偏离教程的隐含假设时那些藏在框架里的假设会一个接一个地暴露出来。这次实战真正的收获不是最终能跑通而是看清了这条链路每一层在做什么。这些原理一旦内化下次无论是换个区域、换个模型格式还是把这套搬到别的受限环境不再是对着报错碰运气而是能一眼看出问题。