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 <mutex>
13 : #include "log.h"
14 : #include "coll_comm.h"
15 : #include "ccu_device_pub.h"
16 :
17 : namespace hccl {
18 : using HcclUs = std::chrono::steady_clock::time_point;
19 :
20 5 : int32_t ProcessTaskAbortPre(const std::vector<CollComm*>& commVector, const std::chrono::seconds& localtimeout)
21 : {
22 5 : HcclResult ret = HCCL_SUCCESS;
23 5 : bool isUseTimeOut = localtimeout != std::chrono::seconds(0);
24 5 : std::chrono::seconds elapsed{};
25 8 : for (auto& comm : commVector) {
26 5 : 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 3 : ret = comm->Suspend();
32 : }
33 5 : 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 3 : HCCL_INFO("[NsRecovery]finish suspend success");
38 3 : if (isUseTimeOut) {
39 1 : CHK_PRT_RET(
40 : elapsed > localtimeout, HCCL_ERROR("[NsRecovery][suspend] NsRecovery suspend timeOut"),
41 : static_cast<int>(TaskAbortResult::TASK_ABORT_TIMEOUT));
42 : }
43 : }
44 3 : return static_cast<int>(TaskAbortResult::TASK_ABORT_SUCCESS);
45 : }
46 :
47 8 : int32_t ProcessTaskAbortPost(
48 : const std::vector<CollComm*>& commVector, int32_t deviceLogicId, const std::chrono::seconds& localtimeout)
49 : {
50 8 : HcclResult ret = HCCL_SUCCESS;
51 8 : bool isUseTimeOut = localtimeout != std::chrono::seconds(0);
52 8 : std::chrono::seconds elapsed{};
53 8 : if (hcomm::CcuIsInited(deviceLogicId)) {
54 6 : CHK_RET(hcomm::CcuSetTaskKill(deviceLogicId));
55 : } else {
56 2 : HCCL_INFO("[NsRecovery][Callback] CCU not inited, skip CcuSetTaskKill, deviceLogicId[%d]", deviceLogicId);
57 : }
58 12 : for (auto& comm : commVector) {
59 7 : if (isUseTimeOut) {
60 3 : std::chrono::steady_clock::time_point startTime = std::chrono::steady_clock::now();
61 3 : ret = comm->Clean();
62 3 : elapsed = std::chrono::duration_cast<std::chrono::seconds>(std::chrono::steady_clock::now() - startTime);
63 : } else {
64 4 : ret = comm->Clean();
65 : }
66 7 : if (ret != HCCL_SUCCESS && ret != HCCL_E_SUSPENDING) {
67 2 : HCCL_ERROR("[NsRecovery][Callback] finish clean failed, ret = 0x%016llx", HCCL_ERROR_CODE(ret));
68 2 : return static_cast<int>(TaskAbortResult::TASK_ABORT_FAIL);
69 : }
70 5 : HCCL_INFO("[NsRecovery][Callback] finish clean success");
71 5 : if (isUseTimeOut) {
72 2 : CHK_PRT_RET(
73 : elapsed > localtimeout, HCCL_ERROR("[NsRecovery][Callback] NsRecovery Clean timeout"),
74 : static_cast<int>(TaskAbortResult::TASK_ABORT_TIMEOUT));
75 : }
76 : }
77 5 : if (hcomm::CcuIsInited(deviceLogicId)) {
78 3 : CHK_RET(hcomm::CcuSetTaskKillDone(deviceLogicId));
79 : } else {
80 2 : HCCL_INFO("[NsRecovery][Callback] CCU not inited, skip CcuSetTaskKillDone, deviceLogicId[%d]", deviceLogicId);
81 : }
82 5 : return static_cast<int>(TaskAbortResult::TASK_ABORT_SUCCESS);
83 : }
84 :
85 : int32_t
86 14 : ProcessTaskAbortHandleCallback(int32_t deviceLogicId, aclrtDeviceTaskAbortStage stage, uint32_t timeout, void* args)
87 : {
88 14 : HcclUs startut = std::chrono::steady_clock::now();
89 14 : CHK_PTR_NULL(args);
90 13 : auto& commVector = *(static_cast<std::vector<CollComm*>*>(args));
91 13 : HCCL_INFO("[NsRecovery][Callback] ProcessTaskAbortHandleCallback start!");
92 13 : const std::chrono::seconds localtimeout = std::chrono::seconds(timeout);
93 :
94 13 : if (stage == aclrtDeviceTaskAbortStage::ACL_RT_DEVICE_TASK_ABORT_PRE) {
95 5 : auto result = ProcessTaskAbortPre(commVector, localtimeout);
96 5 : if (result != static_cast<int>(TaskAbortResult::TASK_ABORT_SUCCESS)) {
97 2 : return result;
98 : }
99 8 : } else if (stage == aclrtDeviceTaskAbortStage::ACL_RT_DEVICE_TASK_ABORT_POST) {
100 8 : auto result = ProcessTaskAbortPost(commVector, deviceLogicId, localtimeout);
101 8 : if (result != static_cast<int>(TaskAbortResult::TASK_ABORT_SUCCESS)) {
102 3 : return result;
103 : }
104 : }
105 8 : HcclUs endut = std::chrono::steady_clock::now();
106 8 : auto execTime = std::chrono::duration_cast<std::chrono::microseconds>(endut - startut).count();
107 8 : HCCL_RUN_INFO("[NsRecovery][Callback] ProcessTaskAbortHandleCallback success, take time:[%lld]us", execTime);
108 8 : return static_cast<int>(TaskAbortResult::TASK_ABORT_SUCCESS);
109 : }
110 :
111 20 : HcclTaskAbortHandler::HcclTaskAbortHandler()
112 : {
113 20 : std::string name = "HCOMM";
114 20 : Hccl::HrtDeviceAbortRegCallBack(ProcessTaskAbortHandleCallback, static_cast<void*>(&commVector_), name);
115 20 : }
116 :
117 20 : HcclTaskAbortHandler::~HcclTaskAbortHandler()
118 : {
119 20 : std::string name = "HCOMM";
120 20 : Hccl::HrtDeviceAbortRegCallBack(nullptr, nullptr, name);
121 20 : }
122 :
123 168 : HcclResult HcclTaskAbortHandler::Register(CollComm* communicator)
124 : {
125 168 : std::lock_guard<std::mutex> lock(vecMutex_);
126 168 : commVector_.push_back(communicator);
127 168 : HCCL_INFO("HcclTaskAbortHandler::Register success, commVector_ size is [%zu]", commVector_.size());
128 :
129 168 : return HCCL_SUCCESS;
130 168 : }
131 :
132 174 : HcclResult HcclTaskAbortHandler::UnRegister(CollComm* communicator)
133 : {
134 174 : std::lock_guard<std::mutex> lock(vecMutex_);
135 174 : HCCL_INFO("HcclTaskAbortHandler::UnRegister Begin, commVector_ size is [%zu]", commVector_.size());
136 174 : auto it = std::find(commVector_.begin(), commVector_.end(), communicator);
137 174 : if (it != commVector_.end()) {
138 168 : commVector_.erase(it);
139 : } else {
140 6 : HCCL_WARNING("HcclTaskAbortHandler::UnRegister, comm not found.");
141 : }
142 174 : HCCL_INFO("HcclTaskAbortHandler::UnRegister finish, commVector_ size is [%zu]", commVector_.size());
143 174 : return HCCL_SUCCESS;
144 174 : }
145 : } // namespace hccl
|