ps/grpc/grpc_ps_server.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #include <grpcpp/ext/proto_server_reflection_plugin.h> | ||
| 2 | #include <grpcpp/grpcpp.h> | ||
| 3 | #include <grpcpp/health_check_service_interface.h> | ||
| 4 | |||
| 5 | #include <cstdint> | ||
| 6 | #include <cstring> | ||
| 7 | #include <fstream> | ||
| 8 | #include <future> | ||
| 9 | #include <optional> | ||
| 10 | #include <stdexcept> | ||
| 11 | #include <string> | ||
| 12 | #include <thread> | ||
| 13 | #include <vector> | ||
| 14 | |||
| 15 | #include "base/array.h" | ||
| 16 | #include "base/base.h" | ||
| 17 | #include "base/flatc.h" | ||
| 18 | #include "base/init.h" | ||
| 19 | #include "base/timer.h" | ||
| 20 | #include "ps.grpc.pb.h" | ||
| 21 | #include "ps.pb.h" | ||
| 22 | #include "ps/base/base_ps_server.h" | ||
| 23 | #include "ps/base/cache_ps_impl.h" | ||
| 24 | #include "ps/base/parameters.h" | ||
| 25 | #include "src/base/config.h" | ||
| 26 | |||
| 27 | #ifdef ENABLE_PERF_REPORT | ||
| 28 | # include <chrono> | ||
| 29 | # include <cstdlib> | ||
| 30 | |||
| 31 | # include "base/report/report_client.h" | ||
| 32 | #else | ||
| 33 | # include "../report_client.h" | ||
| 34 | #endif | ||
| 35 | |||
| 36 | using grpc::Server; | ||
| 37 | using grpc::ServerBuilder; | ||
| 38 | using grpc::ServerContext; | ||
| 39 | using grpc::Status; | ||
| 40 | |||
| 41 | using recstoreps::CommandRequest; | ||
| 42 | using recstoreps::CommandResponse; | ||
| 43 | using recstoreps::GetParameterRequest; | ||
| 44 | using recstoreps::GetParameterResponse; | ||
| 45 | using recstoreps::InitEmbeddingTableRequest; | ||
| 46 | using recstoreps::InitEmbeddingTableResponse; | ||
| 47 | using recstoreps::PSCommand; | ||
| 48 | using recstoreps::PutParameterRequest; | ||
| 49 | using recstoreps::PutParameterResponse; | ||
| 50 | using recstoreps::UpdateParameterRequest; | ||
| 51 | using recstoreps::UpdateParameterResponse; | ||
| 52 | |||
| 53 | DEFINE_string(config_path, "", "config file path"); | ||
| 54 | DEFINE_int32(grpc_local_shard_id, | ||
| 55 | -1, | ||
| 56 | "Only start the specified shard in multi-shard gRPC mode; " | ||
| 57 | "-1 means start all configured shards"); | ||
| 58 | |||
| 59 | namespace { | ||
| 60 | |||
| 61 | ✗ | void AppendShardSuffixIfPresent( | |
| 62 | nlohmann::json& config_node, const char* key, int shard_id) { | ||
| 63 | ✗ | if (!config_node.contains(key) || !config_node[key].is_string()) { | |
| 64 | ✗ | return; | |
| 65 | } | ||
| 66 | ✗ | config_node[key] = | |
| 67 | ✗ | config_node[key].get<std::string>() + "_" + std::to_string(shard_id); | |
| 68 | } | ||
| 69 | |||
| 70 | ✗ | void AppendShardSuffixToNestedFilePaths(nlohmann::json& node, int shard_id) { | |
| 71 | ✗ | if (node.is_object()) { | |
| 72 | ✗ | for (auto& item : node.items()) { | |
| 73 | ✗ | if (item.key() == "file_path" && item.value().is_string()) { | |
| 74 | ✗ | item.value() = | |
| 75 | ✗ | item.value().get<std::string>() + "_" + std::to_string(shard_id); | |
| 76 | ✗ | continue; | |
| 77 | } | ||
| 78 | ✗ | AppendShardSuffixToNestedFilePaths(item.value(), shard_id); | |
| 79 | ✗ | } | |
| 80 | ✗ | return; | |
| 81 | } | ||
| 82 | ✗ | if (node.is_array()) { | |
| 83 | ✗ | for (auto& item : node) { | |
| 84 | ✗ | AppendShardSuffixToNestedFilePaths(item, shard_id); | |
| 85 | } | ||
| 86 | } | ||
| 87 | } | ||
| 88 | |||
| 89 | std::vector<nlohmann::json> | ||
| 90 | ✗ | SelectGRPCShardConfigs(const nlohmann::json& cache_ps_config, | |
| 91 | const std::optional<int>& local_shard_id) { | ||
| 92 | ✗ | std::vector<nlohmann::json> selected; | |
| 93 | ✗ | if (!cache_ps_config.contains("servers") || | |
| 94 | ✗ | !cache_ps_config["servers"].is_array()) { | |
| 95 | ✗ | return selected; | |
| 96 | } | ||
| 97 | |||
| 98 | ✗ | for (const auto& server_config : cache_ps_config["servers"]) { | |
| 99 | ✗ | if (!local_shard_id.has_value()) { | |
| 100 | ✗ | selected.push_back(server_config); | |
| 101 | ✗ | continue; | |
| 102 | } | ||
| 103 | ✗ | if (!server_config.contains("shard") || | |
| 104 | ✗ | !server_config["shard"].is_number_integer()) { | |
| 105 | ✗ | continue; | |
| 106 | } | ||
| 107 | ✗ | if (server_config["shard"].get<int>() == *local_shard_id) { | |
| 108 | ✗ | selected.push_back(server_config); | |
| 109 | } | ||
| 110 | } | ||
| 111 | ✗ | return selected; | |
| 112 | ✗ | } | |
| 113 | |||
| 114 | } // namespace | ||
| 115 | |||
| 116 | class ParameterServiceImpl final | ||
| 117 | : public recstoreps::ParameterService::Service { | ||
| 118 | public: | ||
| 119 | ✗ | ParameterServiceImpl(CachePS* cache_ps, int shard_id) { | |
| 120 | ✗ | cache_ps_ = cache_ps; | |
| 121 | ✗ | shard_id_ = shard_id; | |
| 122 | ✗ | start_time_ = std::chrono::steady_clock::now(); | |
| 123 | ✗ | } | |
| 124 | ✗ | void ResetMetrics() { | |
| 125 | ✗ | total_get_requests_ = 0; | |
| 126 | ✗ | total_put_requests_ = 0; | |
| 127 | ✗ | total_get_keys_ = 0; | |
| 128 | ✗ | total_put_keys_ = 0; | |
| 129 | ✗ | total_get_bytes_ = 0; | |
| 130 | ✗ | total_put_bytes_ = 0; | |
| 131 | ✗ | start_time_ = std::chrono::steady_clock::now(); | |
| 132 | ✗ | } | |
| 133 | ✗ | void PrintMetrics(const std::string& table_name = "grpc_ps_server_metrics", | |
| 134 | const std::string& unique_id = "default_server") { | ||
| 135 | ✗ | auto now = std::chrono::steady_clock::now(); | |
| 136 | ✗ | double elapsed_s = std::chrono::duration<double>(now - start_time_).count(); | |
| 137 | ✗ | if (elapsed_s > 0) { | |
| 138 | double overall_qps = | ||
| 139 | ✗ | (total_get_requests_ + total_put_requests_) / elapsed_s; | |
| 140 | double overall_throughput_mbps = | ||
| 141 | ✗ | ((total_get_bytes_ + total_put_bytes_) / 1024.0 / 1024.0) / elapsed_s; | |
| 142 | |||
| 143 | // Report QPS and throughput metrics | ||
| 144 | |||
| 145 | // report(table_name.c_str(), unique_id.c_str(), "overall_qps", | ||
| 146 | // overall_qps); report(table_name.c_str(), | ||
| 147 | // unique_id.c_str(), | ||
| 148 | // "overall_throughput_mbps", | ||
| 149 | // overall_throughput_mbps); | ||
| 150 | } | ||
| 151 | ✗ | } | |
| 152 | |||
| 153 | private: | ||
| 154 | ✗ | Status GetParameter(ServerContext* context, | |
| 155 | const GetParameterRequest* request, | ||
| 156 | GetParameterResponse* reply) override { | ||
| 157 | #ifdef ENABLE_PERF_REPORT | ||
| 158 | auto start_time = std::chrono::high_resolution_clock::now(); | ||
| 159 | #endif | ||
| 160 | ✗ | base::ConstArray<uint64_t> keys_array(request->keys()); | |
| 161 | ✗ | bool isPerf = request->has_perf() && request->perf(); | |
| 162 | ✗ | if (isPerf) { | |
| 163 | ✗ | xmh::PerfCounter::Record("PS Get Keys", keys_array.Size()); | |
| 164 | } | ||
| 165 | ✗ | xmh::Timer timer_ps_get_req("PS GetParameter Req"); | |
| 166 | ✗ | ParameterCompressor compressor(std::numeric_limits<int>::max()); | |
| 167 | ✗ | std::vector<std::string> blocks; | |
| 168 | ✗ | RECSTORE_LOG_EVERY_MS(INFO, 1000) | |
| 169 | ✗ | << "[PS] Getting " << keys_array.Size() << " keys"; | |
| 170 | ✗ | int total_dim = 0; | |
| 171 | #ifdef ENABLE_PERF_REPORT | ||
| 172 | auto cache_start_time = std::chrono::high_resolution_clock::now(); | ||
| 173 | #endif | ||
| 174 | ✗ | std::vector<ParameterPack> packs; | |
| 175 | ✗ | packs.reserve(keys_array.Size()); | |
| 176 | ✗ | cache_ps_->GetParameterRun2Completion(keys_array, packs, 0); | |
| 177 | |||
| 178 | ✗ | for (auto& pack : packs) { | |
| 179 | ✗ | compressor.AddItem(pack, &blocks); | |
| 180 | ✗ | total_dim += pack.dim; | |
| 181 | } | ||
| 182 | #ifdef ENABLE_PERF_REPORT | ||
| 183 | auto cache_end_time = std::chrono::high_resolution_clock::now(); | ||
| 184 | auto cache_duration = | ||
| 185 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 186 | cache_end_time - cache_start_time) | ||
| 187 | .count(); | ||
| 188 | double start_us_for_cache = | ||
| 189 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 190 | start_time.time_since_epoch()) | ||
| 191 | .count(); | ||
| 192 | std::string report_id_for_cache = | ||
| 193 | "grpc_server::GetParameter|" + | ||
| 194 | std::to_string(static_cast<uint64_t>(start_us_for_cache)); | ||
| 195 | report("embread_stages", | ||
| 196 | report_id_for_cache.c_str(), | ||
| 197 | "cache_lookup_us", | ||
| 198 | static_cast<double>(cache_duration)); | ||
| 199 | #endif | ||
| 200 | |||
| 201 | ✗ | compressor.ToBlock(&blocks); | |
| 202 | ✗ | CHECK_EQ(blocks.size(), 1); | |
| 203 | ✗ | reply->mutable_parameter_value()->swap(blocks[0]); | |
| 204 | ✗ | total_get_requests_++; | |
| 205 | ✗ | total_get_keys_ += keys_array.Size(); | |
| 206 | ✗ | total_get_bytes_ += total_dim * sizeof(float); | |
| 207 | |||
| 208 | ✗ | if (isPerf) { | |
| 209 | ✗ | timer_ps_get_req.end(); | |
| 210 | } else { | ||
| 211 | ✗ | timer_ps_get_req.destroy(); | |
| 212 | } | ||
| 213 | |||
| 214 | #ifdef ENABLE_PERF_REPORT | ||
| 215 | auto end_time = std::chrono::high_resolution_clock::now(); | ||
| 216 | double start_us = | ||
| 217 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 218 | start_time.time_since_epoch()) | ||
| 219 | .count(); | ||
| 220 | auto duration = | ||
| 221 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 222 | end_time - start_time) | ||
| 223 | .count(); | ||
| 224 | |||
| 225 | std::string report_id = "grpc_server::GetParameter|" + | ||
| 226 | std::to_string(static_cast<uint64_t>(start_us)); | ||
| 227 | |||
| 228 | std::string op_latency_key = | ||
| 229 | "EmbRead|" + std::to_string(static_cast<uint64_t>(start_us)); | ||
| 230 | report("op_latency", | ||
| 231 | op_latency_key.c_str(), | ||
| 232 | "recserver_us", | ||
| 233 | static_cast<double>(duration)); | ||
| 234 | |||
| 235 | report("embread_stages", | ||
| 236 | report_id.c_str(), | ||
| 237 | "duration_us", | ||
| 238 | static_cast<double>(duration)); | ||
| 239 | |||
| 240 | report("embread_stages", | ||
| 241 | report_id.c_str(), | ||
| 242 | "request_size", | ||
| 243 | static_cast<double>(keys_array.Size())); | ||
| 244 | |||
| 245 | std::string unique_id = | ||
| 246 | "embread_debug|" + std::to_string(static_cast<uint64_t>(start_us)); | ||
| 247 | FlameGraphData grpc_server_data = { | ||
| 248 | "grpc_ps_server::GetParameter", | ||
| 249 | start_us, | ||
| 250 | 2, // level | ||
| 251 | static_cast<double>(duration), | ||
| 252 | static_cast<double>(duration)}; | ||
| 253 | report_flame_graph( | ||
| 254 | "emb_read_flame_map", unique_id.c_str(), grpc_server_data); | ||
| 255 | #endif | ||
| 256 | |||
| 257 | ✗ | return Status::OK; | |
| 258 | ✗ | } | |
| 259 | |||
| 260 | ✗ | Status Command(ServerContext* context, | |
| 261 | const CommandRequest* request, | ||
| 262 | CommandResponse* reply) override { | ||
| 263 | ✗ | if (request->command() == PSCommand::CLEAR_PS) { | |
| 264 | ✗ | LOG(WARNING) << "[PS Command] Clear All"; | |
| 265 | ✗ | cache_ps_->Clear(); | |
| 266 | ✗ | } else if (request->command() == PSCommand::RELOAD_PS) { | |
| 267 | ✗ | LOG(WARNING) << "[PS Command] Reload PS"; | |
| 268 | ✗ | CHECK_NE(request->arg1().size(), 0); | |
| 269 | ✗ | CHECK_NE(request->arg2().size(), 0); | |
| 270 | ✗ | CHECK_EQ(request->arg1().size(), 1); | |
| 271 | ✗ | LOG(WARNING) << "model_config_path = " << request->arg1()[0]; | |
| 272 | ✗ | for (int i = 0; i < request->arg2().size(); i++) { | |
| 273 | ✗ | LOG(WARNING) << fmt::format("emb_file {}: {}", i, request->arg2()[i]); | |
| 274 | } | ||
| 275 | ✗ | std::vector<std::string> arg1; | |
| 276 | ✗ | for (auto& each : request->arg1()) { | |
| 277 | ✗ | arg1.push_back(each); | |
| 278 | } | ||
| 279 | ✗ | std::vector<std::string> arg2; | |
| 280 | ✗ | for (auto& each : request->arg2()) { | |
| 281 | ✗ | arg2.push_back(each); | |
| 282 | } | ||
| 283 | |||
| 284 | ✗ | cache_ps_->Initialize(arg1, arg2); | |
| 285 | ✗ | } else if (request->command() == PSCommand::SAVE_CHECKPOINT || | |
| 286 | ✗ | request->command() == PSCommand::LOAD_CHECKPOINT) { | |
| 287 | ✗ | if (request->arg1_size() != 1 || request->arg2_size() != 1 || | |
| 288 | ✗ | request->arg1(0).empty() || request->arg2(0).empty()) { | |
| 289 | return Status(grpc::StatusCode::INVALID_ARGUMENT, | ||
| 290 | ✗ | "checkpoint command requires path and metadata"); | |
| 291 | } | ||
| 292 | try { | ||
| 293 | ✗ | nlohmann::json metadata = nlohmann::json::parse(request->arg2(0)); | |
| 294 | ✗ | if (!metadata.is_object() || !metadata.contains("identity") || | |
| 295 | ✗ | !metadata["identity"].is_object() || | |
| 296 | ✗ | !metadata.contains("checkpoint_id") || | |
| 297 | ✗ | !metadata["checkpoint_id"].is_string()) { | |
| 298 | return Status(grpc::StatusCode::INVALID_ARGUMENT, | ||
| 299 | ✗ | "invalid checkpoint metadata"); | |
| 300 | } | ||
| 301 | ✗ | metadata["shard_id"] = shard_id_; | |
| 302 | ✗ | const std::string shard_metadata = metadata.dump(); | |
| 303 | ✗ | const bool save = request->command() == PSCommand::SAVE_CHECKPOINT; | |
| 304 | const bool ok = | ||
| 305 | ✗ | save ? cache_ps_->SaveCheckpoint(request->arg1(0), shard_metadata) | |
| 306 | ✗ | : cache_ps_->LoadCheckpoint(request->arg1(0), shard_metadata); | |
| 307 | ✗ | if (!ok) { | |
| 308 | return Status( | ||
| 309 | grpc::StatusCode::FAILED_PRECONDITION, | ||
| 310 | ✗ | save ? "checkpoint save failed" : "checkpoint load failed"); | |
| 311 | } | ||
| 312 | reply->set_reply("ok"); | ||
| 313 | ✗ | } catch (const nlohmann::json::exception& e) { | |
| 314 | ✗ | return Status(grpc::StatusCode::INVALID_ARGUMENT, e.what()); | |
| 315 | ✗ | } catch (const std::exception& e) { | |
| 316 | ✗ | return Status(grpc::StatusCode::FAILED_PRECONDITION, e.what()); | |
| 317 | ✗ | } | |
| 318 | ✗ | } else if (request->command() == PSCommand::LOAD_FAKE_DATA) { | |
| 319 | ✗ | if (request->arg1_size() != 1 || | |
| 320 | ✗ | static_cast<size_t>(request->arg1(0).size()) != sizeof(int64_t)) { | |
| 321 | ✗ | LOG(ERROR) << "LOAD_FAKE_DATA: arg1 must be one " << sizeof(int64_t) | |
| 322 | ✗ | << "-byte int64_t (requested reply payload size)"; | |
| 323 | return Status(grpc::StatusCode::INVALID_ARGUMENT, | ||
| 324 | ✗ | "LOAD_FAKE_DATA invalid arg1 size"); | |
| 325 | } | ||
| 326 | ✗ | int64_t payload_bytes = 0; | |
| 327 | ✗ | std::memcpy(&payload_bytes, request->arg1(0).data(), sizeof(int64_t)); | |
| 328 | ✗ | if (payload_bytes < 0) { | |
| 329 | ✗ | LOG(ERROR) << "LOAD_FAKE_DATA: payload_bytes must be non-negative, got " | |
| 330 | ✗ | << payload_bytes; | |
| 331 | return Status(grpc::StatusCode::INVALID_ARGUMENT, | ||
| 332 | ✗ | "payload_bytes must be non-negative"); | |
| 333 | } | ||
| 334 | ✗ | constexpr int64_t kMaxReplyPayload = 16 * 1024 * 1024; | |
| 335 | ✗ | if (payload_bytes > kMaxReplyPayload) { | |
| 336 | ✗ | LOG(ERROR) << "LOAD_FAKE_DATA: payload_bytes " << payload_bytes | |
| 337 | ✗ | << " exceeds cap " << kMaxReplyPayload; | |
| 338 | ✗ | return Status(grpc::StatusCode::INVALID_ARGUMENT, "payload too large"); | |
| 339 | } | ||
| 340 | ✗ | std::string fake(static_cast<size_t>(payload_bytes), '\xab'); | |
| 341 | ✗ | reply->set_reply(std::move(fake)); | |
| 342 | ✗ | } else if (request->command() == PSCommand::DUMP_FAKE_DATA) { | |
| 343 | ✗ | if (request->arg1_size() != 1 || | |
| 344 | ✗ | static_cast<size_t>(request->arg1(0).size()) != sizeof(int64_t)) { | |
| 345 | ✗ | LOG(ERROR) << "DUMP_FAKE_DATA: arg1 must be one " << sizeof(int64_t) | |
| 346 | ✗ | << "-byte int64_t (payload bytes n)"; | |
| 347 | return Status(grpc::StatusCode::INVALID_ARGUMENT, | ||
| 348 | ✗ | "DUMP_FAKE_DATA invalid arg1 size"); | |
| 349 | } | ||
| 350 | ✗ | int64_t n = 0; | |
| 351 | ✗ | std::memcpy(&n, request->arg1(0).data(), sizeof(int64_t)); | |
| 352 | ✗ | if (n <= 0) { | |
| 353 | ✗ | LOG(ERROR) << "DUMP_FAKE_DATA: n must be positive"; | |
| 354 | return Status(grpc::StatusCode::INVALID_ARGUMENT, | ||
| 355 | ✗ | "DUMP_FAKE_DATA n must be positive"); | |
| 356 | } | ||
| 357 | ✗ | if (n % static_cast<int64_t>(sizeof(float)) != 0) { | |
| 358 | ✗ | LOG(ERROR) << "DUMP_FAKE_DATA: n must be a multiple of " | |
| 359 | ✗ | << sizeof(float); | |
| 360 | return Status(grpc::StatusCode::INVALID_ARGUMENT, | ||
| 361 | ✗ | "DUMP_FAKE_DATA n must be multiple of sizeof(float)"); | |
| 362 | } | ||
| 363 | ✗ | constexpr int64_t kMaxDumpBytes = 64 * 1024 * 1024; | |
| 364 | ✗ | if (n > kMaxDumpBytes) { | |
| 365 | ✗ | LOG(ERROR) << "DUMP_FAKE_DATA: n exceeds cap " << kMaxDumpBytes; | |
| 366 | return Status( | ||
| 367 | ✗ | grpc::StatusCode::INVALID_ARGUMENT, "DUMP_FAKE_DATA n exceeds cap"); | |
| 368 | } | ||
| 369 | // Receive fake data payload (used for write bandwidth benchmarking) | ||
| 370 | reply->set_reply("ok"); | ||
| 371 | } else { | ||
| 372 | ✗ | LOG(FATAL) << "invalid command"; | |
| 373 | } | ||
| 374 | ✗ | return Status::OK; | |
| 375 | } | ||
| 376 | |||
| 377 | ✗ | Status PutParameter(ServerContext* context, | |
| 378 | const PutParameterRequest* request, | ||
| 379 | PutParameterResponse* reply) override { | ||
| 380 | #ifdef ENABLE_PERF_REPORT | ||
| 381 | auto start_time = std::chrono::high_resolution_clock::now(); | ||
| 382 | #endif | ||
| 383 | const ParameterCompressReader* reader = | ||
| 384 | reinterpret_cast<const ParameterCompressReader*>( | ||
| 385 | ✗ | request->parameter_value().data()); | |
| 386 | ✗ | int size = reader->item_size(); | |
| 387 | ✗ | LOG(INFO) << "[PS] PutParameter: " << size << " keys"; | |
| 388 | ✗ | uint64_t total_bytes = 0; | |
| 389 | |||
| 390 | ✗ | for (int i = 0; i < size; i++) { | |
| 391 | ✗ | total_bytes += reader->item(i)->dim * sizeof(float); | |
| 392 | } | ||
| 393 | ✗ | cache_ps_->PutParameter(reader, 0); | |
| 394 | ✗ | LOG(INFO) << "[PS] PutParameter done: " << size << " keys"; | |
| 395 | ✗ | total_put_requests_++; | |
| 396 | ✗ | total_put_keys_ += size; | |
| 397 | ✗ | total_put_bytes_ += total_bytes; | |
| 398 | |||
| 399 | #ifdef ENABLE_PERF_REPORT | ||
| 400 | auto end_time = std::chrono::high_resolution_clock::now(); | ||
| 401 | double start_us_for_key = | ||
| 402 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 403 | start_time.time_since_epoch()) | ||
| 404 | .count(); | ||
| 405 | auto duration = | ||
| 406 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 407 | end_time - start_time) | ||
| 408 | .count(); | ||
| 409 | std::string op_latency_key = | ||
| 410 | "EmbWrite|" + std::to_string(static_cast<uint64_t>(start_us_for_key)); | ||
| 411 | report("op_latency", | ||
| 412 | op_latency_key.c_str(), | ||
| 413 | "recserver_us", | ||
| 414 | static_cast<double>(duration)); | ||
| 415 | #endif | ||
| 416 | |||
| 417 | ✗ | return Status::OK; | |
| 418 | } | ||
| 419 | |||
| 420 | ✗ | Status SparseGradientCommit( | |
| 421 | ServerContext*, | ||
| 422 | const recstoreps::SparseGradientCommitRequest* request, | ||
| 423 | recstoreps::SparseGradientCommitResponse* reply) override { | ||
| 424 | try { | ||
| 425 | ✗ | reply->set_committed_step(cache_ps_->SparseGradientCommit( | |
| 426 | ✗ | request->table_name(), request->payload(), 0)); | |
| 427 | ✗ | } catch (const std::exception& error) { | |
| 428 | ✗ | reply->set_error(error.what()); | |
| 429 | ✗ | } | |
| 430 | ✗ | return Status::OK; | |
| 431 | } | ||
| 432 | |||
| 433 | ✗ | Status UpdateParameter(ServerContext* context, | |
| 434 | const UpdateParameterRequest* request, | ||
| 435 | UpdateParameterResponse* reply) override { | ||
| 436 | #ifdef ENABLE_PERF_REPORT | ||
| 437 | auto start_time = std::chrono::high_resolution_clock::now(); | ||
| 438 | uint64_t trace_id = 0; | ||
| 439 | const auto trace_it = | ||
| 440 | context->client_metadata().find("x-recstore-trace-id"); | ||
| 441 | if (trace_it != context->client_metadata().end()) { | ||
| 442 | std::string trace_id_str( | ||
| 443 | trace_it->second.data(), trace_it->second.length()); | ||
| 444 | trace_id = static_cast<uint64_t>( | ||
| 445 | std::strtoull(trace_id_str.c_str(), nullptr, 10)); | ||
| 446 | } | ||
| 447 | #endif | ||
| 448 | ✗ | bool success = false; | |
| 449 | ✗ | int size = 0; | |
| 450 | ✗ | std::string table_name; | |
| 451 | #ifdef ENABLE_PERF_REPORT | ||
| 452 | auto before_cache_update_time = std::chrono::high_resolution_clock::now(); | ||
| 453 | #endif | ||
| 454 | try { | ||
| 455 | ✗ | table_name = request->table_name(); | |
| 456 | const ParameterCompressReader* reader = | ||
| 457 | reinterpret_cast<const ParameterCompressReader*>( | ||
| 458 | ✗ | request->gradients().data()); | |
| 459 | ✗ | size = reader->item_size(); | |
| 460 | |||
| 461 | #ifdef ENABLE_PERF_REPORT | ||
| 462 | before_cache_update_time = std::chrono::high_resolution_clock::now(); | ||
| 463 | #endif | ||
| 464 | ✗ | success = cache_ps_->UpdateParameter(table_name, reader, 0); | |
| 465 | |||
| 466 | ✗ | RECSTORE_LOG_EVERY_MS(INFO, 2000) | |
| 467 | ✗ | << "UpdateParameter: table=" << table_name << ", keys=" << size; | |
| 468 | |||
| 469 | ✗ | reply->set_success(success); | |
| 470 | ✗ | } catch (const std::exception& e) { | |
| 471 | ✗ | LOG(ERROR) << "UpdateParameter error: " << e.what(); | |
| 472 | ✗ | reply->set_success(false); | |
| 473 | ✗ | } | |
| 474 | |||
| 475 | #ifdef ENABLE_PERF_REPORT | ||
| 476 | auto end_time = std::chrono::high_resolution_clock::now(); | ||
| 477 | double start_us_for_key = | ||
| 478 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 479 | start_time.time_since_epoch()) | ||
| 480 | .count(); | ||
| 481 | auto duration = | ||
| 482 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 483 | end_time - start_time) | ||
| 484 | .count(); | ||
| 485 | std::string op_latency_key = | ||
| 486 | "EmbUpdate|" + std::to_string(static_cast<uint64_t>(start_us_for_key)); | ||
| 487 | report("op_latency", | ||
| 488 | op_latency_key.c_str(), | ||
| 489 | "recserver_us", | ||
| 490 | static_cast<double>(duration)); | ||
| 491 | |||
| 492 | auto backend_update_duration = | ||
| 493 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 494 | end_time - before_cache_update_time) | ||
| 495 | .count(); | ||
| 496 | const uint64_t effective_trace_id = | ||
| 497 | trace_id == 0 ? static_cast<uint64_t>(start_us_for_key) : trace_id; | ||
| 498 | std::string update_stage_id = | ||
| 499 | "grpc_server::EmbUpdate|" + std::to_string(effective_trace_id); | ||
| 500 | report("embupdate_stages", | ||
| 501 | update_stage_id.c_str(), | ||
| 502 | "server_total_us", | ||
| 503 | static_cast<double>(duration)); | ||
| 504 | report("embupdate_stages", | ||
| 505 | update_stage_id.c_str(), | ||
| 506 | "server_backend_update_us", | ||
| 507 | static_cast<double>(backend_update_duration)); | ||
| 508 | report("embupdate_stages", | ||
| 509 | update_stage_id.c_str(), | ||
| 510 | "server_request_size", | ||
| 511 | static_cast<double>(size)); | ||
| 512 | report("embupdate_stages", | ||
| 513 | update_stage_id.c_str(), | ||
| 514 | "server_success", | ||
| 515 | success ? 1.0 : 0.0); | ||
| 516 | #endif | ||
| 517 | |||
| 518 | ✗ | return Status::OK; | |
| 519 | ✗ | } | |
| 520 | |||
| 521 | ✗ | Status InitEmbeddingTable(ServerContext* context, | |
| 522 | const InitEmbeddingTableRequest* request, | ||
| 523 | InitEmbeddingTableResponse* reply) override { | ||
| 524 | #ifdef ENABLE_PERF_REPORT | ||
| 525 | auto start_time = std::chrono::high_resolution_clock::now(); | ||
| 526 | #endif | ||
| 527 | try { | ||
| 528 | ✗ | if (request->has_config_payload()) { | |
| 529 | ✗ | auto payload = request->config_payload(); | |
| 530 | ✗ | nlohmann::json cfg = nlohmann::json::parse(payload); | |
| 531 | ✗ | uint64_t num_embeddings = cfg.value("num_embeddings", 0); | |
| 532 | ✗ | uint64_t embedding_dim = cfg.value("embedding_dim", 0); | |
| 533 | ✗ | uint64_t table_id = cfg.value("table_id", 0); | |
| 534 | ✗ | RECSTORE_LOG_EVERY_MS(INFO, 2000) | |
| 535 | ✗ | << "InitEmbeddingTable: table=" << request->table_name() | |
| 536 | ✗ | << ", num_embeddings=" << num_embeddings | |
| 537 | ✗ | << ", embedding_dim=" << embedding_dim << ", table_id=" << table_id; | |
| 538 | |||
| 539 | ✗ | const int tag = cache_ps_->InitTable( | |
| 540 | request->table_name(), num_embeddings, embedding_dim, table_id); | ||
| 541 | ✗ | reply->set_success(tag >= 0); | |
| 542 | ✗ | if (tag >= 0) { | |
| 543 | ✗ | reply->set_tag(tag); | |
| 544 | } | ||
| 545 | ✗ | } else { | |
| 546 | ✗ | LOG(WARNING) << "InitEmbeddingTable called without config_payload"; | |
| 547 | ✗ | reply->set_success(false); | |
| 548 | } | ||
| 549 | ✗ | } catch (const std::exception& e) { | |
| 550 | ✗ | LOG(ERROR) << "InitEmbeddingTable error: " << e.what(); | |
| 551 | ✗ | reply->set_success(false); | |
| 552 | ✗ | } | |
| 553 | |||
| 554 | #ifdef ENABLE_PERF_REPORT | ||
| 555 | auto end_time = std::chrono::high_resolution_clock::now(); | ||
| 556 | double start_us_for_key = | ||
| 557 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 558 | start_time.time_since_epoch()) | ||
| 559 | .count(); | ||
| 560 | auto duration = | ||
| 561 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 562 | end_time - start_time) | ||
| 563 | .count(); | ||
| 564 | std::string op_latency_key = | ||
| 565 | "InitEmbeddingTable|" + | ||
| 566 | std::to_string(static_cast<uint64_t>(start_us_for_key)); | ||
| 567 | report("op_latency", | ||
| 568 | op_latency_key.c_str(), | ||
| 569 | "recserver_us", | ||
| 570 | static_cast<double>(duration)); | ||
| 571 | #endif | ||
| 572 | |||
| 573 | ✗ | return Status::OK; | |
| 574 | } | ||
| 575 | |||
| 576 | private: | ||
| 577 | CachePS* cache_ps_; | ||
| 578 | int shard_id_ = 0; | ||
| 579 | std::atomic<uint64_t> total_get_requests_{0}; | ||
| 580 | std::atomic<uint64_t> total_put_requests_{0}; | ||
| 581 | std::atomic<uint64_t> total_get_keys_{0}; | ||
| 582 | std::atomic<uint64_t> total_put_keys_{0}; | ||
| 583 | std::atomic<uint64_t> total_get_bytes_{0}; | ||
| 584 | std::atomic<uint64_t> total_put_bytes_{0}; | ||
| 585 | std::chrono::steady_clock::time_point start_time_; | ||
| 586 | }; | ||
| 587 | |||
| 588 | namespace recstore { | ||
| 589 | class GRPCParameterServer : public BaseParameterServer { | ||
| 590 | public: | ||
| 591 | ✗ | GRPCParameterServer() = default; | |
| 592 | |||
| 593 | ✗ | void Run() { | |
| 594 | // Check whether multi-shard mode is configured | ||
| 595 | ✗ | int num_shards = 1; // default: single shard | |
| 596 | ✗ | if (config_["cache_ps"].contains("num_shards")) { | |
| 597 | ✗ | num_shards = config_["cache_ps"]["num_shards"]; | |
| 598 | } | ||
| 599 | const std::optional<int> local_shard_id = | ||
| 600 | ✗ | FLAGS_grpc_local_shard_id >= 0 | |
| 601 | ✗ | ? std::make_optional(FLAGS_grpc_local_shard_id) | |
| 602 | ✗ | : std::nullopt; | |
| 603 | |||
| 604 | ✗ | if (num_shards > 1) { | |
| 605 | // Multi-server startup | ||
| 606 | std::cout << "Starting distributed parameter server (gRPC), number " | ||
| 607 | ✗ | "of shards: " | |
| 608 | ✗ | << num_shards << std::endl; | |
| 609 | |||
| 610 | ✗ | if (!config_["cache_ps"].contains("servers")) { | |
| 611 | ✗ | LOG(FATAL) << "num_shards > 1 but cache_ps.servers is missing"; | |
| 612 | ✗ | return; | |
| 613 | } | ||
| 614 | |||
| 615 | ✗ | const auto& cache_ps_config = config_["cache_ps"]; | |
| 616 | ✗ | auto servers = SelectGRPCShardConfigs(cache_ps_config, local_shard_id); | |
| 617 | ✗ | const auto configured_servers = cache_ps_config["servers"]; | |
| 618 | ✗ | if (configured_servers.size() != num_shards) { | |
| 619 | ✗ | LOG(FATAL) << "servers count (" << configured_servers.size() | |
| 620 | ✗ | << ") does not match num_shards (" << num_shards << ")"; | |
| 621 | return; | ||
| 622 | } | ||
| 623 | ✗ | if (local_shard_id.has_value() && servers.empty()) { | |
| 624 | ✗ | LOG(FATAL) << "grpc_local_shard_id=" << *local_shard_id | |
| 625 | ✗ | << " is not present in cache_ps.servers"; | |
| 626 | return; | ||
| 627 | } | ||
| 628 | ✗ | if (!local_shard_id.has_value() && | |
| 629 | ✗ | servers.size() != configured_servers.size()) { | |
| 630 | ✗ | LOG(FATAL) << "Selected shard count (" << servers.size() | |
| 631 | ✗ | << ") does not match configured server count (" | |
| 632 | ✗ | << configured_servers.size() << ")"; | |
| 633 | return; | ||
| 634 | } | ||
| 635 | |||
| 636 | ✗ | std::vector<std::thread> server_threads; | |
| 637 | |||
| 638 | ✗ | for (auto& server_config : servers) { | |
| 639 | ✗ | server_threads.emplace_back([this, server_config]() { | |
| 640 | try { | ||
| 641 | ✗ | std::string host = server_config["host"]; | |
| 642 | ✗ | int port = server_config["port"]; | |
| 643 | ✗ | int shard = server_config["shard"]; | |
| 644 | |||
| 645 | ✗ | std::string server_address = host + ":" + std::to_string(port); | |
| 646 | |||
| 647 | ✗ | nlohmann::json shard_config = config_["cache_ps"]; | |
| 648 | ✗ | if (shard_config.contains("base_kv_config") && | |
| 649 | ✗ | shard_config["base_kv_config"].is_object()) { | |
| 650 | ✗ | auto& base_kv_config = shard_config["base_kv_config"]; | |
| 651 | ✗ | AppendShardSuffixIfPresent(base_kv_config, "path", shard); | |
| 652 | ✗ | AppendShardSuffixIfPresent(base_kv_config, "rocksdb_path", shard); | |
| 653 | ✗ | AppendShardSuffixToNestedFilePaths(base_kv_config, shard); | |
| 654 | ✗ | LOG(INFO) << "gRPC shard " << shard | |
| 655 | ✗ | << " using base_kv_config: " << base_kv_config.dump(); | |
| 656 | } | ||
| 657 | |||
| 658 | ✗ | auto cache_ps = std::make_unique<CachePS>(shard_config); | |
| 659 | ✗ | ParameterServiceImpl service(cache_ps.get(), shard); | |
| 660 | |||
| 661 | ✗ | grpc::EnableDefaultHealthCheckService(true); | |
| 662 | ✗ | grpc::reflection::InitProtoReflectionServerBuilderPlugin(); | |
| 663 | ✗ | ServerBuilder builder; | |
| 664 | ✗ | builder.AddListeningPort( | |
| 665 | ✗ | server_address, grpc::InsecureServerCredentials()); | |
| 666 | ✗ | builder.RegisterService(&service); | |
| 667 | ✗ | builder.SetMaxReceiveMessageSize(-1); // Unlimited | |
| 668 | ✗ | builder.SetMaxSendMessageSize(-1); // Unlimited | |
| 669 | ✗ | std::unique_ptr<Server> server(builder.BuildAndStart()); | |
| 670 | |||
| 671 | ✗ | if (!server) { | |
| 672 | std::string err_msg = fmt::format( | ||
| 673 | "FATAL: Failed to start gRPC server shard {} " | ||
| 674 | "on {}. " | ||
| 675 | "Port might be in use or invalid " | ||
| 676 | "configuration. " | ||
| 677 | "Check if port {} is already occupied.", | ||
| 678 | shard, | ||
| 679 | server_address, | ||
| 680 | ✗ | port); | |
| 681 | ✗ | std::cerr << err_msg << std::endl; | |
| 682 | ✗ | LOG(FATAL) << err_msg; | |
| 683 | return; | ||
| 684 | ✗ | } | |
| 685 | ✗ | std::cout << "Server shard " << shard << " listening on " | |
| 686 | ✗ | << server_address << std::endl; | |
| 687 | ✗ | server->Wait(); | |
| 688 | ✗ | } catch (const std::exception& e) { | |
| 689 | std::cerr << "FATAL: Uncaught exception in shard thread: " | ||
| 690 | ✗ | << e.what() << std::endl; | |
| 691 | ✗ | LOG(FATAL) << "Uncaught exception in shard thread: " << e.what(); | |
| 692 | ✗ | } catch (...) { | |
| 693 | ✗ | std::cerr << "FATAL: Unknown exception in shard thread" | |
| 694 | ✗ | << std::endl; | |
| 695 | ✗ | LOG(FATAL) << "Unknown exception in shard thread"; | |
| 696 | ✗ | } | |
| 697 | }); | ||
| 698 | } | ||
| 699 | |||
| 700 | // Wait for all server threads | ||
| 701 | ✗ | for (auto& t : server_threads) { | |
| 702 | ✗ | t.join(); | |
| 703 | } | ||
| 704 | ✗ | } else { | |
| 705 | // Single-server startup | ||
| 706 | ✗ | std::cout << "Starting single parameter server" << std::endl; | |
| 707 | ✗ | std::string server_address("0.0.0.0:15000"); | |
| 708 | ✗ | auto cache_ps = std::make_unique<CachePS>(config_["cache_ps"]); | |
| 709 | ✗ | ParameterServiceImpl service(cache_ps.get(), 0); | |
| 710 | |||
| 711 | ✗ | std::atomic<bool> metrics_running{true}; | |
| 712 | ✗ | std::thread metrics_thread([&service, &metrics_running]() { | |
| 713 | ✗ | while (metrics_running) { | |
| 714 | ✗ | std::this_thread::sleep_for(std::chrono::seconds(10)); | |
| 715 | ✗ | service.PrintMetrics(); | |
| 716 | ✗ | service.ResetMetrics(); | |
| 717 | } | ||
| 718 | ✗ | }); | |
| 719 | |||
| 720 | ✗ | grpc::EnableDefaultHealthCheckService(true); | |
| 721 | ✗ | grpc::reflection::InitProtoReflectionServerBuilderPlugin(); | |
| 722 | ✗ | ServerBuilder builder; | |
| 723 | ✗ | builder.AddListeningPort( | |
| 724 | ✗ | server_address, grpc::InsecureServerCredentials()); | |
| 725 | ✗ | builder.RegisterService(&service); | |
| 726 | ✗ | builder.SetMaxReceiveMessageSize(-1); // Unlimited | |
| 727 | ✗ | builder.SetMaxSendMessageSize(-1); // Unlimited | |
| 728 | ✗ | std::unique_ptr<Server> server(builder.BuildAndStart()); | |
| 729 | ✗ | std::cerr << "sever built succesfully" << std::endl; | |
| 730 | ✗ | if (!server) { | |
| 731 | std::string err_msg = fmt::format( | ||
| 732 | "FATAL: Failed to start gRPC server on {}. " | ||
| 733 | "Port might be in use or invalid configuration.", | ||
| 734 | ✗ | server_address); | |
| 735 | ✗ | std::cerr << err_msg << std::endl; | |
| 736 | ✗ | LOG(FATAL) << err_msg; | |
| 737 | metrics_running = false; | ||
| 738 | if (metrics_thread.joinable()) { | ||
| 739 | metrics_thread.join(); | ||
| 740 | } | ||
| 741 | return; | ||
| 742 | ✗ | } | |
| 743 | ✗ | std::cout << "Server listening on " << server_address << std::endl; | |
| 744 | ✗ | server->Wait(); | |
| 745 | |||
| 746 | ✗ | metrics_running = false; | |
| 747 | ✗ | if (metrics_thread.joinable()) { | |
| 748 | ✗ | metrics_thread.join(); | |
| 749 | } | ||
| 750 | ✗ | } | |
| 751 | } | ||
| 752 | }; | ||
| 753 | |||
| 754 | FACTORY_REGISTER(BaseParameterServer, GRPCParameterServer, GRPCParameterServer); | ||
| 755 | |||
| 756 | } // namespace recstore | ||
| 757 | |||
| 758 | #ifndef RECSTORE_NO_SERVER_MAIN | ||
| 759 | ✗ | int main(int argc, char** argv) { | |
| 760 | ✗ | base::Init(&argc, &argv); | |
| 761 | ✗ | xmh::Reporter::StartReportThread(2000); | |
| 762 | const std::string config_path = | ||
| 763 | ✗ | FLAGS_config_path.empty() | |
| 764 | ✗ | ? base::ResolveRecStoreConfigPath().string() | |
| 765 | ✗ | : FLAGS_config_path; | |
| 766 | ✗ | std::ifstream config_file(config_path); | |
| 767 | ✗ | if (!config_file.is_open()) { | |
| 768 | ✗ | throw std::runtime_error("Cannot open config file: " + config_path); | |
| 769 | } | ||
| 770 | ✗ | nlohmann::json ex; | |
| 771 | ✗ | config_file >> ex; | |
| 772 | ✗ | recstore::GRPCParameterServer ps; | |
| 773 | ✗ | std::cout << "Parameter server config: " << ex.dump(2) << std::endl; | |
| 774 | ✗ | ps.Init(ex); | |
| 775 | ✗ | ps.Run(); | |
| 776 | ✗ | return 0; | |
| 777 | ✗ | } | |
| 778 | #endif | ||
| 779 |