使用 Amazon SageMaker Operator 简化机器学习任务管理

Amazon SageMaker Operator 是一款工具，它可以帮助数据科学家和开发人员利用 Kubernetes 的接口来创建和管理 SageMaker 任务，如模型训练、超参数优化、批量转换及实时推理等。SageMaker Operator 使熟悉 Kubernetes 的开发和运维人员能够通过 kubectl 命令行或 Kubernetes API 接口轻松管理 SageMaker 服务，就像在 Kubernetes 和 AWS SageMaker 服务之间架起了一座桥梁，使得开发者可以快速高效地使用 SageMaker 服务，而无需投入太多精力学习新工具。

应用场景

SageMaker Operator 适用于多种场景：

1. 现有项目引入机器学习：对于已经熟练使用 Kubernetes 平台的项目团队来说，如果需要引入机器学习部分，而不想单独维护一套机器学习平台，那么可以充分利用现有的 Kubernetes 能力，结合第三方工具和服务，迅速实现业务创新，降低成本，提高效率。

2. 统一管理机器学习模块：当项目团队已经在 Kubernetes 平台上实现了 DevOps，并且许多业务模块都实现了持续集成和发布时，机器学习部分作为业务系统的一部分，团队希望能够将其与其他业务模块一起进行工程化管理，最终实现业务流程的全自动化。

技术原理

SageMaker Operator 的工作原理与其他第三方 Kubernetes Operator 类似，由 Kubernetes CRD（自定义资源定义）和 Controller 组成。安装 SageMaker Operator 后，它会在 Kubernetes 平台上注册一些与 SageMaker 相关的 CRD，并以 Pod 形式在 Kubernetes Worker Node 上运行 SageMaker Operator Controller。当开发人员需要创建机器学习任务时，他们需要在 YAML 文件中定义与 CRD 资源相关的配置参数，然后通过 kubectl 提交命令来创建这些资源。SageMaker Controller 会监听这些资源的变化，并调用 SageMaker 服务，完成机器学习任务的创建。

实验演示

接下来，我们将展示如何利用 AWS SageMaker Operator 来创建和管理机器学习任务。在这个示例中，我们将使用以下服务：

• AWS S3：用于存储训练数据集和机器学习模型。
• AWS EKS：用于运行 SageMaker Operator 和其他与 Kubernetes 相关的应用。
• AWS SageMaker：用于自动化创建和管理数据处理、模型训练、模型优化和模型部署等任务。
• AWS IAM：用于控制 SageMaker Operator 和 SageMaker 对 AWS 资源的访问权限。
• AWS CloudWatch：用于记录 SageMaker 机器学习任务的日志，并将其推送到 AWS EKS。

实验内容

我们将搭建 EKS 集群，配置 IAM 权限，并安装 SageMaker Operator。然后，我们将使用 kubectl 命令创建 SageMaker 机器学习训练任务和部署任务。任务完成后，我们将在 SageMaker 界面上查看训练和部署的状态，并确认模型已自动上传至 S3。

实验步骤

1. EKS 集群搭建

EKS 是 AWS 的 Kubernetes 托管集群，可以通过图形界面或命令行创建。在这个实验中，我们将在 ap-southeast-1 区域部署 Kubernetes 集群，该集群将包含一个 Worker Node Group，其中包含三个 c5.xlarge 计算实例。默认情况下，节点将均匀分布在该区域的多个可用区中，如果有定制化需求，可以通过指定命令行参数或编写配置文件来自定义可用区。

1.1 创建集群

bash eksctl create cluster --region <region-name> --name=<cluster-name> --nodes-min=3 --nodes-max=5 --node-type=c5.xlarge --ssh-access --ssh-public-key <ssh-keyname>

1.2 查看 Kubernetes 节点状态

bash kubectl get node

2. 创建 OpenID Connect Provider

OpenID Connect Provider 用于将 Kubernetes 内的 Service Account 与 AWS IAM 角色关联起来。SageMaker Operator Controller 会以 Pod 形式运行在 EKS 节点上，该 Pod 在监听资源创建信息后会调用 SageMaker 服务。成功调用 SageMaker 的前提是该 Pod 需要具有相应的权限。Kubernetes 内 Pod 的权限分配通常是通过 Service Account 传递的，因此 OpenID Connect Provider 的作用是将集群外的 IAM 角色与集群内的 Service Account 关联起来。

2.1 获取 OIDC ID

bash aws eks describe-cluster --name ${CLUSTER_NAME} --region ${AWS_REGION} --query cluster.identity.oidc.issuer --output text

2.2 创建名为 trust.json 的文件

json { "Version": "2012-10-17", "Statement": [ { "Effect": "Allow", "Principal": { "Federated": "arn:aws:iam::<AWS account number>:oidc-provider/oidc.eks.<EKS Cluster region>.amazonaws.com/id/<OIDC ID>" }, "Action": "sts:AssumeRoleWithWebIdentity", "Condition": { "StringEquals": { "oidc.eks.<EKS Cluster region>.amazonaws.com/id/<OIDC ID>:aud": "sts.amazonaws.com", "oidc.eks.<EKS Cluster region>.amazonaws.com/id/<OIDC ID>:sub": "system:serviceaccount:sagemaker-k8s-operator-system:sagemaker-k8s-operator-default" } } } ] }

2.3 创建 IAM role

bash aws iam create-role --role-name <role name> --assume-role-policy-document file://trust.json --output text

2.4 赋予角色 SageMaker FullAccess 权限

bash aws iam attach-role-policy --role-name <role name> --policy-arn arn:aws:iam::aws:policy/AmazonSageMakerFullAccess

3. 安装 SageMaker Operator

SageMaker Operator 可以通过 YAML 文件或 Helm 方式安装。在这个实验中，我们将通过 YAML 文件安装。SageMaker Operator 由 CRD 和 Controller 组成，其中 Controller 会以 Pod 形式运行在 EKS Worker Node 上。

3.1 下载用于安装 SageMaker Operator 的 YAML 文件

bash wget https://raw.githubusercontent.com/aws/amazon-sagemaker-operator-for-k8s/master/release/rolebased/installer.yaml

3.2 更新 YAML 文件中的 eks.amazonaws.com/role-arn 参数项

3.3 安装 SageMaker Operator

bash kubectl apply -f installer.yaml

3.4 查看新创建的 CRD

bash kubectl get crd | grep sagemaker

3.5 查看新创建的 SageMaker Operator Controller

bash kubectl -n sagemaker-k8s-operator-system get pods

4. 安装 SageMaker Operator 日志插件

SageMaker Operator 日志插件是一个可选组件，它允许用户通过 kubectl 命令查看 SageMaker 任务的日志。安装该插件可以方便用户对日志进行管理。

4.1 安装 SageMaker 日志插件

bash export os="linux" wget https://amazon-sagemaker-operator-for-k8s-us-east-1.s3.amazonaws.com/kubectl-smlogs-plugin/v1/${os}.amd64.tar.gz tar xvzf ${os}.amd64.tar.gz mkdir ~/sagemaker-k8s-bin cp ./kubectl-smlogs.${os}.amd64/kubectl-smlogs ~/sagemaker-k8s-bin/ echo 'export PATH=$PATH:~/sagemaker-k8s-bin' >> ~/.bashrc source ~/.bashrc

4.2 验证插件安装是否成功

5. 创建机器学习训练任务

接下来，我们将创建一个 SageMaker 机器学习模型训练任务，以 XGBoost 算法为例，对 MNIST 数据集进行训练和部署。

5.1 生成数据集并上传至 S3 存储桶

bash wget https://raw.githubusercontent.com/aws/amazon-sagemaker-operator-for-k8s/master/scripts/upload_xgboost_mnist_dataset/upload_xgboost_mnist_dataset chmod +x upload_xgboost_mnist_dataset ./upload_xgboost_mnist_dataset --s3-bucket <BUCKET_NAME> --s3-prefix xgboost-mnist

5.2 创建 SageMaker role

bash export assume_role_policy_document='{ "Version": "2012-10-17", "Statement": [{ "Effect": "Allow", "Principal": { "Service": "sagemaker.amazonaws.com" }, "Action": "sts:AssumeRole" }] }' aws iam create-role --role-name <execution role name> --assume-role-policy-document file://<(echo "$assume_role_policy_document") aws iam attach-role-policy --role-name <execution role name> --policy-arn arn:aws:iam::aws:policy/AmazonSageMakerFullAccess

5.3 准备机器学习训练任务 YAML 文件

yaml apiVersion: sagemaker.aws.amazon.com/v1 kind: TrainingJob metadata: name: xgboost-mnist spec: roleArn: <Your SageMaker Role> region: <Your-Region> algorithmSpecification: trainingImage: <Docker Image> trainingInputMode: File outputDataConfig: s3OutputPath: <s3://your-bucket/test> inputDataConfig: - channelName: train dataSource: s3DataSource: s3DataType: S3Prefix s3Uri: <s3://your-bucket/train> s3DataDistributionType: FullyReplicated contentType: text/csv compressionType: None - channelName: validation dataSource: s3DataSource: s3DataType: S3Prefix s3Uri: <s3://your-bucket/validate> s3DataDistributionType: FullyReplicated contentType: text/csv compressionType: None resourceConfig: instanceCount: 1 instanceType: ml.m4.xlarge volumeSizeInGB: 5 hyperParameters: - name: max_depth value: "5" - name: eta value: "0.2" - name: gamma value: "4" - name: min_child_weight value: "6" - name: silent value: "0" - name: objective value: multi:softmax - name: num_class value: "10" - name: num_round value: "10" stoppingCondition: maxRuntimeInSeconds: 86400

5.4 创建 SageMaker 训练任务

bash kubectl apply -f train.yaml

5.5 查看 SageMaker 训练任务状态

bash kubectl describe trainingjob xgboost-mnist kubectl get trainingjob xgboost-mnist

5.6 查看 SageMaker 控制台上的训练任务状态

6. 模型部署

训练完成后，我们将创建模型部署任务，SageMaker Operator 将自动调用 SageMaker 服务完成模型的部署。

6.1 创建模型部署 YAML 文件

yaml apiVersion: sagemaker.aws.amazon.com/v1 kind: HostingDeployment metadata: name: hosting-deployment spec: region: <Your-Region> productionVariants: - variantName: AllTraffic modelName: xgboost-model initialInstanceCount: 1 instanceType: ml.r5.large initialVariantWeight: 1 models: - name: xgboost-model executionRoleArn: <Your SageMaker Role> containers: - containerHostname: xgboost modelDataUrl: <S3://Your-Bucket/Model> image: <Docker-Image>

6.2 查看模型部署任务状态

bash kubectl get hostingdeployment hosting-deployment

6.3 查看 SageMaker 控制台上的模型部署任务状态

总结

通过使用 SageMaker Operator，项目团队可以利用熟悉的 Kubernetes 命令和接口管理复杂的机器学习任务和流程，而无需对现有基础设施和平台进行重大更改。这大大提升了效率并降低了成本。在实际项目中，机器学习部分通常只是业务系统的一部分，因此在设计时需要全面考虑 CI/CD、监控运维、安全控制和服务集成等因素。AWS 提供了高度集成的服务，可以在实验案例基础上进一步扩展，集成新的服务，从而构建一个端到端的解决方案。