ps/brpc/brpc_ps_server.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #include "brpc_ps_server.h" | ||
| 2 | |||
| 3 | #include <brpc/server.h> | ||
| 4 | #include <fmt/format.h> | ||
| 5 | #include <gflags/gflags.h> | ||
| 6 | |||
| 7 | #include <chrono> | ||
| 8 | #include <cerrno> | ||
| 9 | #include <cstdint> | ||
| 10 | #include <cstring> | ||
| 11 | #include <fstream> | ||
| 12 | #include <memory> | ||
| 13 | #include <stdexcept> | ||
| 14 | #include <string> | ||
| 15 | #include <thread> | ||
| 16 | #include <vector> | ||
| 17 | |||
| 18 | #include "base/array.h" | ||
| 19 | #include "base/base.h" | ||
| 20 | #include "base/factory.h" | ||
| 21 | #include "base/flatc.h" | ||
| 22 | #include "base/log.h" | ||
| 23 | #include "base/timer.h" | ||
| 24 | #include "ps/base/base_ps_server.h" | ||
| 25 | #include "ps/base/cache_ps_impl.h" | ||
| 26 | #include "ps/base/parameters.h" | ||
| 27 | #include "ps_brpc.pb.h" | ||
| 28 | #include "recstore_config.h" | ||
| 29 | #include "src/base/config.h" | ||
| 30 | |||
| 31 | #ifdef ENABLE_PERF_REPORT | ||
| 32 | # include <chrono> | ||
| 33 | # include <cstdlib> | ||
| 34 | # include "base/report/report_client.h" | ||
| 35 | #endif | ||
| 36 | |||
| 37 | using recstoreps_brpc::CommandRequest; | ||
| 38 | using recstoreps_brpc::CommandResponse; | ||
| 39 | using recstoreps_brpc::GetParameterRequest; | ||
| 40 | using recstoreps_brpc::GetParameterResponse; | ||
| 41 | using recstoreps_brpc::InitEmbeddingTableRequest; | ||
| 42 | using recstoreps_brpc::InitEmbeddingTableResponse; | ||
| 43 | using recstoreps_brpc::PSCommand; | ||
| 44 | using recstoreps_brpc::PutParameterRequest; | ||
| 45 | using recstoreps_brpc::PutParameterResponse; | ||
| 46 | using recstoreps_brpc::UpdateParameterRequest; | ||
| 47 | using recstoreps_brpc::UpdateParameterResponse; | ||
| 48 | |||
| 49 | DEFINE_string(brpc_config_path, "", "config file path"); | ||
| 50 | DEFINE_int32(brpc_server_port, 15000, "bRPC server port"); | ||
| 51 | DEFINE_int32(local_shard_id, | ||
| 52 | -1, | ||
| 53 | "Only start the specified shard in multi-shard bRPC mode; " | ||
| 54 | "-1 means start all configured shards"); | ||
| 55 | DEFINE_int32(brpc_server_num_threads, | ||
| 56 | 0, | ||
| 57 | "Number of threads for bRPC server, 0 means auto"); | ||
| 58 | DEFINE_bool(brpc_ps_use_rdma, | ||
| 59 | false, | ||
| 60 | "Use RDMA transport for the brpc PS server (requires brpc built " | ||
| 61 | "with WITH_RDMA=ON)."); | ||
| 62 | |||
| 63 | namespace { | ||
| 64 | |||
| 65 | // Mirror the client: allow enabling RDMA and selecting the HCA via env vars so | ||
| 66 | // the ps_server launcher does not require extra command-line flags. | ||
| 67 | // RECSTORE_BRPC_USE_RDMA=1 -> enable RDMA transport | ||
| 68 | // RECSTORE_BRPC_RDMA_DEVICE=mlx5_0 -> select the HCA (maps to -rdma_device) | ||
| 69 | ✗ | bool ResolveBrpcServerUseRdmaFromEnv(bool fallback) { | |
| 70 | ✗ | const char* value = std::getenv("RECSTORE_BRPC_USE_RDMA"); | |
| 71 | ✗ | if (value == nullptr || *value == '\0') { | |
| 72 | ✗ | return fallback; | |
| 73 | } | ||
| 74 | ✗ | return std::string(value) != "0"; | |
| 75 | } | ||
| 76 | |||
| 77 | ✗ | void ApplyBrpcServerRdmaDeviceFromEnv() { | |
| 78 | ✗ | const char* dev = std::getenv("RECSTORE_BRPC_RDMA_DEVICE"); | |
| 79 | ✗ | if (dev != nullptr && *dev != '\0') { | |
| 80 | ✗ | google::SetCommandLineOption("rdma_device", dev); | |
| 81 | } | ||
| 82 | ✗ | } | |
| 83 | |||
| 84 | } // namespace | ||
| 85 | |||
| 86 | namespace recstore { | ||
| 87 | |||
| 88 | namespace { | ||
| 89 | |||
| 90 | ✗ | void AppendShardSuffixIfPresent( | |
| 91 | nlohmann::json& config_node, const char* key, int shard_id) { | ||
| 92 | ✗ | if (!config_node.contains(key) || !config_node[key].is_string()) { | |
| 93 | ✗ | return; | |
| 94 | } | ||
| 95 | ✗ | config_node[key] = | |
| 96 | ✗ | config_node[key].get<std::string>() + "_" + std::to_string(shard_id); | |
| 97 | } | ||
| 98 | |||
| 99 | ✗ | void AppendShardSuffixToNestedFilePaths(nlohmann::json& node, int shard_id) { | |
| 100 | ✗ | if (node.is_object()) { | |
| 101 | ✗ | for (auto& item : node.items()) { | |
| 102 | ✗ | if (item.key() == "file_path" && item.value().is_string()) { | |
| 103 | ✗ | item.value() = | |
| 104 | ✗ | item.value().get<std::string>() + "_" + std::to_string(shard_id); | |
| 105 | ✗ | continue; | |
| 106 | } | ||
| 107 | ✗ | AppendShardSuffixToNestedFilePaths(item.value(), shard_id); | |
| 108 | ✗ | } | |
| 109 | ✗ | return; | |
| 110 | } | ||
| 111 | ✗ | if (node.is_array()) { | |
| 112 | ✗ | for (auto& item : node) { | |
| 113 | ✗ | AppendShardSuffixToNestedFilePaths(item, shard_id); | |
| 114 | } | ||
| 115 | } | ||
| 116 | } | ||
| 117 | |||
| 118 | ✗ | bool ExtractPayloadBytes( | |
| 119 | const brpc::Controller* cntl, | ||
| 120 | const std::string& proto_bytes, | ||
| 121 | std::string* payload_storage, | ||
| 122 | const char** payload_data, | ||
| 123 | int* payload_size) { | ||
| 124 | ✗ | if (!cntl->request_attachment().empty()) { | |
| 125 | ✗ | payload_storage->clear(); | |
| 126 | ✗ | cntl->request_attachment().copy_to(payload_storage); | |
| 127 | ✗ | *payload_data = payload_storage->data(); | |
| 128 | ✗ | *payload_size = payload_storage->size(); | |
| 129 | ✗ | return true; | |
| 130 | } | ||
| 131 | ✗ | if (!proto_bytes.empty()) { | |
| 132 | ✗ | *payload_data = proto_bytes.data(); | |
| 133 | ✗ | *payload_size = proto_bytes.size(); | |
| 134 | ✗ | return true; | |
| 135 | } | ||
| 136 | ✗ | *payload_data = nullptr; | |
| 137 | ✗ | *payload_size = 0; | |
| 138 | ✗ | return false; | |
| 139 | } | ||
| 140 | |||
| 141 | std::vector<nlohmann::json> | ||
| 142 | 6 | SelectShardConfigsInternal(const nlohmann::json& cache_ps_config, | |
| 143 | const std::optional<int>& local_shard_id) { | ||
| 144 | 6 | std::vector<nlohmann::json> selected; | |
| 145 |
3/6✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 6 times.
|
12 | if (!cache_ps_config.contains("servers") || |
| 146 |
2/4✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 6 times.
|
6 | !cache_ps_config["servers"].is_array()) { |
| 147 | ✗ | return selected; | |
| 148 | } | ||
| 149 | |||
| 150 |
6/10✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✓ Branch 6 taken 12 times.
✗ Branch 7 not taken.
✓ Branch 9 taken 12 times.
✗ Branch 10 not taken.
✓ Branch 12 taken 18 times.
✗ Branch 13 not taken.
✓ Branch 14 taken 12 times.
✓ Branch 15 taken 6 times.
|
18 | for (const auto& server_config : cache_ps_config["servers"]) { |
| 151 |
2/2✓ Branch 1 taken 4 times.
✓ Branch 2 taken 8 times.
|
12 | if (!local_shard_id.has_value()) { |
| 152 |
1/2✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
|
4 | selected.push_back(server_config); |
| 153 | 4 | continue; | |
| 154 | } | ||
| 155 |
3/6✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 8 times.
|
16 | if (!server_config.contains("shard") || |
| 156 |
2/4✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 8 times.
|
8 | !server_config["shard"].is_number_integer()) { |
| 157 | ✗ | continue; | |
| 158 | } | ||
| 159 |
4/6✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 8 times.
✗ Branch 5 not taken.
✓ Branch 7 taken 2 times.
✓ Branch 8 taken 6 times.
|
8 | if (server_config["shard"].get<int>() == *local_shard_id) { |
| 160 |
1/2✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
|
2 | selected.push_back(server_config); |
| 161 | } | ||
| 162 | } | ||
| 163 | 6 | return selected; | |
| 164 | ✗ | } | |
| 165 | |||
| 166 | } // namespace | ||
| 167 | |||
| 168 | std::vector<nlohmann::json> | ||
| 169 | 6 | SelectBRPCShardConfigs(const nlohmann::json& cache_ps_config, | |
| 170 | const std::optional<int>& local_shard_id) { | ||
| 171 | 6 | return SelectShardConfigsInternal(cache_ps_config, local_shard_id); | |
| 172 | } | ||
| 173 | |||
| 174 | ✗ | BRPCParameterServiceImpl::BRPCParameterServiceImpl(CachePS* cache_ps) | |
| 175 | ✗ | : cache_ps_(cache_ps) { | |
| 176 | ✗ | start_time_ = std::chrono::steady_clock::now(); | |
| 177 | ✗ | } | |
| 178 | |||
| 179 | ✗ | void BRPCParameterServiceImpl::ResetMetrics() { | |
| 180 | ✗ | total_get_requests_ = 0; | |
| 181 | ✗ | total_put_requests_ = 0; | |
| 182 | ✗ | total_get_keys_ = 0; | |
| 183 | ✗ | total_put_keys_ = 0; | |
| 184 | ✗ | total_get_bytes_ = 0; | |
| 185 | ✗ | total_put_bytes_ = 0; | |
| 186 | ✗ | start_time_ = std::chrono::steady_clock::now(); | |
| 187 | ✗ | } | |
| 188 | |||
| 189 | ✗ | void BRPCParameterServiceImpl::PrintMetrics(const std::string& table_name, | |
| 190 | const std::string& unique_id) { | ||
| 191 | ✗ | auto now = std::chrono::steady_clock::now(); | |
| 192 | ✗ | double elapsed_s = std::chrono::duration<double>(now - start_time_).count(); | |
| 193 | ✗ | if (elapsed_s > 0) { | |
| 194 | double overall_qps = | ||
| 195 | ✗ | (total_get_requests_ + total_put_requests_) / elapsed_s; | |
| 196 | double overall_throughput_mbps = | ||
| 197 | ✗ | ((total_get_bytes_ + total_put_bytes_) / 1024.0 / 1024.0) / elapsed_s; | |
| 198 | |||
| 199 | // Report QPS and throughput metrics | ||
| 200 | // report(table_name.c_str(), unique_id.c_str(), "overall_qps", | ||
| 201 | // overall_qps); report(table_name.c_str(), | ||
| 202 | // unique_id.c_str(), | ||
| 203 | // "overall_throughput_mbps", | ||
| 204 | // overall_throughput_mbps); | ||
| 205 | } | ||
| 206 | ✗ | } | |
| 207 | |||
| 208 | ✗ | void BRPCParameterServiceImpl::GetParameter( | |
| 209 | google::protobuf::RpcController* controller, | ||
| 210 | const GetParameterRequest* request, | ||
| 211 | GetParameterResponse* response, | ||
| 212 | google::protobuf::Closure* done) { | ||
| 213 | ✗ | brpc::ClosureGuard done_guard(done); | |
| 214 | ✗ | brpc::Controller* cntl = static_cast<brpc::Controller*>(controller); | |
| 215 | |||
| 216 | #ifdef ENABLE_PERF_REPORT | ||
| 217 | auto start_time = std::chrono::high_resolution_clock::now(); | ||
| 218 | uint64_t trace_id = cntl->log_id(); | ||
| 219 | std::string unique_id = "embread_debug" + std::to_string(trace_id); | ||
| 220 | #endif | ||
| 221 | ✗ | std::string keys_storage; | |
| 222 | ✗ | const char* keys_data = nullptr; | |
| 223 | ✗ | int keys_size = 0; | |
| 224 | ✗ | ExtractPayloadBytes( | |
| 225 | cntl, request->keys(), &keys_storage, &keys_data, &keys_size); | ||
| 226 | ✗ | base::ConstArray<uint64_t> keys_array; | |
| 227 | ✗ | keys_array.SetData(keys_data, keys_size); | |
| 228 | ✗ | if (keys_size % static_cast<int>(sizeof(uint64_t)) != 0) { | |
| 229 | ✗ | LOG(ERROR) << "GetParameter invalid keys payload size=" << keys_size; | |
| 230 | ✗ | return; | |
| 231 | } | ||
| 232 | ✗ | bool isPerf = request->has_perf() && request->perf(); | |
| 233 | |||
| 234 | ✗ | if (isPerf) { | |
| 235 | ✗ | xmh::PerfCounter::Record("PS Get Keys", keys_array.Size()); | |
| 236 | } | ||
| 237 | |||
| 238 | ✗ | xmh::Timer timer_ps_get_req("PS GetParameter Req"); | |
| 239 | ✗ | ParameterCompressor compressor; | |
| 240 | |||
| 241 | ✗ | RECSTORE_LOG_EVERY_MS(INFO, 1000) | |
| 242 | ✗ | << "[bRPC PS] Getting " << keys_array.Size() << " keys"; | |
| 243 | |||
| 244 | ✗ | int total_dim = 0; | |
| 245 | |||
| 246 | #ifdef ENABLE_PERF_REPORT | ||
| 247 | auto cache_loop_start = std::chrono::high_resolution_clock::now(); | ||
| 248 | #endif | ||
| 249 | ✗ | std::vector<ParameterPack> packs; | |
| 250 | ✗ | packs.reserve(keys_array.Size()); | |
| 251 | ✗ | cache_ps_->GetParameterRun2Completion(keys_array, packs, 0); | |
| 252 | |||
| 253 | { | ||
| 254 | ✗ | int est_bytes = 0; | |
| 255 | ✗ | for (const auto& pack : packs) { | |
| 256 | ✗ | est_bytes += ParameterCompressItem::GetSize(pack.dim); | |
| 257 | } | ||
| 258 | ✗ | compressor.Reserve(static_cast<int>(packs.size()), est_bytes); | |
| 259 | } | ||
| 260 | |||
| 261 | ✗ | for (auto& pack : packs) { | |
| 262 | ✗ | compressor.AddItem(pack, nullptr); | |
| 263 | ✗ | total_dim += pack.dim; | |
| 264 | } | ||
| 265 | #ifdef ENABLE_PERF_REPORT | ||
| 266 | auto cache_loop_end = std::chrono::high_resolution_clock::now(); | ||
| 267 | auto cache_loop_duration = | ||
| 268 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 269 | cache_loop_end - cache_loop_start) | ||
| 270 | .count(); | ||
| 271 | double cache_loop_start_us = | ||
| 272 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 273 | cache_loop_start.time_since_epoch()) | ||
| 274 | .count(); | ||
| 275 | |||
| 276 | std::string report_id_cache = | ||
| 277 | "brpc_server::GetParameter|" + | ||
| 278 | std::to_string(static_cast<uint64_t>(cache_loop_start_us)); | ||
| 279 | report("embread_stages", | ||
| 280 | report_id_cache.c_str(), | ||
| 281 | "cache_lookup_us", | ||
| 282 | static_cast<double>(cache_loop_duration)); | ||
| 283 | |||
| 284 | FlameGraphData cache_loop_fg = { | ||
| 285 | "brpc_server::CacheGet_Loop", | ||
| 286 | cache_loop_start_us, | ||
| 287 | 4, // level | ||
| 288 | static_cast<double>(cache_loop_duration), | ||
| 289 | static_cast<double>(cache_loop_duration)}; | ||
| 290 | if (trace_id != 0) { | ||
| 291 | std::string cache_unique_id = | ||
| 292 | "embread_debug|" + | ||
| 293 | std::to_string(static_cast<uint64_t>(cache_loop_start_us)); | ||
| 294 | report_flame_graph( | ||
| 295 | "emb_read_flame_map", cache_unique_id.c_str(), cache_loop_fg); | ||
| 296 | } | ||
| 297 | |||
| 298 | auto toblock_start = std::chrono::high_resolution_clock::now(); | ||
| 299 | #endif | ||
| 300 | |||
| 301 | ✗ | compressor.AppendToIOBuf(&cntl->response_attachment()); | |
| 302 | |||
| 303 | #ifdef ENABLE_PERF_REPORT | ||
| 304 | auto toblock_end = std::chrono::high_resolution_clock::now(); | ||
| 305 | auto toblock_duration = | ||
| 306 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 307 | toblock_end - toblock_start) | ||
| 308 | .count(); | ||
| 309 | double toblock_start_us = | ||
| 310 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 311 | toblock_start.time_since_epoch()) | ||
| 312 | .count(); | ||
| 313 | FlameGraphData toblock_fg = { | ||
| 314 | "brpc_server::Compressor_ToBlock", | ||
| 315 | toblock_start_us, | ||
| 316 | 4, // level | ||
| 317 | static_cast<double>(toblock_duration), | ||
| 318 | static_cast<double>(toblock_duration)}; | ||
| 319 | if (trace_id != 0) { | ||
| 320 | std::string toblock_unique_id = | ||
| 321 | "embread_debug|" + | ||
| 322 | std::to_string(static_cast<uint64_t>(toblock_start_us)); | ||
| 323 | report_flame_graph( | ||
| 324 | "emb_read_flame_map", toblock_unique_id.c_str(), toblock_fg); | ||
| 325 | } | ||
| 326 | #endif | ||
| 327 | |||
| 328 | ✗ | total_get_requests_++; | |
| 329 | ✗ | total_get_keys_ += keys_array.Size(); | |
| 330 | ✗ | total_get_bytes_ += total_dim * sizeof(float); | |
| 331 | |||
| 332 | ✗ | if (isPerf) { | |
| 333 | ✗ | timer_ps_get_req.end(); | |
| 334 | } else { | ||
| 335 | ✗ | timer_ps_get_req.destroy(); | |
| 336 | } | ||
| 337 | |||
| 338 | #ifdef ENABLE_PERF_REPORT | ||
| 339 | auto end_time = std::chrono::high_resolution_clock::now(); | ||
| 340 | auto duration = | ||
| 341 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 342 | end_time - start_time) | ||
| 343 | .count(); | ||
| 344 | report("ps_server_latency", | ||
| 345 | "GetParameter", | ||
| 346 | "latency_us", | ||
| 347 | static_cast<double>(duration)); | ||
| 348 | |||
| 349 | double start_us_for_key = | ||
| 350 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 351 | start_time.time_since_epoch()) | ||
| 352 | .count(); | ||
| 353 | std::string op_latency_key = | ||
| 354 | "EmbRead|" + std::to_string(static_cast<uint64_t>(start_us_for_key)); | ||
| 355 | report("op_latency", | ||
| 356 | op_latency_key.c_str(), | ||
| 357 | "recserver_us", | ||
| 358 | static_cast<double>(duration)); | ||
| 359 | |||
| 360 | double start_us = | ||
| 361 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 362 | start_time.time_since_epoch()) | ||
| 363 | .count(); | ||
| 364 | |||
| 365 | std::string report_id = "brpc_server::GetParameter|" + | ||
| 366 | std::to_string(static_cast<uint64_t>(start_us)); | ||
| 367 | |||
| 368 | report("embread_stages", | ||
| 369 | report_id.c_str(), | ||
| 370 | "duration_us", | ||
| 371 | static_cast<double>(duration)); | ||
| 372 | |||
| 373 | report("embread_stages", | ||
| 374 | report_id.c_str(), | ||
| 375 | "request_size", | ||
| 376 | static_cast<double>(keys_array.Size())); | ||
| 377 | |||
| 378 | FlameGraphData fg_data = { | ||
| 379 | "brpc_server::GetParameter", | ||
| 380 | start_us, | ||
| 381 | 3, // level | ||
| 382 | static_cast<double>(duration), | ||
| 383 | static_cast<double>(duration)}; | ||
| 384 | if (trace_id != 0) { | ||
| 385 | std::string req_unique_id = | ||
| 386 | "embread_debug|" + std::to_string(static_cast<uint64_t>(start_us)); | ||
| 387 | report_flame_graph("emb_read_flame_map", req_unique_id.c_str(), fg_data); | ||
| 388 | } | ||
| 389 | #endif | ||
| 390 | ✗ | } | |
| 391 | |||
| 392 | ✗ | void BRPCParameterServiceImpl::Command( | |
| 393 | google::protobuf::RpcController* controller, | ||
| 394 | const CommandRequest* request, | ||
| 395 | CommandResponse* response, | ||
| 396 | google::protobuf::Closure* done) { | ||
| 397 | ✗ | brpc::ClosureGuard done_guard(done); | |
| 398 | ✗ | brpc::Controller* cntl = static_cast<brpc::Controller*>(controller); | |
| 399 | |||
| 400 | ✗ | if (request->command() == recstoreps_brpc::PSCommand::CLEAR_PS) { | |
| 401 | ✗ | LOG(WARNING) << "[PS Command] Clear All"; | |
| 402 | ✗ | cache_ps_->Clear(); | |
| 403 | ✗ | } else if (request->command() == recstoreps_brpc::PSCommand::RELOAD_PS) { | |
| 404 | ✗ | LOG(WARNING) << "[PS Command] Reload PS"; | |
| 405 | ✗ | CHECK_NE(request->arg1().size(), 0); | |
| 406 | ✗ | CHECK_NE(request->arg2().size(), 0); | |
| 407 | ✗ | CHECK_EQ(request->arg1().size(), 1); | |
| 408 | ✗ | LOG(WARNING) << "model_config_path = " << request->arg1()[0]; | |
| 409 | ✗ | for (int i = 0; i < request->arg2().size(); i++) { | |
| 410 | ✗ | LOG(WARNING) << fmt::format("emb_file {}: {}", i, request->arg2()[i]); | |
| 411 | } | ||
| 412 | ✗ | std::vector<std::string> arg1; | |
| 413 | ✗ | for (auto& each : request->arg1()) { | |
| 414 | ✗ | arg1.push_back(each); | |
| 415 | } | ||
| 416 | ✗ | std::vector<std::string> arg2; | |
| 417 | ✗ | for (auto& each : request->arg2()) { | |
| 418 | ✗ | arg2.push_back(each); | |
| 419 | } | ||
| 420 | ✗ | cache_ps_->Initialize(arg1, arg2); | |
| 421 | ✗ | } else if (request->command() == recstoreps_brpc::PSCommand::LOAD_FAKE_DATA) { | |
| 422 | ✗ | if (request->arg1_size() != 1 || | |
| 423 | ✗ | static_cast<size_t>(request->arg1(0).size()) != sizeof(int64_t)) { | |
| 424 | ✗ | LOG(ERROR) << "LOAD_FAKE_DATA: arg1 must be one " << sizeof(int64_t) | |
| 425 | ✗ | << "-byte int64_t (requested reply payload size)"; | |
| 426 | ✗ | cntl->SetFailed(EINVAL, "LOAD_FAKE_DATA invalid arg1 size"); | |
| 427 | ✗ | return; | |
| 428 | } | ||
| 429 | ✗ | int64_t payload_bytes = 0; | |
| 430 | ✗ | std::memcpy(&payload_bytes, request->arg1(0).data(), sizeof(int64_t)); | |
| 431 | ✗ | if (payload_bytes < 0) { | |
| 432 | ✗ | LOG(ERROR) << "LOAD_FAKE_DATA: payload_bytes must be non-negative, got " | |
| 433 | ✗ | << payload_bytes; | |
| 434 | ✗ | cntl->SetFailed( | |
| 435 | EINVAL, "LOAD_FAKE_DATA payload_bytes must be non-negative"); | ||
| 436 | ✗ | return; | |
| 437 | } | ||
| 438 | ✗ | constexpr int64_t kMaxReplyPayload = 16 * 1024 * 1024; | |
| 439 | ✗ | if (payload_bytes > kMaxReplyPayload) { | |
| 440 | ✗ | LOG(ERROR) << "LOAD_FAKE_DATA: payload_bytes " << payload_bytes | |
| 441 | ✗ | << " exceeds cap " << kMaxReplyPayload; | |
| 442 | ✗ | cntl->SetFailed(EINVAL, "LOAD_FAKE_DATA payload too large"); | |
| 443 | ✗ | return; | |
| 444 | } | ||
| 445 | ✗ | std::string fake(static_cast<size_t>(payload_bytes), '\xab'); | |
| 446 | ✗ | response->set_reply(std::move(fake)); | |
| 447 | ✗ | } else if (request->command() == recstoreps_brpc::PSCommand::DUMP_FAKE_DATA) { | |
| 448 | ✗ | if (request->arg1_size() != 1 || | |
| 449 | ✗ | static_cast<size_t>(request->arg1(0).size()) != sizeof(int64_t)) { | |
| 450 | ✗ | LOG(ERROR) << "DUMP_FAKE_DATA: arg1 must be one " << sizeof(int64_t) | |
| 451 | ✗ | << "-byte int64_t (payload bytes n)"; | |
| 452 | ✗ | cntl->SetFailed(EINVAL, "DUMP_FAKE_DATA invalid arg1 size"); | |
| 453 | ✗ | return; | |
| 454 | } | ||
| 455 | ✗ | int64_t n = 0; | |
| 456 | ✗ | std::memcpy(&n, request->arg1(0).data(), sizeof(int64_t)); | |
| 457 | ✗ | if (n <= 0) { | |
| 458 | ✗ | LOG(ERROR) << "DUMP_FAKE_DATA: n must be positive"; | |
| 459 | ✗ | cntl->SetFailed(EINVAL, "DUMP_FAKE_DATA n must be positive"); | |
| 460 | ✗ | return; | |
| 461 | } | ||
| 462 | ✗ | if (n % static_cast<int64_t>(sizeof(float)) != 0) { | |
| 463 | ✗ | LOG(ERROR) << "DUMP_FAKE_DATA: n must be a multiple of " << sizeof(float); | |
| 464 | ✗ | cntl->SetFailed( | |
| 465 | EINVAL, "DUMP_FAKE_DATA n must be multiple of sizeof(float)"); | ||
| 466 | ✗ | return; | |
| 467 | } | ||
| 468 | ✗ | constexpr int64_t kMaxDumpBytes = 64 * 1024 * 1024; | |
| 469 | ✗ | if (n > kMaxDumpBytes) { | |
| 470 | ✗ | LOG(ERROR) << "DUMP_FAKE_DATA: n exceeds cap " << kMaxDumpBytes; | |
| 471 | ✗ | cntl->SetFailed(EINVAL, "DUMP_FAKE_DATA n exceeds cap"); | |
| 472 | ✗ | return; | |
| 473 | } | ||
| 474 | response->set_reply("ok"); | ||
| 475 | } else { | ||
| 476 | ✗ | LOG(FATAL) << "invalid command"; | |
| 477 | } | ||
| 478 | ✗ | } | |
| 479 | |||
| 480 | ✗ | void BRPCParameterServiceImpl::PutParameter( | |
| 481 | google::protobuf::RpcController* controller, | ||
| 482 | const PutParameterRequest* request, | ||
| 483 | PutParameterResponse* response, | ||
| 484 | google::protobuf::Closure* done) { | ||
| 485 | ✗ | brpc::ClosureGuard done_guard(done); | |
| 486 | |||
| 487 | ✗ | brpc::Controller* cntl = static_cast<brpc::Controller*>(controller); | |
| 488 | |||
| 489 | ✗ | std::string payload_storage; | |
| 490 | ✗ | const char* payload_data = nullptr; | |
| 491 | ✗ | int payload_size = 0; | |
| 492 | ✗ | if (!ExtractPayloadBytes( | |
| 493 | cntl, | ||
| 494 | request->parameter_value(), | ||
| 495 | &payload_storage, | ||
| 496 | &payload_data, | ||
| 497 | &payload_size)) { | ||
| 498 | ✗ | LOG(ERROR) << "PutParameter empty payload"; | |
| 499 | ✗ | return; | |
| 500 | } | ||
| 501 | |||
| 502 | #ifdef ENABLE_PERF_REPORT | ||
| 503 | auto start_time = std::chrono::high_resolution_clock::now(); | ||
| 504 | #endif | ||
| 505 | |||
| 506 | ✗ | const ParameterCompressReader* reader = | |
| 507 | reinterpret_cast<const ParameterCompressReader*>(payload_data); | ||
| 508 | ✗ | if (!reader->Valid(payload_size)) { | |
| 509 | ✗ | LOG(ERROR) << "PutParameter invalid payload, size=" << payload_size; | |
| 510 | ✗ | return; | |
| 511 | } | ||
| 512 | ✗ | int size = reader->item_size(); | |
| 513 | ✗ | uint64_t total_bytes = 0; | |
| 514 | |||
| 515 | ✗ | for (int i = 0; i < size; i++) { | |
| 516 | ✗ | cache_ps_->PutSingleParameter(reader->item(i), 0); | |
| 517 | ✗ | total_bytes += reader->item(i)->dim * sizeof(float); | |
| 518 | } | ||
| 519 | |||
| 520 | ✗ | total_put_requests_++; | |
| 521 | ✗ | total_put_keys_ += size; | |
| 522 | ✗ | total_put_bytes_ += total_bytes; | |
| 523 | |||
| 524 | #ifdef ENABLE_PERF_REPORT | ||
| 525 | auto end_time = std::chrono::high_resolution_clock::now(); | ||
| 526 | auto duration = | ||
| 527 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 528 | end_time - start_time) | ||
| 529 | .count(); | ||
| 530 | report("ps_server_latency", | ||
| 531 | "PutParameter", | ||
| 532 | "latency_us", | ||
| 533 | static_cast<double>(duration)); | ||
| 534 | |||
| 535 | double start_us_for_key = | ||
| 536 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 537 | start_time.time_since_epoch()) | ||
| 538 | .count(); | ||
| 539 | std::string op_latency_key = | ||
| 540 | "EmbWrite|" + std::to_string(static_cast<uint64_t>(start_us_for_key)); | ||
| 541 | report("op_latency", | ||
| 542 | op_latency_key.c_str(), | ||
| 543 | "recserver_us", | ||
| 544 | static_cast<double>(duration)); | ||
| 545 | #endif | ||
| 546 | ✗ | } | |
| 547 | |||
| 548 | ✗ | void BRPCParameterServiceImpl::UpdateParameter( | |
| 549 | google::protobuf::RpcController* controller, | ||
| 550 | const UpdateParameterRequest* request, | ||
| 551 | UpdateParameterResponse* reply, | ||
| 552 | google::protobuf::Closure* done) { | ||
| 553 | ✗ | brpc::ClosureGuard done_guard(done); | |
| 554 | |||
| 555 | ✗ | brpc::Controller* cntl = static_cast<brpc::Controller*>(controller); | |
| 556 | |||
| 557 | #ifdef ENABLE_PERF_REPORT | ||
| 558 | auto start_time = std::chrono::high_resolution_clock::now(); | ||
| 559 | uint64_t trace_id = 0; | ||
| 560 | const std::string* header_trace = | ||
| 561 | cntl->http_request().GetHeader("x-recstore-trace-id"); | ||
| 562 | if (header_trace != nullptr && !header_trace->empty()) { | ||
| 563 | trace_id = static_cast<uint64_t>( | ||
| 564 | std::strtoull(header_trace->c_str(), nullptr, 10)); | ||
| 565 | } | ||
| 566 | #endif | ||
| 567 | ✗ | bool success = false; | |
| 568 | ✗ | int size = 0; | |
| 569 | #ifdef ENABLE_PERF_REPORT | ||
| 570 | auto before_cache_update_time = std::chrono::high_resolution_clock::now(); | ||
| 571 | #endif | ||
| 572 | |||
| 573 | try { | ||
| 574 | ✗ | const std::string& table_name = request->table_name(); | |
| 575 | |||
| 576 | ✗ | std::string payload_storage; | |
| 577 | ✗ | const char* payload_data = nullptr; | |
| 578 | ✗ | int payload_size = 0; | |
| 579 | ✗ | if (!ExtractPayloadBytes( | |
| 580 | cntl, | ||
| 581 | request->gradients(), | ||
| 582 | &payload_storage, | ||
| 583 | &payload_data, | ||
| 584 | &payload_size)) { | ||
| 585 | ✗ | throw std::runtime_error("UpdateParameter empty gradients payload"); | |
| 586 | } | ||
| 587 | |||
| 588 | ✗ | const ParameterCompressReader* reader = | |
| 589 | reinterpret_cast<const ParameterCompressReader*>(payload_data); | ||
| 590 | ✗ | if (!reader->Valid(payload_size)) { | |
| 591 | ✗ | throw std::runtime_error("UpdateParameter invalid gradients payload"); | |
| 592 | } | ||
| 593 | ✗ | size = reader->item_size(); | |
| 594 | |||
| 595 | #ifdef ENABLE_PERF_REPORT | ||
| 596 | before_cache_update_time = std::chrono::high_resolution_clock::now(); | ||
| 597 | #endif | ||
| 598 | ✗ | success = cache_ps_->UpdateParameter(table_name, reader, 0); | |
| 599 | |||
| 600 | ✗ | RECSTORE_LOG_EVERY_MS(INFO, 2000) | |
| 601 | ✗ | << "UpdateParameter: table=" << table_name << ", keys=" << size; | |
| 602 | |||
| 603 | ✗ | reply->set_success(success); | |
| 604 | ✗ | } catch (const std::exception& e) { | |
| 605 | ✗ | LOG(ERROR) << "UpdateParameter error: " << e.what(); | |
| 606 | ✗ | reply->set_success(false); | |
| 607 | ✗ | } | |
| 608 | |||
| 609 | #ifdef ENABLE_PERF_REPORT | ||
| 610 | auto end_time = std::chrono::high_resolution_clock::now(); | ||
| 611 | auto duration = | ||
| 612 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 613 | end_time - start_time) | ||
| 614 | .count(); | ||
| 615 | report("ps_server_latency", | ||
| 616 | "UpdateParameter", | ||
| 617 | "latency_us", | ||
| 618 | static_cast<double>(duration)); | ||
| 619 | |||
| 620 | double start_us_for_key = | ||
| 621 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 622 | start_time.time_since_epoch()) | ||
| 623 | .count(); | ||
| 624 | std::string op_latency_key = | ||
| 625 | "EmbUpdate|" + std::to_string(static_cast<uint64_t>(start_us_for_key)); | ||
| 626 | report("op_latency", | ||
| 627 | op_latency_key.c_str(), | ||
| 628 | "recserver_us", | ||
| 629 | static_cast<double>(duration)); | ||
| 630 | |||
| 631 | auto backend_update_duration = | ||
| 632 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 633 | end_time - before_cache_update_time) | ||
| 634 | .count(); | ||
| 635 | const uint64_t effective_trace_id = | ||
| 636 | trace_id == 0 ? static_cast<uint64_t>(start_us_for_key) : trace_id; | ||
| 637 | std::string update_stage_id = | ||
| 638 | "brpc_server::EmbUpdate|" + std::to_string(effective_trace_id); | ||
| 639 | report("embupdate_stages", | ||
| 640 | update_stage_id.c_str(), | ||
| 641 | "server_total_us", | ||
| 642 | static_cast<double>(duration)); | ||
| 643 | report("embupdate_stages", | ||
| 644 | update_stage_id.c_str(), | ||
| 645 | "server_backend_update_us", | ||
| 646 | static_cast<double>(backend_update_duration)); | ||
| 647 | report("embupdate_stages", | ||
| 648 | update_stage_id.c_str(), | ||
| 649 | "server_request_size", | ||
| 650 | static_cast<double>(size)); | ||
| 651 | report("embupdate_stages", | ||
| 652 | update_stage_id.c_str(), | ||
| 653 | "server_success", | ||
| 654 | success ? 1.0 : 0.0); | ||
| 655 | #endif | ||
| 656 | ✗ | } | |
| 657 | |||
| 658 | ✗ | void BRPCParameterServiceImpl::InitEmbeddingTable( | |
| 659 | google::protobuf::RpcController* controller, | ||
| 660 | const InitEmbeddingTableRequest* request, | ||
| 661 | InitEmbeddingTableResponse* reply, | ||
| 662 | google::protobuf::Closure* done) { | ||
| 663 | ✗ | brpc::ClosureGuard done_guard(done); | |
| 664 | |||
| 665 | #ifdef ENABLE_PERF_REPORT | ||
| 666 | auto start_time = std::chrono::high_resolution_clock::now(); | ||
| 667 | #endif | ||
| 668 | |||
| 669 | try { | ||
| 670 | ✗ | if (request->has_config_payload()) { | |
| 671 | ✗ | auto payload = request->config_payload(); | |
| 672 | ✗ | nlohmann::json cfg = nlohmann::json::parse(payload); | |
| 673 | ✗ | uint64_t num_embeddings = cfg.value("num_embeddings", 0); | |
| 674 | ✗ | uint64_t embedding_dim = cfg.value("embedding_dim", 0); | |
| 675 | ✗ | RECSTORE_LOG_EVERY_MS(INFO, 2000) | |
| 676 | ✗ | << "InitEmbeddingTable: table=" << request->table_name() | |
| 677 | ✗ | << ", num_embeddings=" << num_embeddings | |
| 678 | ✗ | << ", embedding_dim=" << embedding_dim; | |
| 679 | |||
| 680 | ✗ | bool init_success = cache_ps_->InitTable( | |
| 681 | request->table_name(), num_embeddings, embedding_dim); | ||
| 682 | ✗ | reply->set_success(init_success); | |
| 683 | ✗ | } else { | |
| 684 | ✗ | LOG(WARNING) << "InitEmbeddingTable called without config_payload"; | |
| 685 | ✗ | reply->set_success(false); | |
| 686 | } | ||
| 687 | ✗ | } catch (const std::exception& e) { | |
| 688 | ✗ | LOG(ERROR) << "InitEmbeddingTable error: " << e.what(); | |
| 689 | ✗ | reply->set_success(false); | |
| 690 | ✗ | } | |
| 691 | |||
| 692 | #ifdef ENABLE_PERF_REPORT | ||
| 693 | auto end_time = std::chrono::high_resolution_clock::now(); | ||
| 694 | auto duration = | ||
| 695 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 696 | end_time - start_time) | ||
| 697 | .count(); | ||
| 698 | report("ps_server_latency", | ||
| 699 | "InitEmbeddingTable", | ||
| 700 | "latency_us", | ||
| 701 | static_cast<double>(duration)); | ||
| 702 | |||
| 703 | double start_us_for_key = | ||
| 704 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 705 | start_time.time_since_epoch()) | ||
| 706 | .count(); | ||
| 707 | std::string op_latency_key = | ||
| 708 | "InitEmbeddingTable|" + | ||
| 709 | std::to_string(static_cast<uint64_t>(start_us_for_key)); | ||
| 710 | report("op_latency", | ||
| 711 | op_latency_key.c_str(), | ||
| 712 | "recserver_us", | ||
| 713 | static_cast<double>(duration)); | ||
| 714 | #endif | ||
| 715 | ✗ | } | |
| 716 | |||
| 717 | class BRPCParameterServer : public BaseParameterServer { | ||
| 718 | public: | ||
| 719 | ✗ | BRPCParameterServer() = default; | |
| 720 | |||
| 721 | ✗ | void Run() { | |
| 722 | // Check whether multi-shard mode is configured | ||
| 723 | ✗ | int num_shards = 1; // default: single shard | |
| 724 | ✗ | if (config_["cache_ps"].contains("num_shards")) { | |
| 725 | ✗ | num_shards = config_["cache_ps"]["num_shards"]; | |
| 726 | } | ||
| 727 | const std::optional<int> local_shard_id = | ||
| 728 | ✗ | FLAGS_local_shard_id >= 0 | |
| 729 | ✗ | ? std::make_optional(FLAGS_local_shard_id) | |
| 730 | ✗ | : std::nullopt; | |
| 731 | |||
| 732 | ✗ | if (num_shards > 1) { | |
| 733 | // Multi-server startup | ||
| 734 | std::cout | ||
| 735 | ✗ | << "Starting distributed parameter server (bRPC), number of shards: " | |
| 736 | ✗ | << num_shards << std::endl; | |
| 737 | |||
| 738 | ✗ | if (!config_["cache_ps"].contains("servers")) { | |
| 739 | ✗ | LOG(FATAL) << "num_shards > 1 but cache_ps.servers is missing"; | |
| 740 | ✗ | return; | |
| 741 | } | ||
| 742 | |||
| 743 | ✗ | const auto& cache_ps_config = config_["cache_ps"]; | |
| 744 | auto servers = | ||
| 745 | ✗ | SelectShardConfigsInternal(cache_ps_config, local_shard_id); | |
| 746 | ✗ | const auto configured_servers = cache_ps_config["servers"]; | |
| 747 | ✗ | if (configured_servers.size() != num_shards) { | |
| 748 | ✗ | LOG(FATAL) << "servers 配置数量 (" << configured_servers.size() | |
| 749 | ✗ | << ") 与 num_shards (" << num_shards << ") 不匹配"; | |
| 750 | return; | ||
| 751 | } | ||
| 752 | ✗ | if (local_shard_id.has_value() && servers.empty()) { | |
| 753 | ✗ | LOG(FATAL) << "local_shard_id=" << *local_shard_id | |
| 754 | ✗ | << " is not present in cache_ps.servers"; | |
| 755 | return; | ||
| 756 | } | ||
| 757 | ✗ | if (!local_shard_id.has_value() && | |
| 758 | ✗ | servers.size() != configured_servers.size()) { | |
| 759 | ✗ | LOG(FATAL) << "Selected shard count (" << servers.size() | |
| 760 | ✗ | << ") does not match configured server count (" | |
| 761 | ✗ | << configured_servers.size() << ")"; | |
| 762 | return; | ||
| 763 | } | ||
| 764 | |||
| 765 | ✗ | std::vector<std::thread> server_threads; | |
| 766 | |||
| 767 | ✗ | for (auto& server_config : servers) { | |
| 768 | ✗ | server_threads.emplace_back([this, server_config]() { | |
| 769 | ✗ | std::string host = server_config["host"]; | |
| 770 | ✗ | int port = server_config["port"]; | |
| 771 | ✗ | int shard = server_config["shard"]; | |
| 772 | |||
| 773 | ✗ | std::string server_address = host + ":" + std::to_string(port); | |
| 774 | |||
| 775 | ✗ | nlohmann::json shard_config = config_["cache_ps"]; | |
| 776 | ✗ | shard_config["num_shards"] = 1; | |
| 777 | ✗ | shard_config["servers"] = nlohmann::json::array({server_config}); | |
| 778 | ✗ | if (shard_config.contains("base_kv_config") && | |
| 779 | ✗ | shard_config["base_kv_config"].is_object()) { | |
| 780 | ✗ | auto& base_kv_config = shard_config["base_kv_config"]; | |
| 781 | ✗ | AppendShardSuffixIfPresent(base_kv_config, "path", shard); | |
| 782 | ✗ | AppendShardSuffixIfPresent(base_kv_config, "rocksdb_path", shard); | |
| 783 | ✗ | AppendShardSuffixToNestedFilePaths(base_kv_config, shard); | |
| 784 | ✗ | LOG(INFO) << "bRPC shard " << shard | |
| 785 | ✗ | << " using base_kv_config: " << base_kv_config.dump(); | |
| 786 | } | ||
| 787 | |||
| 788 | ✗ | auto cache_ps = std::make_unique<CachePS>(shard_config); | |
| 789 | auto service = | ||
| 790 | ✗ | std::make_unique<BRPCParameterServiceImpl>(cache_ps.get()); | |
| 791 | |||
| 792 | ✗ | brpc::Server server; | |
| 793 | ✗ | brpc::ServerOptions options; | |
| 794 | ✗ | options.num_threads = FLAGS_brpc_server_num_threads; | |
| 795 | #if BRPC_WITH_RDMA | ||
| 796 | options.use_rdma = ResolveBrpcServerUseRdmaFromEnv( | ||
| 797 | FLAGS_brpc_ps_use_rdma); | ||
| 798 | if (options.use_rdma) { | ||
| 799 | ApplyBrpcServerRdmaDeviceFromEnv(); | ||
| 800 | } | ||
| 801 | #endif | ||
| 802 | |||
| 803 | ✗ | if (server.AddService( | |
| 804 | ✗ | service.get(), brpc::SERVER_DOESNT_OWN_SERVICE) != 0) { | |
| 805 | ✗ | LOG(ERROR) << "Failed to add service!"; | |
| 806 | ✗ | return; | |
| 807 | } | ||
| 808 | |||
| 809 | ✗ | if (server.Start(server_address.c_str(), &options) != 0) { | |
| 810 | ✗ | LOG(ERROR) << "Failed to start bRPC server at " << server_address; | |
| 811 | ✗ | return; | |
| 812 | } | ||
| 813 | |||
| 814 | ✗ | std::cout << "bRPC Server shard " << shard << " listening on " | |
| 815 | ✗ | << server_address << std::endl; | |
| 816 | ✗ | server.RunUntilAskedToQuit(); | |
| 817 | ✗ | }); | |
| 818 | } | ||
| 819 | |||
| 820 | // Wait for all server threads | ||
| 821 | ✗ | for (auto& t : server_threads) { | |
| 822 | ✗ | t.join(); | |
| 823 | } | ||
| 824 | ✗ | } else { | |
| 825 | // Single-server startup | ||
| 826 | ✗ | std::cout << "Starting single parameter server (bRPC)" << std::endl; | |
| 827 | std::string server_address = | ||
| 828 | ✗ | "0.0.0.0:" + std::to_string(FLAGS_brpc_server_port); | |
| 829 | ✗ | auto cache_ps = std::make_unique<CachePS>(config_["cache_ps"]); | |
| 830 | ✗ | auto service = std::make_unique<BRPCParameterServiceImpl>(cache_ps.get()); | |
| 831 | |||
| 832 | ✗ | std::atomic<bool> metrics_running{true}; | |
| 833 | ✗ | std::thread metrics_thread([&service, &metrics_running]() { | |
| 834 | ✗ | while (metrics_running) { | |
| 835 | ✗ | std::this_thread::sleep_for(std::chrono::seconds(10)); | |
| 836 | ✗ | service->PrintMetrics(); | |
| 837 | ✗ | service->ResetMetrics(); | |
| 838 | } | ||
| 839 | ✗ | }); | |
| 840 | |||
| 841 | ✗ | brpc::Server server; | |
| 842 | ✗ | brpc::ServerOptions options; | |
| 843 | ✗ | options.num_threads = FLAGS_brpc_server_num_threads; | |
| 844 | #if BRPC_WITH_RDMA | ||
| 845 | options.use_rdma = ResolveBrpcServerUseRdmaFromEnv(FLAGS_brpc_ps_use_rdma); | ||
| 846 | if (options.use_rdma) { | ||
| 847 | ApplyBrpcServerRdmaDeviceFromEnv(); | ||
| 848 | } | ||
| 849 | #endif | ||
| 850 | |||
| 851 | ✗ | if (server.AddService(service.get(), brpc::SERVER_DOESNT_OWN_SERVICE) != | |
| 852 | 0) { | ||
| 853 | ✗ | LOG(ERROR) << "Failed to add service!"; | |
| 854 | ✗ | metrics_running = false; | |
| 855 | ✗ | if (metrics_thread.joinable()) { | |
| 856 | ✗ | metrics_thread.join(); | |
| 857 | } | ||
| 858 | ✗ | return; | |
| 859 | } | ||
| 860 | |||
| 861 | ✗ | if (server.Start(server_address.c_str(), &options) != 0) { | |
| 862 | ✗ | LOG(ERROR) << "Failed to start bRPC server at " << server_address; | |
| 863 | ✗ | metrics_running = false; | |
| 864 | ✗ | if (metrics_thread.joinable()) { | |
| 865 | ✗ | metrics_thread.join(); | |
| 866 | } | ||
| 867 | ✗ | return; | |
| 868 | } | ||
| 869 | |||
| 870 | ✗ | std::cout << "bRPC Server listening on " << server_address << std::endl; | |
| 871 | ✗ | server.RunUntilAskedToQuit(); | |
| 872 | |||
| 873 | ✗ | metrics_running = false; | |
| 874 | ✗ | if (metrics_thread.joinable()) { | |
| 875 | ✗ | metrics_thread.join(); | |
| 876 | } | ||
| 877 | ✗ | } | |
| 878 | } | ||
| 879 | }; | ||
| 880 | |||
| 881 | FACTORY_REGISTER(BaseParameterServer, BRPCParameterServer, BRPCParameterServer); | ||
| 882 | |||
| 883 | } // namespace recstore | ||
| 884 | |||
| 885 | #ifndef RECSTORE_NO_SERVER_MAIN | ||
| 886 | ✗ | int main(int argc, char** argv) { | |
| 887 | ✗ | gflags::ParseCommandLineFlags(&argc, &argv, true); | |
| 888 | |||
| 889 | const std::string config_path = | ||
| 890 | ✗ | FLAGS_brpc_config_path.empty() | |
| 891 | ✗ | ? base::ResolveRecStoreConfigPath().string() | |
| 892 | ✗ | : FLAGS_brpc_config_path; | |
| 893 | ✗ | std::ifstream config_file(config_path); | |
| 894 | ✗ | if (!config_file.is_open()) { | |
| 895 | ✗ | throw std::runtime_error("Cannot open config file: " + config_path); | |
| 896 | } | ||
| 897 | ✗ | nlohmann::json ex; | |
| 898 | ✗ | config_file >> ex; | |
| 899 | |||
| 900 | ✗ | recstore::BRPCParameterServer ps; | |
| 901 | ✗ | std::cout << "bRPC Parameter server config: " << ex.dump(2) << std::endl; | |
| 902 | ✗ | ps.Init(ex); | |
| 903 | ✗ | ps.Run(); | |
| 904 | |||
| 905 | ✗ | return 0; | |
| 906 | ✗ | } | |
| 907 | #endif | ||
| 908 |