aws-samples/sample-nova-multi-turn-rl-infra

# Nova 多轮强化学习 CDK 基础设施 用于在 SageMaker HyperPod 上使用多轮强化学习进行 Amazon Nova 模型定制的 AWS 基础设施。 ## 架构 两阶段部署： 1. `cdk deploy` 配置基础基础设施（VPC、IAM、S3、EKS、HyperPod、ECS、Step Functions、Lambda、EventBridge）。 2. 将 `.jsonl` 文件上传到 S3 的 `training-data/` 前缀会通过 EventBridge 触发 Step Functions 工作流。 ``` EventBridge (S3 upload) │ ▼ S3Trigger Lambda ──▶ Step Functions pipeline │ ▼ BootstrapEcrImage (CodeBuild) → DeployRftInfra (Lambda) → StartRewardWorkers (Lambda) → ValidateUploadData (Lambda) → SubmitTraining (CodeBuild, hyperpod CLI) → Success (on failure → PipelineFailed → SNS alert) ``` 步骤 0 和 4 在 CodeBuild 中运行（需要 Docker / hyperpod CLI）。步骤 1–3 在打包为 Docker 镜像的 Lambda 中运行，并使用 Nova SDK。每一步都有幂等性检查、指数退避重试机制，以及统一路由到失败状态的兜底策略。 由 SDK 管理的资源（conversation-proxy Lambda、SQS FIFO 队列、DynamoDB 状态表、ECS Fargate reward worker）是在工作流运行时创建的，而不是由 CDK 创建的。 ## 前置条件 - Python 3.12+，AWS CDK v2 (`npm install -g aws-cdk`)，AWS CLI v2 - Docker 或 [Finch](https://github.com/runfinch/finch) 用于构建 Lambda 镜像 - 拥有 `ml.p5.48xlarge` 配额的 AWS 账户 ## 快速开始 ``` pip install -r requirements.txt # 要求：编辑 cdk.json 并设置 project_tag + sdk_resource_prefix cdk deploy --require-approval never # 上传会自动触发 pipeline aws s3 cp training-data.jsonl s3://BUCKET_NAME/training-data/training-data.jsonl ``` 在部署时覆盖参数： ``` cdk deploy -c instance_count=1 -c max_steps=20 -c global_batch_size=128 ``` ## 配置 所有参数都位于 `cdk.json` 的 `context` 下。 **基础设施：** `project_tag` 和 `sdk_resource_prefix` 是必填前缀。关键默认值：`instance_type=ml.p5.48xlarge`，`instance_count=10`（`generation_replicas=4` 时至少为 10），`nova_model=NOVA_LITE_2`，`region=us-east-1`，`eks_kubernetes_version=1.32`，`vf_env_id=wordle`。 **训练覆盖参数**（贯穿整个工作流事件）：`training_method`、`max_steps`、`generation_replicas`、`global_batch_size`、`max_new_tokens`、`max_length`、`training_timeout`。 **目标现有集群：** 设置 `target_cluster_name`、`target_eks_cluster_name`、`target_instance_type`、`target_instance_count` 以针对预先存在的 HyperPod 集群运行。 ### 训练数据格式 基于元数据的 JSONL，每行一条记录；`data_s3_path` 必须是一个 `.jsonl` 文件，而不是目录： ``` {"id": "wordle_train_001", "metadata": {"prompt": "Guess the 5-letter word", "answer": "crane"}} ``` ### 自定义 reward 环境 1. 创建包含 `load_environment()` 函数的 `custom-environments//` 目录（参见 `my-custom-env`）。 2. 在 `cdk.json` 中设置 `use_custom_env: "true"` 和 `custom_env_id: ""`。 3. `cdk deploy` 会将该环境打包并上传至 S3。 `client/nova_async_client.py` 提供了一个兼容 OpenAI 的异步客户端，该客户端使用 SigV4 认证，适用于需要调用 Nova 推理端点的自定义 reward 函数。 ## 监控 **工作流失败告警** — 订阅导出为 `AlertTopicArn` 的 SNS topic： ``` aws sns subscribe --topic-arn \ --protocol email --notification-endpoint you@example.com ``` **EventBridge 触发失败** — 检查导出为 `TriggerDlqUrl` 的 DLQ。 **CloudWatch Logs：** - `/aws/vendedlogs/states/nova-rl-*` — Step Functions 执行日志 - `/aws/lambda/*{DeployRftInfra,StartRewardWorkers,ValidateUploadData}*` — 步骤 1–3 - `/aws/codebuild/nova-rl-submit-training` — 步骤 4 - `/aws/sagemaker/Clusters//*` — HyperPod 容器日志 ## 清理 1. 停止正在运行的 Step Functions 执行和 ECS reward worker。 2. 删除 SDK CloudFormation 堆栈 (`-*-NovaForgeSDK`)。 3. `cdk destroy`。 由于 GuardDuty 管理的 ENI，VPC 删除可能会失败——请先删除 VPC endpoint，等待约 2 分钟，然后重试。 ## 安全 - 启用了 `cdk-nag` `AwsSolutionsChecks`，并附有已记录的豁免项 - S3 存储桶：KMS 加密、阻止公共访问、强制 SSL、已启用版本控制 - 使用 KMS CMK 并为 CloudWatch Logs、SNS、EKS secrets、CodeBuild artifacts 开启轮换 - SNS 和 SQS DLQ 强制使用 SSL - IAM 权限策略在 SDK 命名允许的范围内限定于 `sdk_resource_prefix` - Lambda 容器以非 root 用户身份运行 - 用于 ECR、EKS、STS、CloudWatch Logs、Lambda、SQS 的 VPC 接口 endpoint **EKS 公共端点：** HyperPod 需要一个 PUBLIC 端点，因为其 `/readyz` 健康检查是由客户 VPC 外部的 AWS 托管服务账号运行的。`PUBLIC_AND_PRIVATE` 和仅私有配置会导致集群创建失败。上游 CIDR 范围未公开，因此无法限制 `public_access_cidrs`。访问受到 IAM 身份验证的保护。 ## 项目结构 ``` app.py CDK app entry point cdk.json CDK config and context defaults requirements.txt training-data.jsonl Sample data (300 Wordle samples) stacks/ nova_rl_stack.py VPC, EKS, HyperPod, S3, ECS, IAM orchestration.py Step Functions, Lambdas, EventBridge iam_policies.py Shared IAM policy builders lambdas/ Dockerfile Shared image for SDK Lambdas shared/infra_utils.py RFT infrastructure kwargs builder deploy_rft_infra/ Step 1 start_reward_workers/ Step 2 validate_upload_data/ Step 3 submit_training/ Step 4 (CodeBuild) s3_trigger/ EventBridge → Step Functions helm_install_trigger/ Custom Resource for Helm install lifecycle-scripts/on_create.sh HyperPod node bootstrap client/nova_async_client.py Async client with SigV4 (BYOO) custom-environments/ Custom reward environments ``` ## 故障排除 | 问题 | 解决方案 | |-------|------------| | `cdk deploy` 因配额失败 | 通过 Service Quotas 申请 `ml.p5.48xlarge` | | 上传未触发工作流 | 验证键名是否匹配 `training-data/*.jsonl`；检查 `TriggerDlqUrl` | | `data_s3_path` 验证错误 | 路径必须以 `.jsonl` 结尾（不能是目录） | | `SubmitTraining` 失败 | 检查 CodeBuild 日志；验证 `hyperpod connect-cluster` 是否成功 | | 训练 pod 发生 crash-loop | 查看 CloudWatch 中的 `/aws/sagemaker/Clusters//*` | | Reward worker 空闲 | 检查 ECS 任务日志；确认 SQS 队列深度 | | VPC 删除失败 | 删除 VPC endpoint，等待 2 分钟以释放 GuardDuty ENI | | HyperPod 报错 `NotStabilized` | EKS 端点必须为 PUBLIC（不能是 PUBLIC_AND_PRIVATE） |

标签：AWS CDK, SageMaker HyperPod, 大模型微调, 强化学习, 机器学习工程, 自动化流水线, 请求拦截, 逆向工具