在机器学习训练中实现硬件可恢复性是降低风险并确保训练过程不间断的关键。通过实现预防性健康监测和自动恢复机制,组织可以创建一个能够在遇到硬件故障或其他问题时保持训练过程完整性的容错环境。
在本文中,我们将介绍 AWS Neuron 节点问题探测器及其恢复 DaemonSet,适用于在 Amazon Elastic Kubernetes ServiceEKS上运行的 AWS Trainium 和 AWS Inferentia。该组件通过监控日志快速检测稀有故障事件。当 Neuron 设备出现故障时,它将标记出现问题的工作节点为不健康,并及时用新工作节点替换它们。通过加快故障检测和修复的速度,它提高了机器学习训练的可靠性,并减少了由于硬件故障造成的时间和成本浪费。
该解决方案基于节点问题检测器和恢复 DaemonSet,这是一个用于自动检测和报告 Kubernetes 集群中各种节点级别问题的强大工具。
节点问题检测器组件将持续监控工作节点上的内核消息kmsg日志。如果检测到与 Neuron 设备相关的错误信息,它将把 Kubernetes API 服务器上的 NodeCondition 更改为 NeuronHasError。
节点恢复代理是一个独立的组件,它会定期检查节点问题探测器暴露的 Prometheus 指标。当它发现与 Neuron 设备相关的问题时,它会采取自动化措施。首先,它会在相关的自动扩缩组中将受影响的实例标记为不健康,这将触发该扩缩组停止该实例并启动一个替代实例。此外,节点恢复代理还会发布 Amazon CloudWatch 指标,以便用户监控和警报这些事件。
以下图展示了解决方案的架构和工作流:
在接下来的实践中,我们将创建一个 EKS 集群并部署 Neuron 插件以检测节点问题,在节点中注入错误消息。我们将观察失败的节点被停止并被新节点替换,并在 CloudWatch 中找到指示错误的指标。
在开始之前,请确保您在本地安装了以下工具:
最新版本的 AWS Command Line InterfaceAWS CLIeksctlkubectlTerraformSession Manager 插件请按照以下步骤配置节点问题检测和恢复插件:
使用 EKS Terraform 模块创建 EKS 集群:
bashgit clone https//githubcom/awslabs/dataoneksgit
export TFVARregion=useast2export TFVARtrn132xldesiredsize=4export TFVARtrn132xlminsize=4cd dataoneks/aiml/trainiuminferentia/ ampamp chmod x installsh/installsh
aws eks region useast2 describecluster name trainiuminferentia
aws eks region useast2 updatekubeconfig name trainiuminferentia
kubectl get nodes
安装节点问题检测器插件和服务账户的 AWS 身份与访问管理IAM角色。
创建相应的 IAM 策略。如下所示,更新 Resource 键值以匹配您的包含 Trainium 和 AWS Inferentia 节点的节点组 ARN,并更新 ec2ResourceTag/awsautoscalinggroupName 键值以匹配自动扩缩组名称。您可以从 Amazon EKS 控制台获取这些值。选择导航窗格中的“集群”,打开 trainiuminferentia 集群,选择节点组,找到您的节点组。

json
{ Version 20121017 Statement [ { Action [ autoscalingSetInstanceHealth autoscalingDescribeAutoScalingInstances ] Effect Allow Resource lt适用于 Neuron 节点的 Auto Scaling 组 ARNgt } { Action [ ec2DescribeInstances ] Effect Allow Resource Condition { ForAllValuesStringEquals { ec2ResourceTag/awsautoscalinggroupName lt对应 Neuron 节点的 Auto Scaling 组名称gt } } } { Action [ cloudwatchPutMetricData ] Effect Allow Resource Condition { StringEquals { cloudwatchNamespace NeuronHealthCheck } } } ]}
该组件将作为 DaemonSet 安装到您的 EKS 集群中。
bashaws iam createpolicy policyname NeuronProblemDetectorPolicy policydocument file//npdpolicytrimmedjson
请注意 ARN。
bashCLUSTERNAME=trainiuminferentia # 您的 EKS 集群名称 AWSREGION=useast2ACCOUNTID=(aws sts getcalleridentity query Account output text)POLICYARN=arnawsiamACCOUNTIDpolicy/NeuronProblemDetectorPolicy
eksctl create addon cluster CLUSTERNAME name ekspodidentityagent region AWSREGION
eksctl create podidentityassociation cluster CLUSTERNAME namespace neuronhealthchecksystem serviceaccountname nodeproblemdetector permissionpolicyarns=POLICYARN region AWSREGION
安装 Neuron NPD 和恢复插件bashkubectl create ns neuronhealthchecksystemcurl https//rawgithubusercontentcom/awsneuron/awsneuronsdk/215b421ac448d85f89be056e27e29842a6b03c9c/src/k8/neuronproblemdetector/k8sneuronproblemdetectorandrecoveryyml kubectl apply f curl https//rawgithubusercontentcom/awsneuron/awsneuronsdk/215b421ac448d85f89be056e27e29842a6b03c9c/src/k8/neuronproblemdetector/k8sneuronproblemdetectorandrecoveryrbacyml kubectl apply f curl https//rawgithubusercontentcom/awsneuron/awsneuronsdk/215b421ac448d85f89be056e27e29842a6b03c9c/src/k8/neuronproblemdetector/k8sneuronproblemdetectorandrecoveryconfigyml kubectl apply f
预计结果包含 4 个 Neuron 节点的集群:
bashkubectl get pod n neuronhealthchecksystem
如果 succesvol,则所有 pod 应该处于 Running状态。默认情况下,节点问题检测器不会对失败的节点采取任何操作。如果您希望代理自动终止 EC2 实例,请更新 DaemonSet,如下所示:
tgfish飞鱼加速bashkubectl edit n neuronhealthchecksystem ds/nodeproblemdetector
env name ENABLERECOVERY value true
安装插件后,您可以通过运行 kubectl describe node 来查看 Neuron 条件。我们通过在实例中注入错误日志来模拟设备错误:
bash
kubectl describe node ip1006458151useast2computeinternal grep Conditions A7
如上所示,您可以看到节点条件的变化。
要获取提供者 ID:
bashkubectl describe node ip1006458151useast2computeinternal grep i provider sed E s//([/])/1/
通过 SSH 进入工作节点并模拟 Neuron 设备上的硬件错误:
bashaws ssm startsession target i0381404aa69eae3f6 region useast2
bashsh42sh42 sudo bash[root@ip19216893211 bin]# echo test NEURONHWERR=DMAERROR test gtgt /dev/kmsg
大约 2 分钟后,您会看到错误被识别:
bashkubectl describe node ip1006458151useast2computeinternal grep Conditions A7
现在,节点问题检测器已经检测到错误,恢复代理自动采取措施将此节点设置为不健康状态,从而使 Amazon EKS 将该节点指定为不可调度,驱逐在该节点上的 Pods:
bash
kubectl get node
您可以打开 CloudWatch 控制台并验证 NeuronHealthCheck 的指标。您可以看到 CloudWatch 的 NeuronHasErrorDMAERROR 指标值为 1。
替换后,您会看到一个新的工作节点被创建:
bash
kubectl get node
让我们看一个实际场景,您正在运行一个分布式训练任务,使用 MPI 运算符,如 Llama2 on Trainium,并且其中一个节点出现无法恢复的 Neuron 错误。在未部署该插件之前,训练任务陷入僵局,浪费了时间和计算成本。而部署该插件之后,节点问题检测器会主动将有问题的节点从集群中移除。在训练脚本中,它会定期保存检查点,以确保训练从上一个检查点恢复。
以下截图显示了分布式训练任务的示例日志:
以下截图显示在第 77 步创建的检查点:
在第 86 步时,训练在其中一个节点出现问题后停止。错误是人为注入的测试。
在节点问题检测并通过 Neuron 插件进行替换后,训练过程从第 77 步恢复,即上一个检查点。
虽然 Auto Scaling 组会停止不健康节点,但它们可能会遇到问题,从而阻止替换节点的启动。在这种情况下,训练任务会停滞,需要手动干预。然而,被停止的节点将不会对相关的 EC2 实例产生额外费用。
如果您希望在停止实例之外采取自定义操作,可以创建 CloudWatch 警报,监视指标 NeuronHasErrorDMAERROR、NeuronHasErrorHANGONCOLLECTIVES、NeuronHasErrorHBMUNCORRECTABLEERROR、NeuronHasErrorSRAMUNCORRECTABLEERROR 和 NeuronHasErrorNCUNCORRECTABLEERROR,并使用 CloudWatch Metrics Insights 查询,如 SELECT AVG(NeuronHasErrorDMAERROR) FROM NeuronHealthCheck 来聚合这些值以评估警报。
要清理本文中所有配置的资源,请运行清理脚本:
bash
eksctl delete podidentityassociation serviceaccountname nodeproblemdetector namespace neuronhealthchecksystem cluster CLUSTERNAME region AWSREGION
cd dataoneks/aiml/trainiuminferentia/cleanupsh
在本文中,我们展示了 Neuron 问题检测器和恢复 DaemonSet 是如何在使用 Trainium 和 AWS Inferentia 的 EC2 实例上工作的。如果您正在运行基于 Neuron 的 EC2 实例并使用托管节点或自我管理的节点组,您可以在 EKS 集群中部署检测器和恢复 DaemonSet,从而在节点发生故障时提升机器学习训练工作负载的可靠性和容错性。
Harish Rao 是 AWS 的高级解决方案架构师,专注于大规模分布式 AI 训练和推理,帮助客户利用 AI 的力量推动创新,解决复杂挑战。
Ziwen Ning 是 AWS 的软件开发工程师,专注于通过将 AWS Neuron 与容器化环境和 Kubernetes 集成来提升 AI/ML 体验。
Geeta Gharpure 是 Annapurna ML 工程团队的高级软件开发人员,专注于在 Kubernetes 上运行大规模 AI/ML 工作负载。
Darren Lin 是 AWS 的云原生专业解决方案架构师,专注于 Linux、Kubernetes、容器、可观察性和开源技术等领域。