framework/op.cc
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #include "framework/op.h" | ||
| 2 | #include "framework/common/hierkv_local_runtime.h" | ||
| 3 | #include "framework/common/local_shm_op_component.h" | ||
| 4 | #include "framework/common/op_runtime_support.h" | ||
| 5 | #include "framework/common/ps_client_config_adapter.h" | ||
| 6 | #include "ps/client_factory.h" | ||
| 7 | #include "ps/rdma/rdma_ps_client_adapter.h" | ||
| 8 | #include <algorithm> | ||
| 9 | #include <cctype> | ||
| 10 | #include <cstring> | ||
| 11 | #include <immintrin.h> | ||
| 12 | #include <iostream> | ||
| 13 | #include <stdexcept> | ||
| 14 | #include <vector> | ||
| 15 | #include <unordered_map> | ||
| 16 | #include <filesystem> | ||
| 17 | #include <mutex> | ||
| 18 | #include <memory> | ||
| 19 | #include <numeric> | ||
| 20 | #include <thread> | ||
| 21 | #include <cstdlib> | ||
| 22 | #include <emmintrin.h> | ||
| 23 | #include <string> | ||
| 24 | #include <fstream> | ||
| 25 | #include "base/tensor.h" | ||
| 26 | #include <glog/logging.h> | ||
| 27 | #ifdef ENABLE_PERF_REPORT | ||
| 28 | # include "base/report/report_client.h" | ||
| 29 | #endif | ||
| 30 | |||
| 31 | namespace recstore { | ||
| 32 | |||
| 33 | namespace { | ||
| 34 | 24 | std::string NormalizeBackendName(std::string backend_name) { | |
| 35 | 24 | std::transform( | |
| 36 | backend_name.begin(), | ||
| 37 | backend_name.end(), | ||
| 38 | backend_name.begin(), | ||
| 39 | 136 | [](unsigned char c) { return static_cast<char>(std::tolower(c)); }); | |
| 40 | 24 | return backend_name; | |
| 41 | } | ||
| 42 | |||
| 43 | 58 | bool IsReadWriteSuccess(BasePSClient* client, int ret) { | |
| 44 | (void)client; | ||
| 45 | 58 | return ret == 0; | |
| 46 | } | ||
| 47 | |||
| 48 | 24 | std::string ResolveBackendNameWithHierKV(const json& config) { | |
| 49 | 24 | if (config.contains("cache_ps") && config["cache_ps"].contains("ps_type")) { | |
| 50 | const std::string ps_type = | ||
| 51 | 24 | NormalizeBackendName(config["cache_ps"]["ps_type"].get<std::string>()); | |
| 52 | 24 | if (IsHierKVBackendName(ps_type)) { | |
| 53 | 20 | return ps_type; | |
| 54 | } | ||
| 55 | 24 | } | |
| 56 | 4 | switch (ResolveFrameworkPSClientType(config)) { | |
| 57 | 4 | case PSClientType::kGrpc: | |
| 58 | 4 | return HasFrameworkDistributedClientConfig(config) | |
| 59 | ? "distributed_grpc" | ||
| 60 | 8 | : "grpc"; | |
| 61 | ✗ | case PSClientType::kBrpc: | |
| 62 | ✗ | return HasFrameworkDistributedClientConfig(config) | |
| 63 | ? "distributed_brpc" | ||
| 64 | ✗ | : "brpc"; | |
| 65 | ✗ | case PSClientType::kRdma: | |
| 66 | ✗ | return "rdma"; | |
| 67 | ✗ | case PSClientType::kLocalShm: | |
| 68 | ✗ | return "local_shm"; | |
| 69 | } | ||
| 70 | |||
| 71 | ✗ | return "unknown"; | |
| 72 | } | ||
| 73 | } // namespace | ||
| 74 | |||
| 75 | 70 | void validate_keys(const base::RecTensor& keys) { | |
| 76 | 70 | if (keys.dtype() != base::DataType::UINT64) { | |
| 77 | 2 | throw std::invalid_argument("Keys tensor must have dtype UINT64, but got " + | |
| 78 | 4 | base::DataTypeToString(keys.dtype())); | |
| 79 | } | ||
| 80 | 68 | if (keys.dim() != 1) { | |
| 81 | ✗ | throw std::invalid_argument("Keys tensor must be 1-dimensional, but has " + | |
| 82 | ✗ | std::to_string(keys.dim()) + " dimensions."); | |
| 83 | } | ||
| 84 | 68 | } | |
| 85 | |||
| 86 | 66 | void validate_embeddings(const base::RecTensor& embeddings, | |
| 87 | const std::string& name) { | ||
| 88 | 66 | if (embeddings.dtype() != base::DataType::FLOAT32) { | |
| 89 | ✗ | throw std::invalid_argument( | |
| 90 | ✗ | name + " tensor must have dtype FLOAT32, but got " + | |
| 91 | ✗ | base::DataTypeToString(embeddings.dtype())); | |
| 92 | } | ||
| 93 | 66 | if (embeddings.dim() != 2) { | |
| 94 | ✗ | throw std::invalid_argument( | |
| 95 | ✗ | name + " tensor must be 2-dimensional, but has " + | |
| 96 | ✗ | std::to_string(embeddings.dim()) + " dimensions."); | |
| 97 | } | ||
| 98 | // No fixed embedding dimension check for mock. | ||
| 99 | 66 | } | |
| 100 | |||
| 101 | 2 | void KVClientOp::EmbInit(const base::RecTensor& keys, | |
| 102 | const base::RecTensor& init_values) { | ||
| 103 | 2 | EmbWrite(keys, init_values); | |
| 104 | 2 | } | |
| 105 | |||
| 106 | 2 | void KVClientOp::EmbDelete(const base::RecTensor& keys) { | |
| 107 | 2 | throw std::runtime_error("Not impl"); | |
| 108 | } | ||
| 109 | 2 | bool KVClientOp::EmbExists(const base::RecTensor& keys) { | |
| 110 | 2 | throw std::runtime_error("Not impl"); | |
| 111 | } | ||
| 112 | |||
| 113 | 2 | void KVClientOp::WaitForWrite(uint64_t write_id) { | |
| 114 | 2 | throw std::runtime_error("Not impl"); | |
| 115 | } | ||
| 116 | 2 | void KVClientOp::SaveToFile(const std::string& path) { | |
| 117 | 2 | throw std::runtime_error("Not impl"); | |
| 118 | } | ||
| 119 | 2 | void KVClientOp::LoadFromFile(const std::string& path) { | |
| 120 | 2 | throw std::runtime_error("Not impl"); | |
| 121 | } | ||
| 122 | |||
| 123 | ✗ | bool KVClientOp::SaveCheckpoint(const std::string& path, | |
| 124 | const std::string& metadata) { | ||
| 125 | ✗ | if (path.empty() || metadata.empty()) { | |
| 126 | ✗ | throw std::invalid_argument( | |
| 127 | ✗ | "checkpoint path and metadata must be non-empty"); | |
| 128 | } | ||
| 129 | #ifdef USE_FAKE_KVCLIENT | ||
| 130 | throw std::runtime_error("Checkpoint save is unsupported by fake KVClient"); | ||
| 131 | #else | ||
| 132 | ✗ | if (ps_client_ == nullptr) { | |
| 133 | ✗ | throw std::runtime_error("PS client is not initialized"); | |
| 134 | } | ||
| 135 | ✗ | return ps_client_->SaveCheckpoint(path, metadata); | |
| 136 | #endif | ||
| 137 | } | ||
| 138 | |||
| 139 | ✗ | bool KVClientOp::LoadCheckpoint(const std::string& path, | |
| 140 | const std::string& metadata) { | ||
| 141 | ✗ | if (path.empty() || metadata.empty()) { | |
| 142 | ✗ | throw std::invalid_argument( | |
| 143 | ✗ | "checkpoint path and metadata must be non-empty"); | |
| 144 | } | ||
| 145 | #ifdef USE_FAKE_KVCLIENT | ||
| 146 | throw std::runtime_error("Checkpoint load is unsupported by fake KVClient"); | ||
| 147 | #else | ||
| 148 | ✗ | if (ps_client_ == nullptr) { | |
| 149 | ✗ | throw std::runtime_error("PS client is not initialized"); | |
| 150 | } | ||
| 151 | ✗ | return ps_client_->LoadCheckpoint(path, metadata); | |
| 152 | #endif | ||
| 153 | } | ||
| 154 | |||
| 155 | 2 | uint64_t KVClientOp::EmbWriteAsync(const base::RecTensor& keys, | |
| 156 | const base::RecTensor& values) { | ||
| 157 | 2 | throw std::runtime_error("Not impl"); | |
| 158 | } | ||
| 159 | |||
| 160 | 70 | std::shared_ptr<CommonOp> GetKVClientOp() { | |
| 161 | 70 | static std::shared_ptr<CommonOp> instance; | |
| 162 | static std::once_flag once_flag; | ||
| 163 | 70 | std::call_once(once_flag, []() { | |
| 164 | 6 | instance = std::make_shared<KVClientOp>(); | |
| 165 | 6 | }); | |
| 166 | 70 | return instance; | |
| 167 | } | ||
| 168 | |||
| 169 | } // namespace recstore | ||
| 170 | |||
| 171 | #ifndef USE_FAKE_KVCLIENT | ||
| 172 | |||
| 173 | namespace recstore { | ||
| 174 | |||
| 175 | 36 | KVClientOp::KVClientOp() { | |
| 176 | 36 | if (!ps_client_) { | |
| 177 | try { | ||
| 178 | 24 | json config = GetGlobalConfig(); | |
| 179 | 24 | ps_backend_name_ = ResolveBackendNameWithHierKV(config); | |
| 180 | 24 | if (IsHierKVBackendName(ps_backend_name_)) { | |
| 181 | 20 | ConfigureLogging(); | |
| 182 | 20 | LOG(INFO) << "Initialized local HierKV backend in KVClientOp."; | |
| 183 | 20 | return; | |
| 184 | } | ||
| 185 | 4 | bool use_rdma = false; | |
| 186 | try { | ||
| 187 | 4 | use_rdma = ResolveFrameworkPSClientType(config) == PSClientType::kRdma; | |
| 188 | ✗ | } catch (...) { | |
| 189 | ✗ | use_rdma = false; | |
| 190 | ✗ | } | |
| 191 | |||
| 192 | 4 | if (use_rdma) { | |
| 193 | ✗ | InitializeRdmaProcessRuntime(); | |
| 194 | ✗ | ConfigureLogging(false); | |
| 195 | } else { | ||
| 196 | 4 | ConfigureLogging(); | |
| 197 | } | ||
| 198 | 4 | ps_client_holder_ = create_ps_client_from_config(config); | |
| 199 | 4 | ps_client_ = ps_client_holder_.get(); | |
| 200 | |||
| 201 | 4 | LOG(INFO) << "PS client initialized successfully."; | |
| 202 | 24 | } catch (const std::exception& e) { | |
| 203 | ✗ | LOG(ERROR) << "Failed to initialize PS client: " << std::string(e.what()); | |
| 204 | ✗ | throw; | |
| 205 | ✗ | } | |
| 206 | } | ||
| 207 | ✗ | } | |
| 208 | |||
| 209 | BasePSClient* KVClientOp::ps_client_ = nullptr; | ||
| 210 | std::unique_ptr<BasePSClient> KVClientOp::ps_client_holder_; | ||
| 211 | |||
| 212 | ✗ | void KVClientOp::SetPSConfig(const std::string& host, int port) { | |
| 213 | ✗ | if (IsHierKVBackendName(ps_backend_name_)) { | |
| 214 | ✗ | LOG(INFO) << "HierKV backend ignores set_ps_config host=" << host | |
| 215 | ✗ | << " port=" << port; | |
| 216 | ✗ | return; | |
| 217 | } | ||
| 218 | ✗ | ps_client_holder_.reset(); | |
| 219 | ✗ | ps_client_ = nullptr; | |
| 220 | |||
| 221 | ✗ | json file_config = GetGlobalConfig(); | |
| 222 | ✗ | int final_port = port; | |
| 223 | ✗ | if (final_port <= 0) { | |
| 224 | ✗ | if (file_config.contains("client") && | |
| 225 | ✗ | file_config["client"].contains("port")) { | |
| 226 | ✗ | final_port = file_config["client"]["port"].get<int>(); | |
| 227 | ✗ | } else if (file_config.contains("cache_ps") && | |
| 228 | ✗ | file_config["cache_ps"].contains("servers") && | |
| 229 | ✗ | file_config["cache_ps"]["servers"].is_array() && | |
| 230 | ✗ | !file_config["cache_ps"]["servers"].empty()) { | |
| 231 | ✗ | final_port = file_config["cache_ps"]["servers"][0]["port"].get<int>(); | |
| 232 | } else { | ||
| 233 | ✗ | final_port = 15000; | |
| 234 | } | ||
| 235 | } | ||
| 236 | |||
| 237 | ✗ | std::string final_host = host; | |
| 238 | ✗ | if (final_host.empty()) { | |
| 239 | ✗ | final_host = "127.0.0.1"; | |
| 240 | } | ||
| 241 | |||
| 242 | ✗ | json config = file_config; | |
| 243 | ✗ | config.erase("distributed_client"); | |
| 244 | ✗ | if (!config.contains("client")) { | |
| 245 | ✗ | config["client"] = json::object(); | |
| 246 | } | ||
| 247 | ✗ | config["client"]["host"] = final_host; | |
| 248 | ✗ | config["client"]["port"] = final_port; | |
| 249 | ✗ | config["client"]["shard"] = 0; | |
| 250 | |||
| 251 | ✗ | ps_client_holder_ = create_ps_client_from_config(config); | |
| 252 | ✗ | ps_client_ = ps_client_holder_.get(); | |
| 253 | ✗ | ps_backend_name_ = ResolveBackendNameWithHierKV(config); | |
| 254 | ✗ | LOG(INFO) << "Re-initialized PS client with host=" << final_host | |
| 255 | ✗ | << " port=" << final_port; | |
| 256 | ✗ | } | |
| 257 | |||
| 258 | ✗ | void KVClientOp::SetPSBackend(const std::string& backend) { | |
| 259 | ✗ | if (backend.empty()) { | |
| 260 | ✗ | throw std::invalid_argument("backend must be non-empty"); | |
| 261 | } | ||
| 262 | |||
| 263 | ✗ | const std::string normalized_backend = NormalizeBackendName(backend); | |
| 264 | ✗ | if (IsHierKVBackendName(normalized_backend)) { | |
| 265 | ✗ | ps_client_holder_.reset(); | |
| 266 | ✗ | ps_client_ = nullptr; | |
| 267 | ✗ | ps_backend_name_ = normalized_backend; | |
| 268 | ✗ | LOG(INFO) << "Switched KVClientOp backend to local HierKV runtime."; | |
| 269 | ✗ | return; | |
| 270 | } | ||
| 271 | |||
| 272 | ✗ | json config = GetGlobalConfig(); | |
| 273 | ✗ | if (!config.contains("cache_ps")) { | |
| 274 | ✗ | config["cache_ps"] = json::object(); | |
| 275 | } | ||
| 276 | ✗ | config["cache_ps"]["ps_type"] = NormalizePSType(backend); | |
| 277 | |||
| 278 | ✗ | ps_client_holder_.reset(); | |
| 279 | ✗ | ps_client_ = nullptr; | |
| 280 | ✗ | ps_client_holder_ = create_ps_client_from_config(config); | |
| 281 | ✗ | ps_client_ = ps_client_holder_.get(); | |
| 282 | ✗ | ps_backend_name_ = ResolveBackendNameWithHierKV(config); | |
| 283 | ✗ | LOG(INFO) << "Re-initialized PS client with backend=" << ps_backend_name_; | |
| 284 | ✗ | } | |
| 285 | |||
| 286 | 4 | std::string KVClientOp::CurrentPSBackend() const { return ps_backend_name_; } | |
| 287 | |||
| 288 | 48 | void KVClientOp::EmbRead(const RecTensor& keys, RecTensor& values) { | |
| 289 | 48 | if (IsHierKVBackendName(ps_backend_name_)) { | |
| 290 | 10 | GetHierKVLocalRuntime().Read(keys, values); | |
| 291 | 6 | return; | |
| 292 | } | ||
| 293 | 38 | if (ps_client_ == nullptr) { | |
| 294 | ✗ | throw std::runtime_error("PS client is not initialized. Please call " | |
| 295 | ✗ | "KVClientOp::SetPSClient() first."); | |
| 296 | } | ||
| 297 | |||
| 298 | # ifdef ENABLE_PERF_REPORT | ||
| 299 | auto start_time = std::chrono::high_resolution_clock::now(); | ||
| 300 | double start_us = | ||
| 301 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 302 | start_time.time_since_epoch()) | ||
| 303 | .count(); | ||
| 304 | std::string report_id = | ||
| 305 | "op::EmbRead|" + std::to_string(static_cast<uint64_t>(start_us)); | ||
| 306 | std::string unique_id = | ||
| 307 | "embread_debug|" + std::to_string(static_cast<uint64_t>(start_us)); | ||
| 308 | # endif | ||
| 309 | |||
| 310 | 76 | LOG(INFO) << "EmbRead: keys.shape=" << keys.shape(0) << ", values.shape=[" | |
| 311 | 38 | << values.shape(0) << ", " << values.shape(1) << "]"; | |
| 312 | 76 | LOG(INFO) << "EmbRead: keys.data=" << keys.data_as<uint64_t>() | |
| 313 | 38 | << ", values.data=" << values.data_as<float>(); | |
| 314 | 38 | if (keys.shape(0) > 0) { | |
| 315 | 38 | std::ostringstream oss; | |
| 316 | 38 | oss << "EmbRead: keys start with: "; | |
| 317 | 142 | for (int i = 0; i < std::min((int64_t)10, keys.shape(0)); ++i) | |
| 318 | 104 | oss << keys.data_as<uint64_t>()[i] << ", "; | |
| 319 | 38 | LOG(INFO) << oss.str(); | |
| 320 | 38 | } | |
| 321 | 38 | if (values.shape(0) > 0) { | |
| 322 | 38 | std::ostringstream oss; | |
| 323 | 38 | oss << "EmbRead: values start with: "; | |
| 324 | 140 | for (int i = 0; i < std::min((int64_t)10, values.shape(0)); ++i) { | |
| 325 | 102 | oss << "["; | |
| 326 | 1052 | for (int j = 0; j < std::min((int64_t)10, values.shape(1)); ++j) { | |
| 327 | 950 | oss << values.data_as<float>()[i * values.shape(1) + j] << ", "; | |
| 328 | } | ||
| 329 | 102 | oss << "] "; | |
| 330 | } | ||
| 331 | 38 | LOG(INFO) << oss.str(); | |
| 332 | 38 | } | |
| 333 | 38 | validate_keys(keys); | |
| 334 | 38 | validate_embeddings(values, "Values"); | |
| 335 | |||
| 336 | 38 | const int64_t L = keys.shape(0); | |
| 337 | 38 | if (values.shape(0) != L) { | |
| 338 | 2 | throw std::invalid_argument( | |
| 339 | 4 | "Dimension mismatch: Keys has length " + std::to_string(L) + | |
| 340 | 8 | " but values has length " + std::to_string(values.shape(0))); | |
| 341 | } | ||
| 342 | 36 | const uint64_t* keys_data = keys.data_as<uint64_t>(); | |
| 343 | 36 | base::ConstArray<uint64_t> keys_array(keys_data, L); | |
| 344 | 36 | float* values_data = values.data_as<float>(); | |
| 345 | |||
| 346 | 36 | const int64_t D = values.shape(1); | |
| 347 | 36 | const size_t total = static_cast<size_t>(L) * static_cast<size_t>(D); | |
| 348 | 36 | std::fill_n(values_data, total, 0.0f); | |
| 349 | |||
| 350 | 36 | int ret = ps_client_->GetParameter(keys_array, values); | |
| 351 | 36 | if (!IsReadWriteSuccess(ps_client_, ret)) { | |
| 352 | 2 | throw std::runtime_error("Failed to read embeddings from PS client."); | |
| 353 | } | ||
| 354 | |||
| 355 | # ifdef ENABLE_PERF_REPORT | ||
| 356 | auto end_time = std::chrono::high_resolution_clock::now(); | ||
| 357 | auto duration = | ||
| 358 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 359 | end_time - start_time) | ||
| 360 | .count(); | ||
| 361 | std::string op_latency_key = | ||
| 362 | "EmbRead|" + std::to_string(static_cast<uint64_t>(start_us)); | ||
| 363 | report("op_latency", | ||
| 364 | op_latency_key.c_str(), | ||
| 365 | "recstore_us", | ||
| 366 | static_cast<double>(duration)); | ||
| 367 | |||
| 368 | report("embread_stages", | ||
| 369 | report_id.c_str(), | ||
| 370 | "duration_us", | ||
| 371 | static_cast<double>(duration)); | ||
| 372 | |||
| 373 | report("embread_stages", | ||
| 374 | report_id.c_str(), | ||
| 375 | "request_size", | ||
| 376 | static_cast<double>(keys.shape(0))); | ||
| 377 | |||
| 378 | FlameGraphData op_data = { | ||
| 379 | "op::EmbRead", | ||
| 380 | start_us, | ||
| 381 | 0, // level | ||
| 382 | static_cast<double>(duration), | ||
| 383 | static_cast<double>(duration)}; | ||
| 384 | report_flame_graph("emb_read_flame_map", unique_id.c_str(), op_data); | ||
| 385 | # endif | ||
| 386 | } | ||
| 387 | |||
| 388 | 4 | void KVClientOp::EmbUpdate(const base::RecTensor& keys, | |
| 389 | const base::RecTensor& grads) { | ||
| 390 | 8 | EmbUpdate("default", keys, grads); | |
| 391 | ✗ | } | |
| 392 | |||
| 393 | 10 | void KVClientOp::EmbUpdate(const std::string& table_name, | |
| 394 | const base::RecTensor& keys, | ||
| 395 | const base::RecTensor& grads) { | ||
| 396 | 10 | if (IsHierKVBackendName(ps_backend_name_)) { | |
| 397 | 4 | GetHierKVLocalRuntime().Update(table_name, keys, grads); | |
| 398 | 2 | return; | |
| 399 | } | ||
| 400 | 6 | if (ps_client_ == nullptr) { | |
| 401 | ✗ | throw std::runtime_error("PS client is not initialized. Please call " | |
| 402 | ✗ | "KVClientOp::SetPSClient() first."); | |
| 403 | } | ||
| 404 | |||
| 405 | # ifdef ENABLE_PERF_REPORT | ||
| 406 | auto start_time = std::chrono::high_resolution_clock::now(); | ||
| 407 | const uint64_t trace_id = static_cast<uint64_t>( | ||
| 408 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 409 | start_time.time_since_epoch()) | ||
| 410 | .count()); | ||
| 411 | struct TraceGuard { | ||
| 412 | explicit TraceGuard(uint64_t new_trace_id) | ||
| 413 | : previous_trace_id_(recstore::g_trace_id) { | ||
| 414 | recstore::g_trace_id = new_trace_id; | ||
| 415 | } | ||
| 416 | ~TraceGuard() { recstore::g_trace_id = previous_trace_id_; } | ||
| 417 | uint64_t previous_trace_id_; | ||
| 418 | } trace_guard(trace_id); | ||
| 419 | # endif | ||
| 420 | |||
| 421 | 6 | int64_t validate_done_us = 0; | |
| 422 | 6 | validate_keys(keys); | |
| 423 | 6 | validate_embeddings(grads, "Grads"); | |
| 424 | |||
| 425 | 6 | const int64_t L = keys.shape(0); | |
| 426 | 6 | if (grads.shape(0) != L) { | |
| 427 | ✗ | throw std::invalid_argument( | |
| 428 | ✗ | "Dimension mismatch: Keys has length " + std::to_string(L) + | |
| 429 | ✗ | " but grads has length " + std::to_string(grads.shape(0))); | |
| 430 | } | ||
| 431 | |||
| 432 | 6 | const int64_t D = grads.shape(1); | |
| 433 | 6 | if (D <= 0) { | |
| 434 | ✗ | throw std::invalid_argument( | |
| 435 | ✗ | "Invalid grad dimension D: " + std::to_string(D)); | |
| 436 | } | ||
| 437 | |||
| 438 | # ifdef ENABLE_PERF_REPORT | ||
| 439 | auto validate_done_time = std::chrono::high_resolution_clock::now(); | ||
| 440 | validate_done_us = | ||
| 441 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 442 | validate_done_time - start_time) | ||
| 443 | .count(); | ||
| 444 | # endif | ||
| 445 | |||
| 446 | 6 | const uint64_t* keys_data = keys.data_as<uint64_t>(); | |
| 447 | 6 | base::ConstArray<uint64_t> keys_array(keys_data, L); | |
| 448 | |||
| 449 | 6 | int ret = ps_client_->UpdateParameter(table_name, keys_array, grads); | |
| 450 | 6 | if (ret != 0) { | |
| 451 | 2 | throw std::runtime_error("Failed to update embeddings via PS client."); | |
| 452 | } | ||
| 453 | |||
| 454 | # ifdef ENABLE_PERF_REPORT | ||
| 455 | auto end_time = std::chrono::high_resolution_clock::now(); | ||
| 456 | auto duration = | ||
| 457 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 458 | end_time - start_time) | ||
| 459 | .count(); | ||
| 460 | double start_us = | ||
| 461 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 462 | start_time.time_since_epoch()) | ||
| 463 | .count(); | ||
| 464 | std::string op_latency_key = | ||
| 465 | "EmbUpdate|" + std::to_string(static_cast<uint64_t>(start_us)); | ||
| 466 | report("op_latency", | ||
| 467 | op_latency_key.c_str(), | ||
| 468 | "recstore_us", | ||
| 469 | static_cast<double>(duration)); | ||
| 470 | |||
| 471 | std::string update_stage_id = | ||
| 472 | "op_client::EmbUpdate|" + std::to_string(trace_id); | ||
| 473 | report("embupdate_stages", | ||
| 474 | update_stage_id.c_str(), | ||
| 475 | "op_total_us", | ||
| 476 | static_cast<double>(duration)); | ||
| 477 | report("embupdate_stages", | ||
| 478 | update_stage_id.c_str(), | ||
| 479 | "op_validate_us", | ||
| 480 | static_cast<double>(validate_done_us)); | ||
| 481 | report("embupdate_stages", | ||
| 482 | update_stage_id.c_str(), | ||
| 483 | "request_size", | ||
| 484 | static_cast<double>(L)); | ||
| 485 | report("embupdate_stages", | ||
| 486 | update_stage_id.c_str(), | ||
| 487 | "embedding_dim", | ||
| 488 | static_cast<double>(D)); | ||
| 489 | # endif | ||
| 490 | } | ||
| 491 | |||
| 492 | ✗ | int64_t KVClientOp::SparseGradientCommit( | |
| 493 | const std::string& table, | ||
| 494 | SparseGradientCommitHeader header, | ||
| 495 | const base::RecTensor& keys, | ||
| 496 | const base::RecTensor& grads) { | ||
| 497 | ✗ | if (!ps_client_) | |
| 498 | ✗ | throw std::runtime_error("PS client is not initialized"); | |
| 499 | ✗ | return ps_client_->SparseGradientCommit( | |
| 500 | table, | ||
| 501 | header, | ||
| 502 | ✗ | base::ConstArray<uint64_t>(keys.data_as<uint64_t>(), keys.shape(0)), | |
| 503 | ✗ | grads); | |
| 504 | } | ||
| 505 | |||
| 506 | ✗ | int64_t KVClientOp::SparseGradientCommitAsync( | |
| 507 | const std::string& table, | ||
| 508 | SparseGradientCommitHeader header, | ||
| 509 | const base::RecTensor& keys, | ||
| 510 | const base::RecTensor& grads, | ||
| 511 | std::function<void()> ready) { | ||
| 512 | ✗ | if (!ps_client_) | |
| 513 | ✗ | throw std::runtime_error("PS client is not initialized"); | |
| 514 | ✗ | auto* client = ps_client_; | |
| 515 | ✗ | return sparse_commits_.Submit( | |
| 516 | ✗ | [client, table](SparseGradientCommitHeader h, | |
| 517 | const base::RecTensor& k, | ||
| 518 | ✗ | const base::RecTensor& g) { | |
| 519 | ✗ | return client->SparseGradientCommit( | |
| 520 | ✗ | table, | |
| 521 | h, | ||
| 522 | ✗ | base::ConstArray<uint64_t>(k.data_as<uint64_t>(), k.shape(0)), | |
| 523 | ✗ | g); | |
| 524 | }, | ||
| 525 | header, | ||
| 526 | keys, | ||
| 527 | grads, | ||
| 528 | ✗ | std::move(ready)); | |
| 529 | } | ||
| 530 | |||
| 531 | ✗ | void KVClientOp::WaitSparseGradientCommit(int64_t id) { | |
| 532 | ✗ | sparse_commits_.Wait(id); | |
| 533 | ✗ | } | |
| 534 | |||
| 535 | ✗ | void KVClientOp::CloseSparseGradientCommit(const std::string& table, | |
| 536 | SparseGradientCommitHeader header) { | ||
| 537 | ✗ | base::RecTensor keys({0}, base::DataType::UINT64); | |
| 538 | ✗ | base::RecTensor grads({0, int64_t(header.dim)}, base::DataType::FLOAT32); | |
| 539 | ✗ | SparseGradientCommit(table, header, keys, grads); | |
| 540 | ✗ | } | |
| 541 | |||
| 542 | 2 | uint64_t KVClientOp::EmbUpdateAsync( | |
| 543 | const std::string& table_name, | ||
| 544 | const base::RecTensor& keys, | ||
| 545 | const base::RecTensor& grads, | ||
| 546 | framework::CacheUpdateActions cache) { | ||
| 547 | 2 | if (ps_client_ == nullptr) | |
| 548 | ✗ | throw std::runtime_error("PS client is not initialized"); | |
| 549 | 2 | if (ps_backend_name_ != "rdma") | |
| 550 | 2 | throw std::runtime_error("emb_update_async requires the RDMA backend"); | |
| 551 | ✗ | validate_keys(keys); | |
| 552 | ✗ | validate_embeddings(grads, "Grads"); | |
| 553 | ✗ | if (table_name.empty() || keys.shape(0) != grads.shape(0) || | |
| 554 | ✗ | grads.shape(1) <= 0) | |
| 555 | ✗ | throw std::invalid_argument("Invalid asynchronous embedding update shape"); | |
| 556 | ✗ | const int64_t rows = keys.shape(0); | |
| 557 | ✗ | auto id = ps_client_->SubmitUpdateParameterAsync( | |
| 558 | table_name, | ||
| 559 | ✗ | base::ConstArray<uint64_t>(keys.data_as<uint64_t>(), rows), | |
| 560 | grads); | ||
| 561 | ✗ | pending_updates_.Add(id, rows, std::move(cache)); | |
| 562 | ✗ | return id; | |
| 563 | } | ||
| 564 | |||
| 565 | 2 | void KVClientOp::WaitForEmbUpdate(uint64_t update_id) { | |
| 566 | 2 | if (ps_client_ == nullptr) | |
| 567 | ✗ | throw std::runtime_error("PS client is not initialized"); | |
| 568 | 4 | pending_updates_.Wait(update_id, [this, update_id] { | |
| 569 | ✗ | if (ps_client_->WaitUpdateParameter(update_id) != 0) | |
| 570 | ✗ | throw std::runtime_error("Failed to complete asynchronous RDMA update"); | |
| 571 | ✗ | }); | |
| 572 | ✗ | } | |
| 573 | |||
| 574 | 2 | void KVClientOp::EmbUpdateCached( | |
| 575 | const std::string& table, | ||
| 576 | const base::RecTensor& keys, | ||
| 577 | const base::RecTensor& grads, | ||
| 578 | const framework::CacheUpdateActions& cache, | ||
| 579 | bool local) { | ||
| 580 | 2 | framework::g_cache_aware_op.Update(keys.shape(0), cache, [&] { | |
| 581 | 2 | if (local) | |
| 582 | ✗ | LocalUpdateFlat(table, keys, grads); | |
| 583 | else | ||
| 584 | 2 | EmbUpdate(table, keys, grads); | |
| 585 | 2 | }); | |
| 586 | 2 | } | |
| 587 | |||
| 588 | 18 | void KVClientOp::EmbWriteCached( | |
| 589 | const base::RecTensor& keys, | ||
| 590 | const base::RecTensor& values, | ||
| 591 | const framework::CacheUpdateActions& cache, | ||
| 592 | bool per_key) { | ||
| 593 | 18 | EmbWrite(keys, values); | |
| 594 | 18 | if (per_key) | |
| 595 | ✗ | framework::g_cache_aware_op.AfterUpdate(keys.shape(0), cache, true); | |
| 596 | 18 | else if (cache.enabled && cache.enabled()) | |
| 597 | ✗ | framework::CacheAwareOp::ClearNoThrow(cache.clear); | |
| 598 | 18 | } | |
| 599 | |||
| 600 | 10 | int KVClientOp::InitEmbeddingTableCached( | |
| 601 | const std::string& table_name, | ||
| 602 | const EmbeddingTableConfig& config, | ||
| 603 | const framework::CacheUpdateActions& cache) { | ||
| 604 | 10 | const auto tag = InitEmbeddingTable(table_name, config); | |
| 605 | 10 | framework::CacheAwareOp::AfterReload(tag >= 0, cache); | |
| 606 | 10 | return tag; | |
| 607 | } | ||
| 608 | |||
| 609 | ✗ | bool KVClientOp::LoadCheckpointCached( | |
| 610 | const std::string& path, | ||
| 611 | const std::string& metadata, | ||
| 612 | const framework::CacheUpdateActions& cache) { | ||
| 613 | ✗ | const bool loaded = LoadCheckpoint(path, metadata); | |
| 614 | ✗ | framework::CacheAwareOp::AfterReload(loaded, cache); | |
| 615 | ✗ | return loaded; | |
| 616 | } | ||
| 617 | |||
| 618 | 18 | int KVClientOp::InitEmbeddingTable(const std::string& table_name, | |
| 619 | const EmbeddingTableConfig& config) { | ||
| 620 | 18 | if (IsHierKVBackendName(ps_backend_name_)) { | |
| 621 | 4 | if (!GetHierKVLocalRuntime().InitEmbeddingTable(table_name, config)) { | |
| 622 | ✗ | return -1; | |
| 623 | } | ||
| 624 | 4 | return static_cast<int>(config.table_id); | |
| 625 | } | ||
| 626 | 14 | if (ps_client_ == nullptr) { | |
| 627 | ✗ | throw std::runtime_error("PS client is not initialized. Please call " | |
| 628 | ✗ | "KVClientOp::SetPSClient() first."); | |
| 629 | } | ||
| 630 | |||
| 631 | # ifdef ENABLE_PERF_REPORT | ||
| 632 | auto start_time = std::chrono::high_resolution_clock::now(); | ||
| 633 | # endif | ||
| 634 | 14 | int ret = ps_client_->InitEmbeddingTable(table_name, config); | |
| 635 | # ifdef ENABLE_PERF_REPORT | ||
| 636 | auto end_time = std::chrono::high_resolution_clock::now(); | ||
| 637 | auto duration = | ||
| 638 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 639 | end_time - start_time) | ||
| 640 | .count(); | ||
| 641 | double start_us = | ||
| 642 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 643 | start_time.time_since_epoch()) | ||
| 644 | .count(); | ||
| 645 | std::string op_latency_key = | ||
| 646 | "InitEmbeddingTable|" + std::to_string(static_cast<uint64_t>(start_us)); | ||
| 647 | // report(table_name, key, metric_name, value) | ||
| 648 | report("op_latency", | ||
| 649 | op_latency_key.c_str(), | ||
| 650 | "recstore_us", | ||
| 651 | static_cast<double>(duration)); | ||
| 652 | # endif | ||
| 653 | 14 | return ret; | |
| 654 | } | ||
| 655 | |||
| 656 | 36 | void KVClientOp::EmbWrite(const RecTensor& keys, const RecTensor& values) { | |
| 657 | 36 | if (IsHierKVBackendName(ps_backend_name_)) { | |
| 658 | 14 | GetHierKVLocalRuntime().Write(keys, values); | |
| 659 | 6 | return; | |
| 660 | } | ||
| 661 | 22 | if (ps_client_ == nullptr) { | |
| 662 | ✗ | throw std::runtime_error("PS client is not initialized. Please call " | |
| 663 | ✗ | "KVClientOp::SetPSClient() first."); | |
| 664 | } | ||
| 665 | |||
| 666 | # ifdef ENABLE_PERF_REPORT | ||
| 667 | auto start_time = std::chrono::high_resolution_clock::now(); | ||
| 668 | # endif | ||
| 669 | |||
| 670 | 22 | validate_keys(keys); | |
| 671 | 22 | validate_embeddings(values, "Values"); | |
| 672 | |||
| 673 | 22 | const int64_t L = keys.shape(0); | |
| 674 | 22 | const auto& values_shape = values.shape(); | |
| 675 | 22 | if (values.shape(0) != L) { | |
| 676 | ✗ | throw std::invalid_argument( | |
| 677 | ✗ | "Dimension mismatch: Keys has length " + std::to_string(L) + | |
| 678 | ✗ | " but values has length " + std::to_string(values.shape(0))); | |
| 679 | } | ||
| 680 | 22 | const int64_t D = values.shape(1); | |
| 681 | |||
| 682 | 22 | const uint64_t* keys_data = keys.data_as<uint64_t>(); | |
| 683 | 22 | base::ConstArray<uint64_t> keys_array(keys_data, L); | |
| 684 | |||
| 685 | 22 | const int64_t total_values = L * D; | |
| 686 | 22 | if (values_shape[0] * values_shape[1] != total_values) { | |
| 687 | ✗ | throw std::invalid_argument( | |
| 688 | ✗ | "Values total elements mismatch: expected " + | |
| 689 | ✗ | std::to_string(total_values) + ", but got " + | |
| 690 | ✗ | std::to_string(values_shape[0] * values_shape[1])); | |
| 691 | } | ||
| 692 | |||
| 693 | 22 | if (D <= 0) { | |
| 694 | ✗ | throw std::invalid_argument( | |
| 695 | ✗ | "Invalid embedding dimension D: " + std::to_string(D)); | |
| 696 | } | ||
| 697 | |||
| 698 | 22 | int ret = ps_client_->PutParameter(keys_array, values); | |
| 699 | 22 | if (!IsReadWriteSuccess(ps_client_, ret)) { | |
| 700 | 2 | throw std::runtime_error("Failed to write embeddings to PS client."); | |
| 701 | } | ||
| 702 | |||
| 703 | # ifdef ENABLE_PERF_REPORT | ||
| 704 | auto end_time = std::chrono::high_resolution_clock::now(); | ||
| 705 | auto duration = | ||
| 706 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 707 | end_time - start_time) | ||
| 708 | .count(); | ||
| 709 | double start_us = | ||
| 710 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 711 | start_time.time_since_epoch()) | ||
| 712 | .count(); | ||
| 713 | std::string op_latency_key = | ||
| 714 | "EmbWrite|" + std::to_string(static_cast<uint64_t>(start_us)); | ||
| 715 | report("op_latency", | ||
| 716 | op_latency_key.c_str(), | ||
| 717 | "recstore_us", | ||
| 718 | static_cast<double>(duration)); | ||
| 719 | # endif | ||
| 720 | } | ||
| 721 | |||
| 722 | 4 | void KVClientOp::EmbInit(const base::RecTensor& keys, | |
| 723 | const InitStrategy& strategy) { | ||
| 724 | 4 | validate_keys(keys); | |
| 725 | 2 | } | |
| 726 | |||
| 727 | uint64_t | ||
| 728 | 10 | KVClientOp::EmbPrefetch(const base::RecTensor& keys, const RecTensor& values) { | |
| 729 | 10 | if (IsHierKVBackendName(ps_backend_name_)) { | |
| 730 | 6 | const int64_t embedding_dim = values.dim() == 2 ? values.shape(1) : 0; | |
| 731 | 6 | return GetHierKVLocalRuntime().Prefetch(keys, embedding_dim); | |
| 732 | } | ||
| 733 | 4 | const uint64_t* keys_data = keys.data_as<uint64_t>(); | |
| 734 | 4 | int64_t L = keys.shape(0); | |
| 735 | 4 | base::ConstArray<uint64_t> keys_array(keys_data, L); | |
| 736 | 4 | return ps_client_->PrefetchParameter(keys_array); | |
| 737 | } | ||
| 738 | |||
| 739 | 6 | bool KVClientOp::IsPrefetchDone(uint64_t prefetch_id) { | |
| 740 | 6 | if (IsHierKVBackendName(ps_backend_name_)) { | |
| 741 | 4 | return GetHierKVLocalRuntime().IsPrefetchDone(prefetch_id); | |
| 742 | } | ||
| 743 | 2 | return ps_client_->IsPrefetchDone(prefetch_id); | |
| 744 | } | ||
| 745 | |||
| 746 | 6 | void KVClientOp::WaitForPrefetch(uint64_t prefetch_id) { | |
| 747 | 6 | if (IsHierKVBackendName(ps_backend_name_)) { | |
| 748 | 2 | GetHierKVLocalRuntime().WaitForPrefetch(prefetch_id); | |
| 749 | 2 | return; | |
| 750 | } | ||
| 751 | 4 | ps_client_->WaitForPrefetch(prefetch_id); | |
| 752 | } | ||
| 753 | |||
| 754 | 12 | void KVClientOp::GetPretchResult(uint64_t prefetch_id, RecTensor& values) { | |
| 755 | 12 | if (IsHierKVBackendName(ps_backend_name_)) { | |
| 756 | 6 | GetHierKVLocalRuntime().ConsumePrefetch(prefetch_id, values); | |
| 757 | 6 | return; | |
| 758 | } | ||
| 759 | 6 | if (!ps_client_->GetPrefetchResult(prefetch_id, values)) { | |
| 760 | 2 | throw std::runtime_error("Failed to get prefetch result from PS client."); | |
| 761 | } | ||
| 762 | } | ||
| 763 | |||
| 764 | 2 | bool KVClientOp::IsWriteDone(uint64_t write_id) { | |
| 765 | // return ps_client_->IsWriteDone(write_id); | ||
| 766 | 2 | throw std::runtime_error("Not impl"); | |
| 767 | } | ||
| 768 | |||
| 769 | namespace testing {} // namespace testing | ||
| 770 | |||
| 771 | } // namespace recstore | ||
| 772 | |||
| 773 | #else | ||
| 774 | |||
| 775 | # include "common/op_mock.cc" | ||
| 776 | |||
| 777 | #endif // USE_FAKE_KVCLIENT | ||
| 778 |