使用TaskD替换Elastic Agent

前提条件

• 已完成升级环境检查。
• 训练镜像已安装Elastic Agent。

操作步骤

1. 参考获取软件包章节，下载新版本TaskD组件安装包。
2. 下载完成后，进入安装包所在路径并解压安装包。
3. 执行ls -l命令，回显示例如下。

   -rw-r--r-- 1 root root 6134726 Nov 10 10:32 Ascend-mindxdl-taskd_{version}_linux-aarch64.zip
   -r-------- 1 root root 6205642 Nov  5 23:38 taskd-{version}-py3-none-linux_aarch64.whl

4. 基于已有的训练镜像，卸载Elastic Agent并安装新版本TaskD。
   1. 运行训练镜像。示例如下：

      docker run -it  -v /host/packagepath:/container/packagepath training_image:latest /bin/bash

   2. 卸载已安装的Elastic Agent组件。

      pip uninstall mindx-elastic -y

      回显示例如下，表示卸载成功。

      Successfully uninstalled mindx_elastic-{version}

   3. 删除Elastic Agent使能代码。

      sed -i '/mindx_elastic.api/d' $(pip3 show torch | grep Location | awk -F ' ' '{print $2}')/torch/distributed/run.py

      （可选）执行以下命令，查看对应文件是否已经删除Elastic Agent嵌入代码。

      vi $(pip3 show torch | grep Location | awk -F ' ' '{print $2}')/torch/distributed/run.py

   4. 安装新版本TaskD。

      pip install taskd-{version}-py3-none-linux_aarch64.whl

      回显示例如下，表示安装成功。

      Successfully installed taskd-{version}

      执行以下命令，使能TaskD。

      sed -i '/import os/i import taskd.python.adaptor.patch' $(pip3 show torch | grep Location | awk -F ' ' '{print $2}')/torch/distributed/run.py

   5. 安装了新版本TaskD后，将容器保存为新镜像。

      docker ps

      回显示例如下。

      CONTAINER ID   IMAGE                  COMMAND                  CREATED        STATUS        PORTS     NAMES 
      bb118ca00041    f76142d63d3a           "/bin/bash -c 'sleep…"   2 hours ago    Up 2 hours              k8s_ascend_default-last-test-deepseek2-60b

      将该容器提交为新版本训练容器镜像，注意新镜像的tag与旧镜像不一致。示例如下：

      docker commit bb118ca00041 newimage:latest

5. 检查TaskD是否替换完成。参考检查TaskD章节，检查组件状态是否正常。
6. 修改训练脚本（例如train_start.sh）和任务YAML。
   1. 创建manager.py文件，放在调用训练脚本时的当前目录下，manager.py文件内容如下所示。

      from taskd.api import init_taskd_manager, start_taskd_manager
      import os 
       
      job_id=os.getenv("MINDX_TASK_ID") 
      node_nums=XX         # 用户填入任务节点总数
      proc_per_node=XX     # 用户填入任务每个节点的训练进程数量
        
      init_taskd_manager({"job_id":job_id, "node_nums": node_nums, "proc_per_node": proc_per_node}) 
      start_taskd_manager()

      

      manager.py文件中的参数详细说明请参见def init_taskd_manager(config:dict) -> bool:。

   2. 在训练脚本中增加以下代码拉起TaskD Manager。

      export TASKD_PROCESS_ENABLE="on" 
      # 以PyTorch框架为例
      if [[ "${RANK}" == 0 ]]; then
          export MASTER_ADDR=${POD_IP} 
          python manager.py &           # 具体执行路径由当前路径决定
      fi 
            
      torchrun ...

   3. 在任务YAML中修改容器端口，在所有的Pod下增加TaskD通信使用的端口9601。

      ...
      ports:                          
         - containerPort: 9601              
           name: taskd-port
      ...