Line data Source code
1 : /**
2 : * Copyright (c) 2025 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 :
11 : #include "shared_jetty_channel_pool.h"
12 : #include "my_rank.h"
13 : #include "log.h"
14 : #include <algorithm>
15 : #include <functional>
16 :
17 : namespace hccl {
18 :
19 313 : SharedJettyChannelPool& SharedJettyChannelPool::GetInstance()
20 : {
21 313 : static SharedJettyChannelPool instance;
22 313 : return instance;
23 : }
24 :
25 18 : HcclResult SharedJettyChannelPool::ReturnExistingChannels(
26 : MyRank* myRank, const std::string& tag, const EndpointDescPair& epPair, uint32_t requestedNum,
27 : ChannelHandle* outChannels, uint32_t& returnFromExisting, uint32_t& needCreate)
28 : {
29 18 : std::lock_guard<std::mutex> lock(mtx_);
30 18 : auto& tagMap = rankPools_[myRank];
31 18 : auto& epPairMap = tagMap[tag];
32 18 : auto& epChannels = epPairMap[epPair];
33 :
34 18 : uint32_t available = static_cast<uint32_t>(epChannels.channels.size());
35 18 : if (available >= requestedNum) {
36 1 : returnFromExisting = requestedNum;
37 : } else {
38 17 : returnFromExisting = available;
39 17 : needCreate = requestedNum - available;
40 : }
41 :
42 18 : HCCL_INFO(
43 : "[%s] myRank[%p], tag[%s], available[%u], requested[%u], returnFromExisting[%u], needCreate[%u].", __func__,
44 : myRank, tag.c_str(), available, requestedNum, returnFromExisting, needCreate);
45 :
46 22 : for (uint32_t i = 0; i < returnFromExisting; ++i) {
47 4 : uint32_t idx = epChannels.nextReturnIdx % epChannels.channels.size();
48 4 : outChannels[i] = epChannels.channels[idx];
49 4 : epChannels.nextReturnIdx = (epChannels.nextReturnIdx + 1) % epChannels.channels.size();
50 4 : HCCL_INFO("[%s] return existing channel[%u]: handle[0x%llx].", __func__, i, outChannels[i]);
51 : }
52 18 : return HCCL_SUCCESS;
53 18 : }
54 :
55 21 : HcclResult SharedJettyChannelPool::AcquireChannels(
56 : MyRank* myRank, const std::string& tag, const EndpointDescPair& epPair, uint32_t requestedNum,
57 : const std::function<HcclResult(uint32_t, ChannelHandle*)>& createFunc, ChannelHandle* outChannels,
58 : uint32_t* outReusedCount)
59 : {
60 21 : if (myRank == nullptr || requestedNum == 0 || outChannels == nullptr) {
61 3 : HCCL_ERROR(
62 : "[%s] invalid params, myRank[%p], requestedNum[%u], outChannels[%p].", __func__, myRank, requestedNum,
63 : outChannels);
64 3 : return HCCL_E_PARA;
65 : }
66 :
67 18 : uint32_t returnFromExisting = 0;
68 18 : uint32_t needCreate = 0;
69 :
70 : // 第一段（持锁）：查询已有 channel 并计算需新建数量，取走复用句柄后释放锁。
71 18 : CHK_RET(ReturnExistingChannels(myRank, tag, epPair, requestedNum, outChannels, returnFromExisting, needCreate));
72 :
73 : // 第二段（无锁）：执行建链 I/O，避免阻塞其他 myRank/tag 的并发 Acquire。
74 18 : if (needCreate > 0) {
75 17 : ChannelHandle* newChannels = outChannels + returnFromExisting;
76 17 : HcclResult ret = createFunc(needCreate, newChannels);
77 17 : if (ret != HCCL_SUCCESS) {
78 1 : HCCL_ERROR("[%s] createFunc failed, needCreate[%u], ret[%d].", __func__, needCreate, ret);
79 : // 不回退 nextReturnIdx：第一段与第三段之间可能有并发 Acquire 推进了游标，
80 : // 回退会错误覆盖其他线程的推进量。channel 仍在池中，后续仍可通过取模访问到，
81 : // 本轮调用方收到错误，仅失去部分可复用句柄的租约，不影响正确性。
82 1 : return ret;
83 : }
84 :
85 : // 第三段（持锁）：将新建 channel 回填到池，重新定位条目以规避 rehash 导致的引用失效。
86 16 : std::lock_guard<std::mutex> lock(mtx_);
87 16 : auto& tagMap = rankPools_[myRank];
88 16 : auto& epPairMap = tagMap[tag];
89 16 : auto& epChannels = epPairMap[epPair];
90 45 : for (uint32_t i = 0; i < needCreate; ++i) {
91 29 : epChannels.channels.push_back(newChannels[i]);
92 29 : HCCL_INFO(
93 : "[%s] created new channel[%u]: handle[0x%llx], total channels[%zu].", __func__, i, newChannels[i],
94 : epChannels.channels.size());
95 : }
96 16 : }
97 :
98 17 : if (outReusedCount != nullptr) {
99 4 : *outReusedCount = returnFromExisting;
100 : }
101 17 : return HCCL_SUCCESS;
102 : }
103 :
104 : SharedJettyChannelPool::RankPoolIter
105 233 : SharedJettyChannelPool::CollectMyRankChannelsLocked(MyRank* myRank, std::vector<ChannelHandle>& allChannels)
106 : {
107 233 : auto it = rankPools_.find(myRank);
108 233 : if (it == rankPools_.end()) {
109 229 : return it;
110 : }
111 : // 先统计总数并 reserve，避免 push_back 触发多次 realloc
112 4 : uint32_t totalChannels = 0;
113 9 : for (auto& tagEntry : it->second) {
114 10 : for (auto& epEntry : tagEntry.second) {
115 5 : totalChannels += static_cast<uint32_t>(epEntry.second.channels.size());
116 : }
117 : }
118 4 : allChannels.reserve(totalChannels);
119 9 : for (auto& tagEntry : it->second) {
120 10 : for (auto& epEntry : tagEntry.second) {
121 7 : for (ChannelHandle ch : epEntry.second.channels) {
122 2 : allChannels.push_back(ch);
123 : }
124 : }
125 : }
126 4 : return it;
127 : }
128 :
129 231 : HcclResult SharedJettyChannelPool::DestroyAllByMyRank(MyRank* myRank)
130 : {
131 231 : if (myRank == nullptr) {
132 1 : return HCCL_SUCCESS;
133 : }
134 230 : std::lock_guard<std::mutex> lock(mtx_);
135 230 : std::vector<ChannelHandle> allChannels;
136 230 : auto it = CollectMyRankChannelsLocked(myRank, allChannels);
137 230 : if (it == rankPools_.end()) {
138 227 : return HCCL_SUCCESS;
139 : }
140 :
141 3 : if (!allChannels.empty()) {
142 0 : HcclResult ret = static_cast<HcclResult>(HcommChannelDestroy(allChannels.data(), allChannels.size()));
143 0 : if (ret != HCCL_SUCCESS) {
144 0 : HCCL_ERROR("[%s] HcommChannelDestroy failed, channelNum[%zu], ret[%d].", __func__, allChannels.size(), ret);
145 : }
146 : }
147 :
148 3 : rankPools_.erase(it);
149 3 : HCCL_INFO("[%s] destroyed myRank[%p] shared jetty channels, total[%zu].", __func__, myRank, allChannels.size());
150 3 : return HCCL_SUCCESS;
151 230 : }
152 :
153 3 : HcclResult SharedJettyChannelPool::CheckMyRankDestroy(MyRank* myRank)
154 : {
155 3 : std::lock_guard<std::mutex> lock(mtx_);
156 3 : std::vector<ChannelHandle> allChannels;
157 3 : auto it = CollectMyRankChannelsLocked(myRank, allChannels);
158 3 : if (it == rankPools_.end()) {
159 2 : return HCCL_SUCCESS;
160 : }
161 1 : if (!allChannels.empty()) {
162 1 : HCCL_ERROR(
163 : "[%s] cannot destroy myRank[%p], still has [%zu] shared jetty channels.", __func__, myRank,
164 : allChannels.size());
165 1 : return HCCL_E_UNAVAIL;
166 : }
167 0 : return HCCL_SUCCESS;
168 3 : }
169 :
170 9 : void SharedJettyChannelPool::RemoveChannels(
171 : MyRank* myRank, const std::string& tag, const EndpointDescPair& epPair, const ChannelHandle* channels,
172 : uint32_t channelNum)
173 : {
174 9 : if (myRank == nullptr || channels == nullptr || channelNum == 0) {
175 2 : return;
176 : }
177 8 : std::lock_guard<std::mutex> lock(mtx_);
178 8 : auto tagIt = rankPools_.find(myRank);
179 8 : if (tagIt == rankPools_.end()) {
180 0 : return;
181 : }
182 8 : auto epIt = tagIt->second.find(tag);
183 8 : if (epIt == tagIt->second.end()) {
184 1 : return;
185 : }
186 7 : auto pairIt = epIt->second.find(epPair);
187 7 : if (pairIt == epIt->second.end()) {
188 0 : return;
189 : }
190 7 : auto& epChannels = pairIt->second;
191 17 : for (uint32_t i = 0; i < channelNum; ++i) {
192 10 : auto& vec = epChannels.channels;
193 10 : vec.erase(std::remove(vec.begin(), vec.end(), channels[i]), vec.end());
194 : }
195 : // 重置游标避免取模越界（channels 已收缩）
196 7 : if (epChannels.channels.empty()) {
197 5 : epChannels.nextReturnIdx = 0;
198 : } else {
199 2 : epChannels.nextReturnIdx %= epChannels.channels.size();
200 : }
201 7 : HCCL_INFO(
202 : "[%s] removed [%u] channels from pool, remaining[%zu].", __func__, channelNum, epChannels.channels.size());
203 8 : }
204 :
205 : } // namespace hccl
|
