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