aws-samples/sample-sagemaker-mlflow-rest-apis

# Amazon SageMaker AI MLflow CDK 项目 此 CDK 项目演示了如何通过外部服务访问 Amazon SageMaker AI MLflow REST API。它支持两种 MLflow 部署模式 —— **Tracking Server** 和 **Serverless App** —— 并创建一个 Flask Web 应用程序，作为外部安全访问 MLflow API 的代理，遵循 AWS 最佳实践。 ## 架构概述 该项目由四个主要 Stack 组成： 1. **Networking Stack** - 基础 VPC 基础设施（共享） 2. **SageMaker Domain Stack** - 包含用户配置文件的 Amazon SageMaker AI Studio 域（共享） 3. **Managed MLflow Stack** - MLflow Tracking Server 或 Serverless App（特定类型） 4. **Flask App Stack** - 带有 ALB 的 Flask 代理应用程序，用于访问 MLflow API（特定类型） ### MLflow 部署模式 该项目支持两种 MLflow 部署模式，通过 `mlflowType` CDK context 参数进行选择： | 方面 | Tracking Server (`tracking`) | Serverless App (`serverless`) | |--------|------------------------------|-------------------------------| | 资源 | `CfnMlflowTrackingServer` | Custom Resource（基于 Lambda 的 `create_mlflow_app`） | | API Endpoint | `https://{region}.experiments.sagemaker.aws` | `https://mlflow.sagemaker.{region}.app.aws` | | SigV4 服务 | `sagemaker-mlflow` | `sagemaker` | | 请求标头 | `x-mlflow-sm-tracking-server-arn` | `x-sm-mlflow-app-arn` | | IAM 权限 | `sagemaker-mlflow:*` | `sagemaker:CallMlflowAppApi`, `sagemaker-mlflow:*` | 这两种模式可以在同一个账户/区域中共存 —— Networking 和 Amazon SageMaker AI Domain Stack 是共享的，而 MLflow 和 Flask App Stack 会根据类型进行命名（例如：`sagemaker-infra-mlflow-tracking`、`sagemaker-infra-flaskapp-serverless`）。 ## Stack 依赖关系 ``` NetworkingStack (Foundation - shared) ├── SageMakerDomainStack (depends on NetworkingStack - shared) ├── ManagedMlflowStack (depends on NetworkingStack + SageMakerDomainStack - per mlflowType) └── FlaskAppStack (depends on NetworkingStack + ManagedMlflowStack - per mlflowType) ``` ## 功能 ### Networking Stack - 具有公有、私有和隔离子网的多可用区 VPC - 用于优化成本和安全性的 VPC endpoint - 遵循最小权限原则的安全组 - 用于出站互联网访问的 NAT 网关 ### Amazon SageMaker AI Domain Stack - 仅限 VPC 访问的 Amazon SageMaker AI Studio 域 - 具有适当权限的 IAM 执行角色 - 用于存放 Amazon SageMaker AI 工件的 S3 存储桶 - 默认用户配置文件配置 ### Managed MLflow Stack - **Tracking Server 模式**：具有自动模型注册功能的 Amazon SageMaker AI MLflow Tracking Server - **Serverless App 模式**：通过 `create_mlflow_app` API 预置无服务器 MLflow 应用程序的 Lambda Custom Resource，包含自动打包的 boto3 Lambda layer（在 `cdk synth`/`cdk deploy` 期间自动构建 —— 如果可用，则使用本地 `pip`/`pip3`，否则回退到 Docker） - 具有生命周期策略的 MLflow 工件 S3 存储桶 - 完全托管的后端（无需进行数据库管理） ### Flask App Stack - 支持 Tracking Server 和 Serverless MLflow 两种模式的 Flask 代理应用程序 - 根据 ARN 格式自动检测模式（ARN 中包含 `:mlflow-app/` = Serverless） - 使用适合模式的签名服务和标头进行 SigV4 请求签名 - Application Load Balancer（面向互联网，端口 80 -> EC2 端口 5000） - EC2 实例（Ubuntu 24.04 LTS），通过 SSM Parameter Store 自动解析 AMI - 通过 STS 担任的专用 IAM 角色（`FlaskMlflowRole`），仅具有最小的 MLflow 权限 - 遵循最小权限原则的安全组 - 使用用户数据脚本进行自动部署（S3 辅助程序下载、systemd 服务设置） ## 前置条件 要顺利完成本指南，请确保您具备以下前置条件： - 一个 AWS 账户 - 一台安装了以下工具的工作站： - 配置好 AWS CLI，且具有在 AWS 账户中创建 VPC、EC2、Amazon SageMaker AI、S3、IAM、CloudFormation 和 Application Load Balancer 的权限 - Node.js v18+ - NPM - AWS CDK CLI v2.100.0+ - 带有 `pip`/`pip3` 的 Python 3.x（用于在 Serverless 部署期间自动打包 boto3 Lambda layer）。如果不可用，将回退使用 Docker。 ### 辅助脚本 该项目包含将辅助脚本自动部署到 S3 的流程。请确保 `helpers/` 文件夹中包含： - `app/main.py` - Flask 应用程序（带 SigV4 签名的双模式代理） - `app/aws_utils.py` - AWS 工具（STS AssumeRole、带有可配置服务名称的 SigV4Auth） - `app/requirements.txt` - Python 依赖项 (boto3 >= 1.42.0) - `mlflowproxy.service` - Systemd 服务文件（环境变量由 CDK 注入） - `setup_mlflow_proxy_app.sh` - 安装配置脚本 - `install_python13.sh` - Python 3.13 安装脚本 **注意**：在 CDK 部署期间，会使用 `BucketDeployment` 自动将辅助脚本上传到 S3。 ## 安装 1. 克隆仓库并安装依赖项： ``` npm install ``` 2. 引导 CDK（如果尚未执行）： ``` npx cdk bootstrap aws:/// ``` ## 部署 ### 使用 Tracking Server 部署（默认） ``` npx cdk deploy --all ``` 或显式指定： ``` npx cdk deploy --all -c mlflowType=tracking ``` ### 使用 Serverless MLflow App 部署 ``` npx cdk deploy --all -c mlflowType=serverless ``` ### 同时部署两种模式（在同一账户/区域中共存） ``` npx cdk deploy --all -c mlflowType=tracking npx cdk deploy --all -c mlflowType=serverless ``` ### 部署到其他区域 ``` CDK_DEFAULT_REGION=us-west-2 npx cdk bootstrap aws:///us-west-2 CDK_DEFAULT_REGION=us-west-2 npx cdk deploy --all -c mlflowType=serverless ``` ## 配置 ### 环境变量 - `CDK_DEFAULT_ACCOUNT` - AWS 账户 ID - `CDK_DEFAULT_REGION` - AWS 区域（默认为 us-east-1） ### Flask 应用程序组件 该 Flask 应用程序包含： - **main.py** - 具有双模式 MLflow API 路由（Tracking Server / Serverless）的 Flask 代理应用程序 - **aws_utils.py** - AWS 身份验证、STS AssumeRole 和 SigV4 请求签名 - **mlflowproxy.service** - 用于自动启动的 Systemd 服务（环境变量：`SM_MLFLOW_ROLE_ARN`、`SM_MLFLOW_AWS_REGION`、`SM_MLFLOW_RESOURCE_ARN`） - **setup_mlflow_proxy_app.sh`** - 安装和配置脚本 - **install_python13.sh** - Python 3.13 安装脚本 ## 清理 要销毁特定模式的所有资源： ``` npx cdk destroy --all -c mlflowType=tracking npx cdk destroy --all -c mlflowType=serverless ``` **注意**：某些资源（如 S3 存储桶）具有保留策略，可能需要手动删除。Networking 和 Amazon SageMaker AI Domain Stack 是共享的 —— 请仅在移除两个 MLflow/Flask Stack 后再销毁它们。 ## 故障排除 ### 常见问题 1. **VPC Endpoint 创建失败**：确保该区域支持所有必需的 VPC endpoint 2. **Amazon SageMaker AI Domain 创建超时**：检查安全组规则和子网配置 3. **无法访问 MLflow API**：验证 Flask 应用程序是否正在运行，以及 ALB 健康检查是否通过 4. **Flask 应用程序出现 403 错误**：检查 IAM 角色权限 —— Tracking Server 需要 `sagemaker-mlflow:*`，Serverless 模式额外需要 `sagemaker:CallMlflowAppApi` 5. **ALB 健康检查失败**：验证 EC2 实例是否正在运行，以及 Flask 应用程序是否已在端口 5000 上启动 6. **循环依赖错误**：确保安全组规则不会创建循环引用 7. **EC2 用户数据执行失败**：通过 SSM 检查 `/var/log/user-data.log` —— 常见原因是启动时 dpkg 锁被 `unattended-upgrades` 占用 8. **Serverless MLflow App 创建失败**：boto3 Lambda layer 会在综合阶段自动打包。如果本地 `pip`/`pip3` 不可用，则必须运行 Docker 以进行回退打包。请确保 layer 中安装的 boto3 >= 1.42.0（Lambda 自带的 boto3 对于 `create_mlflow_app` API 来说版本过旧） ## 贡献者 该项目的开发与维护由以下人员完成： - **Manish Garg** - **Ashish Bhatt** - **Ram Yennapusa**

标签：Amazon SageMaker, AWS CDK, Flask, MITM代理, MLflow, MLOps, 人工智能, 用户模式Hook绕过, 自动化攻击, 逆向工具