Skip to content

Commit 78e19c8

Browse files
vegetableysmvegetableysm
committed
Support for specifying RDMA devices when multiple RDMA devices are present.
Signed-off-by: vegetableysm <[email protected]>
1 parent 728477c commit 78e19c8

15 files changed

+231
-111
lines changed

‎CMakeLists.txt

+2
Original file line numberDiff line numberDiff line change
@@ -696,6 +696,7 @@ if (${CMAKE_SYSTEM_NAME} STREQUAL "Linux")
696696
--disable-perf
697697
--disable-efa
698698
--disable-mrail
699+
--with-cuda=no
699700
--enable-verbs > /dev/null
700701
WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}/thirdparty/libfabric
701702
)
@@ -719,6 +720,7 @@ if (${CMAKE_SYSTEM_NAME} STREQUAL "Linux")
719720
--disable-perf
720721
--disable-efa
721722
--disable-mrail
723+
--with-cuda=no
722724
--disable-verbs > /dev/null
723725
WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}/thirdparty/libfabric
724726
)

‎src/client/rpc_client.cc

+34-17
Original file line numberDiff line numberDiff line change
@@ -94,16 +94,18 @@ Status RPCClient::Connect(const std::string& rpc_endpoint) {
9494
Status RPCClient::Connect(const std::string& rpc_endpoint,
9595
std::string const& username,
9696
std::string const& password,
97-
const std::string& rdma_endpoint) {
97+
const std::string& rdma_endpoint,
98+
std::string src_rdma_ednpoint) {
9899
return this->Connect(rpc_endpoint, RootSessionID(), username, password,
99-
rdma_endpoint);
100+
rdma_endpoint, src_rdma_ednpoint);
100101
}
101102

102103
Status RPCClient::Connect(const std::string& rpc_endpoint,
103104
const SessionID session_id,
104105
std::string const& username,
105106
std::string const& password,
106-
const std::string& rdma_endpoint) {
107+
const std::string& rdma_endpoint,
108+
std::string src_rdma_ednpoint) {
107109
size_t pos = rpc_endpoint.find(":");
108110
std::string host, port;
109111
if (pos == std::string::npos) {
@@ -125,28 +127,32 @@ Status RPCClient::Connect(const std::string& rpc_endpoint,
125127

126128
return this->Connect(host, static_cast<uint32_t>(std::stoul(port)),
127129
session_id, username, password, rdma_host,
128-
static_cast<uint32_t>(std::stoul(rdma_port)));
130+
static_cast<uint32_t>(std::stoul(rdma_port)),
131+
src_rdma_ednpoint);
129132
}
130133

131134
Status RPCClient::Connect(const std::string& host, uint32_t port,
132-
const std::string& rdma_host, uint32_t rdma_port) {
135+
const std::string& rdma_host, uint32_t rdma_port,
136+
std::string src_rdma_ednpoint) {
133137
return this->Connect(host, port, RootSessionID(), "", "", rdma_host,
134-
rdma_port);
138+
rdma_port, src_rdma_ednpoint);
135139
}
136140

137141
Status RPCClient::Connect(const std::string& host, uint32_t port,
138142
std::string const& username,
139143
std::string const& password,
140-
const std::string& rdma_host, uint32_t rdma_port) {
144+
const std::string& rdma_host, uint32_t rdma_port,
145+
std::string src_rdma_ednpoint) {
141146
return this->Connect(host, port, RootSessionID(), username, password,
142-
rdma_host, rdma_port);
147+
rdma_host, rdma_port, src_rdma_ednpoint);
143148
}
144149

145150
Status RPCClient::Connect(const std::string& host, uint32_t port,
146151
const SessionID session_id,
147152
std::string const& username,
148153
std::string const& password,
149-
const std::string& rdma_host, uint32_t rdma_port) {
154+
const std::string& rdma_host, uint32_t rdma_port,
155+
std::string src_rdma_ednpoint) {
150156
std::lock_guard<std::recursive_mutex> guard(client_mutex_);
151157
std::string rpc_endpoint = host + ":" + std::to_string(port);
152158
RETURN_ON_ASSERT(!connected_ || rpc_endpoint == rpc_endpoint_);
@@ -183,7 +189,8 @@ Status RPCClient::Connect(const std::string& host, uint32_t port,
183189
instance_id_ = UnspecifiedInstanceID() - 1;
184190

185191
if (rdma_host.length() > 0) {
186-
Status status = ConnectRDMA(rdma_host, rdma_port);
192+
src_rdma_endpoint_ = src_rdma_ednpoint;
193+
Status status = ConnectRDMA(rdma_host, rdma_port, src_rdma_ednpoint);
187194
if (status.ok()) {
188195
rdma_endpoint_ = rdma_host + ":" + std::to_string(rdma_port);
189196
std::cout << "Connected to RPC server: " << rpc_endpoint
@@ -192,33 +199,38 @@ Status RPCClient::Connect(const std::string& host, uint32_t port,
192199
} else {
193200
std::cout << "Connect RDMA server failed! Fall back to RPC mode. Error:"
194201
<< status.message() << std::endl;
202+
std::cout << "Failed src_rdma_ednpoint: " << src_rdma_ednpoint
203+
<< std::endl;
195204
}
196205
}
197206

198207
return Status::OK();
199208
}
200209

201-
Status RPCClient::ConnectRDMA(const std::string& rdma_host,
202-
uint32_t rdma_port) {
210+
Status RPCClient::ConnectRDMA(const std::string& rdma_host, uint32_t rdma_port,
211+
std::string src_rdma_endpoint) {
203212
if (this->rdma_connected_) {
204213
return Status::OK();
205214
}
206215

207216
RETURN_ON_ERROR(RDMAClientCreator::Create(this->rdma_client_, rdma_host,
208-
static_cast<int>(rdma_port)));
217+
static_cast<int>(rdma_port),
218+
src_rdma_endpoint));
209219

210220
int retry = 0;
211221
do {
212-
if (this->rdma_client_->Connect().ok()) {
222+
Status status = this->rdma_client_->Connect();
223+
if (status.ok()) {
213224
break;
214225
}
215226
if (retry == 10) {
216227
return Status::Invalid("Failed to connect to RDMA server.");
217228
}
218229
retry++;
219230
usleep(300 * 1000);
220-
std::cout << "Connect rdma server failed! retry: " << retry << " times."
221-
<< std::endl;
231+
std::cout << "Connect rdma server failed! Error:" + status.message() +
232+
"retry: "
233+
<< retry << " times." << std::endl;
222234
} while (true);
223235
this->rdma_connected_ = true;
224236
return Status::OK();
@@ -272,6 +284,9 @@ Status RPCClient::RDMAReleaseMemInfo(RegisterMemInfo& remote_info) {
272284

273285
Status RPCClient::StopRDMA() {
274286
if (!rdma_connected_) {
287+
RETURN_ON_ERROR(
288+
RDMAClientCreator::Release(RDMAClientCreator::buildConnectionKey(
289+
rdma_endpoint_, src_rdma_endpoint_)));
275290
return Status::OK();
276291
}
277292
rdma_connected_ = false;
@@ -285,7 +300,9 @@ Status RPCClient::StopRDMA() {
285300

286301
RETURN_ON_ERROR(rdma_client_->Stop());
287302
RETURN_ON_ERROR(rdma_client_->Close());
288-
RETURN_ON_ERROR(RDMAClientCreator::Release(rdma_endpoint_));
303+
RETURN_ON_ERROR(
304+
RDMAClientCreator::Release(RDMAClientCreator::buildConnectionKey(
305+
rdma_endpoint_, src_rdma_endpoint_)));
289306

290307
return Status::OK();
291308
}

‎src/client/rpc_client.h

+13-6
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,8 @@ class RPCClient final : public ClientBase {
8888
*/
8989
Status Connect(const std::string& rpc_endpoint, std::string const& username,
9090
std::string const& password,
91-
const std::string& rdma_endpoint = "");
91+
const std::string& rdma_endpoint = "",
92+
std::string src_rdma_ednpoint = "");
9293

9394
/**
9495
* @brief Connect to vineyardd using the given TCP endpoint `rpc_endpoint`.
@@ -104,7 +105,8 @@ class RPCClient final : public ClientBase {
104105
Status Connect(const std::string& rpc_endpoint, const SessionID session_id,
105106
std::string const& username = "",
106107
std::string const& password = "",
107-
const std::string& rdma_endpoint = "");
108+
const std::string& rdma_endpoint = "",
109+
std::string src_rdma_ednpoint = "");
108110

109111
/**
110112
* @brief Connect to vineyardd using the given TCP `host` and `port`.
@@ -117,7 +119,8 @@ class RPCClient final : public ClientBase {
117119
* @return Status that indicates whether the connect has succeeded.
118120
*/
119121
Status Connect(const std::string& host, uint32_t port,
120-
const std::string& rdma_host = "", uint32_t rdma_port = -1);
122+
const std::string& rdma_host = "", uint32_t rdma_port = -1,
123+
std::string src_rdma_ednpoint = "");
121124

122125
/**
123126
* @brief Connect to vineyardd using the given TCP `host` and `port`.
@@ -131,7 +134,8 @@ class RPCClient final : public ClientBase {
131134
*/
132135
Status Connect(const std::string& host, uint32_t port,
133136
std::string const& username, std::string const& password,
134-
const std::string& rdma_host = "", uint32_t rdma_port = -1);
137+
const std::string& rdma_host = "", uint32_t rdma_port = -1,
138+
std::string src_rdma_ednpoint = "");
135139

136140
/**
137141
* @brief Connect to vineyardd using the given TCP `host` and `port`.
@@ -147,7 +151,8 @@ class RPCClient final : public ClientBase {
147151
Status Connect(const std::string& host, uint32_t port,
148152
const SessionID session_id, std::string const& username = "",
149153
std::string const& password = "",
150-
const std::string& rdma_host = "", uint32_t rdma_port = -1);
154+
const std::string& rdma_host = "", uint32_t rdma_port = -1,
155+
std::string src_rdma_ednpoint = "");
151156

152157
/**
153158
* @brief Create a new client using self endpoint.
@@ -436,7 +441,8 @@ class RPCClient final : public ClientBase {
436441
const std::string rdma_endpoint() { return rdma_endpoint_; }
437442

438443
private:
439-
Status ConnectRDMA(const std::string& rdma_host, uint32_t rdma_port);
444+
Status ConnectRDMA(const std::string& rdma_host, uint32_t rdma_port,
445+
std::string src_rdma_endpoint = "");
440446

441447
Status StopRDMA();
442448

@@ -479,6 +485,7 @@ class RPCClient final : public ClientBase {
479485
std::string rdma_endpoint_;
480486
std::shared_ptr<RDMAClient> rdma_client_;
481487
mutable bool rdma_connected_ = false;
488+
std::string src_rdma_endpoint_ = "";
482489

483490
friend class Client;
484491
};

‎src/common/rdma/rdma.cc

+19-4
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ Status IRDMA::RegisterMemory(fid_mr** mr, fid_domain* domain, void* address,
108108
return Status::IOError("Failed to register memory region:" +
109109
std::to_string(ret));
110110
}
111-
CHECK_ERROR(!ret, "Failed to register memory region:" + std::to_string(ret));
111+
CHECK_ERROR(ret, "Failed to register memory region:" + std::to_string(ret));
112112

113113
mr_desc = fi_mr_desc(*mr);
114114

@@ -177,10 +177,25 @@ int IRDMA::GetCompletion(fid_cq* cq, int timeout, void** context) {
177177
return ret < 0 ? ret : 0;
178178
}
179179

180-
void IRDMA::FreeInfo(fi_info* info) {
181-
if (info) {
182-
fi_freeinfo(info);
180+
void IRDMA::FreeInfo(fi_info* info, bool is_hints) {
181+
if (!info) {
182+
return;
183183
}
184+
185+
if (is_hints) {
186+
if (info->src_addr) {
187+
free(info->src_addr);
188+
info->src_addr = nullptr;
189+
info->src_addrlen = 0;
190+
}
191+
if (info->dest_addr) {
192+
free(info->dest_addr);
193+
info->dest_addr = nullptr;
194+
info->dest_addrlen = 0;
195+
}
196+
}
197+
198+
fi_freeinfo(info);
184199
}
185200

186201
} // namespace vineyard

‎src/common/rdma/rdma.h

+1-1
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ class IRDMA {
5555

5656
static int GetCompletion(fid_cq* cq, int timeout, void** context);
5757

58-
static void FreeInfo(fi_info* info);
58+
static void FreeInfo(fi_info* info, bool is_hints);
5959

6060
template <typename FIDType>
6161
static Status CloseResource(FIDType* res, const char* resource_name) {

0 commit comments

Comments
 (0)