GCC Code Coverage Report

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