Line data Source code
1 : /**
2 : * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3 : * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4 : * CANN Open Software License Agreement Version 2.0 (the "License").
5 : * Please refer to the License for details. You may not use this file except in compliance with the License.
6 : * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7 : * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8 : * See LICENSE in the root of the software repository for the full text of the License.
9 : */
10 : #include "task_abort_handler.h"
11 : #include <algorithm>
12 : #include "log.h"
13 : #include "coll_comm.h"
14 : #include "ccu_device_pub.h"
15 :
16 : namespace hccl {
17 : using HcclUs = std::chrono::steady_clock::time_point;
18 : static std::mutex vecMutex;
19 :
20 4 : int32_t ProcessTaskAbortPre(const std::vector<CollComm *> &commVector, const std::chrono::seconds &localtimeout)
21 : {
22 4 : HcclResult ret = HCCL_SUCCESS;
23 4 : bool isUseTimeOut = localtimeout != std::chrono::seconds(0);
24 4 : std::chrono::seconds elapsed{};
25 6 : for (auto& comm : commVector) {
26 4 : if (isUseTimeOut) {
27 2 : std::chrono::steady_clock::time_point startTime = std::chrono::steady_clock::now();
28 2 : ret = comm->Suspend();
29 2 : elapsed = std::chrono::duration_cast<std::chrono::seconds>(std::chrono::steady_clock::now() - startTime);
30 : } else {
31 2 : ret = comm->Suspend();
32 : }
33 4 : if (ret != HCCL_SUCCESS && ret != HCCL_E_SUSPENDING) {
34 2 : HCCL_ERROR("[NsRecovery] finish suspend failed, ret = 0x%016llx", HCCL_ERROR_CODE(ret));
35 2 : return static_cast<int>(TaskAbortResult::TASK_ABORT_FAIL);
36 : }
37 2 : HCCL_INFO("[NsRecovery]finish suspend success");
38 2 : if (isUseTimeOut) {
39 1 : CHK_PRT_RET(elapsed > localtimeout, HCCL_ERROR("[NsRecovery][suspend] NsRecovery suspend timeOut"),
40 : static_cast<int>(TaskAbortResult::TASK_ABORT_TIMEOUT));
41 : }
42 : }
43 2 : return static_cast<int>(TaskAbortResult::TASK_ABORT_SUCCESS);
44 : }
45 :
46 6 : int32_t ProcessTaskAbortPost(const std::vector<CollComm *> &commVector, int32_t deviceLogicId,
47 : const std::chrono::seconds &localtimeout)
48 : {
49 6 : HcclResult ret = HCCL_SUCCESS;
50 6 : bool isUseTimeOut = localtimeout != std::chrono::seconds(0);
51 6 : std::chrono::seconds elapsed{};
52 6 : if (hcomm::CcuIsInited(deviceLogicId)) {
53 4 : CHK_RET(hcomm::CcuSetTaskKill(deviceLogicId));
54 : } else {
55 2 : HCCL_INFO("[NsRecovery][Callback] CCU not inited, skip CcuSetTaskKill, deviceLogicId[%d]", deviceLogicId);
56 : }
57 10 : for (auto& comm : commVector) {
58 6 : if (isUseTimeOut) {
59 3 : std::chrono::steady_clock::time_point startTime = std::chrono::steady_clock::now();
60 3 : ret = comm->Clean();
61 3 : elapsed = std::chrono::duration_cast<std::chrono::seconds>(std::chrono::steady_clock::now() - startTime);
62 : } else {
63 3 : ret = comm->Clean();
64 : }
65 6 : if (ret != HCCL_SUCCESS && ret != HCCL_E_SUSPENDING) {
66 2 : HCCL_ERROR("[NsRecovery][Callback] finish clean failed, ret = 0x%016llx", HCCL_ERROR_CODE(ret));
67 2 : return static_cast<int>(TaskAbortResult::TASK_ABORT_FAIL);
68 : }
69 4 : HCCL_INFO("[NsRecovery][Callback] finish clean success");
70 4 : if (isUseTimeOut) {
71 2 : CHK_PRT_RET(elapsed > localtimeout, HCCL_ERROR("[NsRecovery][Callback] NsRecovery Clean timeout"),
72 : static_cast<int>(TaskAbortResult::TASK_ABORT_TIMEOUT));
73 : }
74 : }
75 4 : if (hcomm::CcuIsInited(deviceLogicId)) {
76 2 : CHK_RET(hcomm::CcuSetTaskKillDone(deviceLogicId));
77 : } else {
78 2 : HCCL_INFO("[NsRecovery][Callback] CCU not inited, skip CcuSetTaskKillDone, deviceLogicId[%d]", deviceLogicId);
79 : }
80 4 : return static_cast<int>(TaskAbortResult::TASK_ABORT_SUCCESS);
81 : }
82 :
83 10 : int32_t ProcessTaskAbortHandleCallback(int32_t deviceLogicId, aclrtDeviceTaskAbortStage stage,
84 : uint32_t timeout, void* args)
85 : {
86 10 : HcclUs startut = std::chrono::steady_clock::now();
87 10 : CHK_PTR_NULL(args);
88 10 : auto &commVector = *(static_cast<std::vector<CollComm *> *>(args));
89 10 : HCCL_INFO("[NsRecovery][Callback] ProcessTaskAbortHandleCallback start!");
90 10 : const std::chrono::seconds localtimeout = std::chrono::seconds(timeout);
91 :
92 10 : if (stage == aclrtDeviceTaskAbortStage::ACL_RT_DEVICE_TASK_ABORT_PRE) {
93 4 : auto result = ProcessTaskAbortPre(commVector, localtimeout);
94 4 : if (result != static_cast<int>(TaskAbortResult::TASK_ABORT_SUCCESS)) {
95 2 : return result;
96 : }
97 6 : } else if (stage == aclrtDeviceTaskAbortStage::ACL_RT_DEVICE_TASK_ABORT_POST) {
98 6 : auto result = ProcessTaskAbortPost(commVector, deviceLogicId, localtimeout);
99 6 : if (result != static_cast<int>(TaskAbortResult::TASK_ABORT_SUCCESS)) {
100 2 : return result;
101 : }
102 : }
103 6 : HcclUs endut = std::chrono::steady_clock::now();
104 6 : auto execTime = std::chrono::duration_cast<std::chrono::microseconds>(endut - startut).count();
105 6 : HCCL_RUN_INFO("[NsRecovery][Callback] ProcessTaskAbortHandleCallback success, take time:[%lld]us", execTime);
106 6 : return static_cast<int>(TaskAbortResult::TASK_ABORT_SUCCESS);
107 : }
108 :
109 6 : HcclTaskAbortHandler::HcclTaskAbortHandler()
110 : {
111 6 : std::string name = "HCOMM";
112 6 : Hccl::HrtDeviceAbortRegCallBack(ProcessTaskAbortHandleCallback, static_cast<void *>(&commVector_), name);
113 6 : }
114 :
115 6 : HcclTaskAbortHandler::~HcclTaskAbortHandler()
116 : {
117 6 : std::string name = "HCOMM";
118 6 : Hccl::HrtDeviceAbortRegCallBack(nullptr, nullptr, name);
119 6 : }
120 :
121 245 : HcclTaskAbortHandler &HcclTaskAbortHandler::GetInstance()
122 : {
123 245 : static HcclTaskAbortHandler handler;
124 245 : return handler;
125 : }
126 :
127 102 : HcclResult HcclTaskAbortHandler::Register(CollComm *communicator)
128 : {
129 102 : std::lock_guard<std::mutex> lock(vecMutex);
130 102 : commVector_.push_back(communicator);
131 102 : HCCL_INFO("HcclTaskAbortHandler::Register success, commVector_ size is [%lu]", commVector_.size());
132 :
133 102 : return HCCL_SUCCESS;
134 102 : }
135 :
136 128 : HcclResult HcclTaskAbortHandler::UnRegister(CollComm *communicator)
137 : {
138 128 : std::lock_guard<std::mutex> lock(vecMutex);
139 128 : HCCL_INFO("HcclTaskAbortHandler::UnRegister Begin, commVector_ size is [%lu]", commVector_.size());
140 128 : auto it = std::find(commVector_.begin(), commVector_.end(), communicator);
141 128 : if (it != commVector_.end()) {
142 102 : commVector_.erase(it);
143 : } else {
144 26 : HCCL_WARNING("HcclTaskAbortHandler::UnRegister, comm not found.");
145 : }
146 128 : HCCL_INFO("HcclTaskAbortHandler::UnRegister finish, commVector_ size is [%lu]", commVector_.size());
147 128 : return HCCL_SUCCESS;
148 128 : }
149 : }
|