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