ps/base/cache_ps_impl.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | |||
| 3 | #include <algorithm> | ||
| 4 | #include <atomic> | ||
| 5 | #include <boost/coroutine2/all.hpp> | ||
| 6 | #include <cstring> | ||
| 7 | #include <cstdint> | ||
| 8 | #include <exception> | ||
| 9 | #include <experimental/filesystem> | ||
| 10 | #include <mutex> | ||
| 11 | #include <random> | ||
| 12 | #include <stdexcept> | ||
| 13 | #include <unordered_map> | ||
| 14 | #include <vector> | ||
| 15 | |||
| 16 | #include "base/array.h" | ||
| 17 | #include "base/factory.h" | ||
| 18 | #include "base/log.h" // NOLINT | ||
| 19 | #include "base/timer.h" | ||
| 20 | #include "parameters.h" | ||
| 21 | #include "storage/kv_engine/base_kv.h" | ||
| 22 | #include "storage/kv_engine/engine_factory.h" | ||
| 23 | #include "storage/kv_engine/engine_selector.h" | ||
| 24 | #include "optimizer/optimizer.h" | ||
| 25 | #include "ps/local_shm/local_shm_stage_report.h" | ||
| 26 | #include "ps/base/sparse_gradient_commit.h" | ||
| 27 | |||
| 28 | #ifdef ENABLE_PERF_REPORT | ||
| 29 | # include <chrono> | ||
| 30 | # include "base/report/report_client.h" | ||
| 31 | #endif | ||
| 32 | |||
| 33 | using boost::coroutines2::coroutine; | ||
| 34 | |||
| 35 | static const int KEY_CNT = 12543670; | ||
| 36 | |||
| 37 | template <typename key_t> | ||
| 38 | struct TaskElement { | ||
| 39 | TaskElement(const base::ConstArray<key_t>& keys, | ||
| 40 | const base::MutableArray<ParameterPack>& packs, | ||
| 41 | std::atomic_bool* promise) | ||
| 42 | : keys(keys), packs(packs), promise(promise) {} | ||
| 43 | |||
| 44 | TaskElement() {} | ||
| 45 | |||
| 46 | base::ConstArray<key_t> keys; | ||
| 47 | base::MutableArray<ParameterPack> packs; | ||
| 48 | std::atomic_bool* promise; | ||
| 49 | }; | ||
| 50 | |||
| 51 | class CachePS { | ||
| 52 | public: | ||
| 53 | using key_t = uint64_t; | ||
| 54 | |||
| 55 | struct FlatGetProfile { | ||
| 56 | std::uint64_t batch_get_ns = 0; | ||
| 57 | std::uint64_t index_lookup_ns = 0; | ||
| 58 | std::uint64_t zero_fill_ns = 0; | ||
| 59 | std::uint64_t row_copy_ns = 0; | ||
| 60 | std::uint64_t rows = 0; | ||
| 61 | std::uint64_t value_bytes = 0; | ||
| 62 | std::uint64_t missing_rows = 0; | ||
| 63 | }; | ||
| 64 | |||
| 65 | using DirectFixedRow = BaseKV::DirectFixedRow; | ||
| 66 | using RDMABackingRegion = BaseKV::RDMABackingRegion; | ||
| 67 | |||
| 68 | 34 | CachePS(json config) { | |
| 69 | 34 | LOG(INFO) << "cache ps config: " << config.dump(2); | |
| 70 | 34 | BaseKVConfig kv_config; | |
| 71 | 34 | kv_config.num_threads_ = config["num_threads"].get<int>(); | |
| 72 | 34 | kv_config.json_config_ = config["base_kv_config"]; | |
| 73 | 34 | auto r = base::ResolveEngine(kv_config); | |
| 74 | 34 | base_kv_.reset(base::Factory<BaseKV, const BaseKVConfig&>::NewInstance( | |
| 75 | r.engine, r.cfg)); | ||
| 76 | 34 | optimizer_ = CreateOptimizer(config.value("optimizer", json::object())); | |
| 77 | 34 | } | |
| 78 | |||
| 79 | 34 | ~CachePS() {} | |
| 80 | |||
| 81 | ✗ | bool Initialize(const std::vector<std::string>& model_config_path, | |
| 82 | const std::vector<std::string>& emb_file_path) { | ||
| 83 | ✗ | LOG(INFO) << "Before Load CKPT"; | |
| 84 | ✗ | LoadCkpt(model_config_path, emb_file_path); | |
| 85 | ✗ | LOG(INFO) << "After Load CKPT"; | |
| 86 | ✗ | return true; | |
| 87 | } | ||
| 88 | |||
| 89 | ✗ | void Clear() { | |
| 90 | ✗ | std::lock_guard<std::mutex> lock(checkpoint_mu_); | |
| 91 | ✗ | base_kv_->clear(); | |
| 92 | ✗ | active_checkpoint_identity_.clear(); | |
| 93 | ✗ | checkpoint_dirty_.store( | |
| 94 | ✗ | base_kv_->CheckpointRecordCount() != 0, std::memory_order_relaxed); | |
| 95 | ✗ | } | |
| 96 | |||
| 97 | void LoadFakeData(int64_t key_capacity, int value_size) { | ||
| 98 | std::lock_guard<std::mutex> lock(checkpoint_mu_); | ||
| 99 | checkpoint_dirty_.store(true, std::memory_order_relaxed); | ||
| 100 | base_kv_->LoadFakeData(key_capacity, value_size); | ||
| 101 | } | ||
| 102 | |||
| 103 | ✗ | bool LoadCkpt(const std::vector<std::string>& model_config_path, | |
| 104 | const std::vector<std::string>& emb_file_path) { | ||
| 105 | // base_kv_->loadCkpt(); | ||
| 106 | // LoadFakeData(KEY_CNT); | ||
| 107 | ✗ | return true; | |
| 108 | } | ||
| 109 | |||
| 110 | ✗ | bool SaveCheckpoint(const std::string& path, const std::string& metadata) { | |
| 111 | ✗ | const std::string identity = CheckpointIdentity(metadata); | |
| 112 | ✗ | std::lock_guard<std::mutex> lock(checkpoint_mu_); | |
| 113 | ✗ | if (!base_kv_->SaveCheckpoint(path, metadata)) { | |
| 114 | ✗ | LOG(ERROR) << "Failed to save checkpoint: " << path; | |
| 115 | ✗ | return false; | |
| 116 | } | ||
| 117 | ✗ | active_checkpoint_identity_ = identity; | |
| 118 | ✗ | checkpoint_dirty_.store(false, std::memory_order_relaxed); | |
| 119 | ✗ | return true; | |
| 120 | ✗ | } | |
| 121 | |||
| 122 | ✗ | bool LoadCheckpoint(const std::string& path, | |
| 123 | const std::string& expected_metadata) { | ||
| 124 | ✗ | const std::string identity = CheckpointIdentity(expected_metadata); | |
| 125 | ✗ | std::lock_guard<std::mutex> lock(checkpoint_mu_); | |
| 126 | ✗ | if (!checkpoint_dirty_.load(std::memory_order_relaxed) && | |
| 127 | ✗ | active_checkpoint_identity_ == identity) { | |
| 128 | ✗ | return true; | |
| 129 | } | ||
| 130 | ✗ | if (checkpoint_dirty_.load(std::memory_order_relaxed)) { | |
| 131 | ✗ | throw std::runtime_error( | |
| 132 | ✗ | "checkpoint load rejected: parameter server has unsaved updates"); | |
| 133 | } | ||
| 134 | ✗ | if (!active_checkpoint_identity_.empty()) { | |
| 135 | ✗ | throw std::runtime_error( | |
| 136 | ✗ | "checkpoint load rejected: active checkpoint identity mismatch"); | |
| 137 | } | ||
| 138 | ✗ | if (base_kv_->CheckpointRecordCount() != 0) { | |
| 139 | ✗ | throw std::runtime_error( | |
| 140 | ✗ | "checkpoint load rejected: parameter server is not fresh"); | |
| 141 | } | ||
| 142 | try { | ||
| 143 | ✗ | if (!base_kv_->LoadCheckpoint(path, expected_metadata)) { | |
| 144 | ✗ | checkpoint_dirty_.store( | |
| 145 | ✗ | base_kv_->CheckpointRecordCount() != 0, std::memory_order_relaxed); | |
| 146 | ✗ | LOG(ERROR) << "Failed to load checkpoint: " << path; | |
| 147 | ✗ | return false; | |
| 148 | } | ||
| 149 | ✗ | } catch (...) { | |
| 150 | ✗ | checkpoint_dirty_.store( | |
| 151 | ✗ | base_kv_->CheckpointRecordCount() != 0, std::memory_order_relaxed); | |
| 152 | ✗ | throw; | |
| 153 | ✗ | } | |
| 154 | ✗ | active_checkpoint_identity_ = identity; | |
| 155 | ✗ | checkpoint_dirty_.store(false, std::memory_order_relaxed); | |
| 156 | ✗ | return true; | |
| 157 | ✗ | } | |
| 158 | |||
| 159 | void PutSingleParameter( | ||
| 160 | const uint64_t key, const void* data, const int dim, const int tid) { | ||
| 161 | base_kv_->Put(key, std::string_view((char*)data, dim * sizeof(float)), tid); | ||
| 162 | MarkCheckpointDirty(); | ||
| 163 | } | ||
| 164 | |||
| 165 | 44 | void PutDenseParameterBatch( | |
| 166 | const uint64_t* keys, | ||
| 167 | const float* values, | ||
| 168 | int key_count, | ||
| 169 | int embedding_dim, | ||
| 170 | int tid) { | ||
| 171 | 44 | if (key_count <= 0 || embedding_dim <= 0) { | |
| 172 | ✗ | return; | |
| 173 | } | ||
| 174 | 44 | base::ConstArray<uint64_t> key_array(keys, key_count); | |
| 175 | 44 | std::vector<base::ConstArray<float>> value_slices; | |
| 176 | 44 | value_slices.reserve(static_cast<std::size_t>(key_count)); | |
| 177 | 284 | for (int i = 0; i < key_count; ++i) { | |
| 178 | 240 | value_slices.emplace_back(values + i * embedding_dim, embedding_dim); | |
| 179 | } | ||
| 180 | 44 | base_kv_->BatchPut(key_array, &value_slices, tid); | |
| 181 | 44 | MarkCheckpointDirty(); | |
| 182 | 44 | } | |
| 183 | |||
| 184 | ✗ | void PutSingleParameter(const ParameterCompressItem* item, int tid) { | |
| 185 | ✗ | auto key = item->key; | |
| 186 | ✗ | auto dim = item->dim; | |
| 187 | ✗ | base_kv_->Put( | |
| 188 | ✗ | key, std::string_view((char*)item->data(), dim * sizeof(float)), tid); | |
| 189 | ✗ | MarkCheckpointDirty(); | |
| 190 | ✗ | } | |
| 191 | |||
| 192 | void PutParameter(coroutine<void>::push_type& sink, | ||
| 193 | const ParameterCompressReader* reader, | ||
| 194 | int tid) { | ||
| 195 | std::vector<uint64_t> keys_vec; | ||
| 196 | std::vector<base::ConstArray<float>> values; | ||
| 197 | for (int i = 0; i < reader->item_size(); i++) { | ||
| 198 | keys_vec.emplace_back(reader->item(i)->key); | ||
| 199 | values.emplace_back( | ||
| 200 | (float*)reader->item(i)->data(), reader->item(i)->dim); | ||
| 201 | } | ||
| 202 | base::ConstArray<uint64_t> keys(keys_vec); | ||
| 203 | |||
| 204 | base_kv_->BatchPut(sink, keys, &values, tid); | ||
| 205 | MarkCheckpointDirty(); | ||
| 206 | } | ||
| 207 | |||
| 208 | ✗ | void PutParameter(const ParameterCompressReader* reader, int tid) { | |
| 209 | ✗ | std::vector<uint64_t> keys_vec; | |
| 210 | ✗ | std::vector<base::ConstArray<float>> values; | |
| 211 | ✗ | for (int i = 0; i < reader->item_size(); i++) { | |
| 212 | ✗ | keys_vec.emplace_back(reader->item(i)->key); | |
| 213 | ✗ | values.emplace_back( | |
| 214 | ✗ | (float*)reader->item(i)->data(), reader->item(i)->dim); | |
| 215 | } | ||
| 216 | ✗ | base::ConstArray<uint64_t> keys(keys_vec); | |
| 217 | |||
| 218 | ✗ | base_kv_->BatchPut(keys, &values, tid); | |
| 219 | ✗ | MarkCheckpointDirty(); | |
| 220 | ✗ | } | |
| 221 | |||
| 222 | bool GetParameterRun2Completion(key_t key, ParameterPack& pack, int tid) { | ||
| 223 | std::vector<uint64_t> keys = {key}; | ||
| 224 | base::ConstArray<uint64_t> keys_array(keys); | ||
| 225 | std::vector<base::ConstArray<float>> values; | ||
| 226 | |||
| 227 | base_kv_->BatchGet(keys_array, &values, tid); | ||
| 228 | base::ConstArray<float> value = values[0]; | ||
| 229 | |||
| 230 | if (value.Size() == 0) { | ||
| 231 | pack.key = key; | ||
| 232 | pack.dim = 0; | ||
| 233 | pack.emb_data = nullptr; | ||
| 234 | RECSTORE_LOG_EVERY_MS(ERROR, 1000) << "key " << key << " not existing"; | ||
| 235 | return false; | ||
| 236 | } | ||
| 237 | pack.key = key; | ||
| 238 | pack.dim = value.Size(); | ||
| 239 | pack.emb_data = value.Data(); | ||
| 240 | // LOG(ERROR) << "Get key " << key << " dim " << pack.dim; | ||
| 241 | return true; | ||
| 242 | } | ||
| 243 | |||
| 244 | ✗ | bool GetParameterRun2Completion(base::ConstArray<uint64_t> keys, | |
| 245 | std::vector<ParameterPack>& packs, | ||
| 246 | int tid) { | ||
| 247 | #ifdef ENABLE_PERF_REPORT | ||
| 248 | auto start_time = std::chrono::high_resolution_clock::now(); | ||
| 249 | #endif | ||
| 250 | ✗ | const auto batch_get_start = std::chrono::steady_clock::now(); | |
| 251 | ✗ | std::vector<base::ConstArray<float>> values; | |
| 252 | ✗ | base_kv_->BatchGet(keys, &values, tid); | |
| 253 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 254 | "cache_ps_get_batch_get_us", | ||
| 255 | recstore::LocalShmElapsedUs(batch_get_start)); | ||
| 256 | |||
| 257 | ✗ | const auto pack_build_start = std::chrono::steady_clock::now(); | |
| 258 | ✗ | for (int i = 0; i < keys.Size(); i++) { | |
| 259 | ✗ | packs.emplace_back(keys[i], values[i].Size(), values[i].Data()); | |
| 260 | } | ||
| 261 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 262 | "cache_ps_get_pack_us", recstore::LocalShmElapsedUs(pack_build_start)); | ||
| 263 | |||
| 264 | #ifdef ENABLE_PERF_REPORT | ||
| 265 | auto end_time = std::chrono::high_resolution_clock::now(); | ||
| 266 | double start_us = | ||
| 267 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 268 | start_time.time_since_epoch()) | ||
| 269 | .count(); | ||
| 270 | auto duration = | ||
| 271 | std::chrono::duration_cast<std::chrono::microseconds>( | ||
| 272 | end_time - start_time) | ||
| 273 | .count(); | ||
| 274 | |||
| 275 | std::string report_id = "cache_ps::GetParameterRun2Completion|" + | ||
| 276 | std::to_string(static_cast<uint64_t>(start_us)); | ||
| 277 | |||
| 278 | report("embread_stages", | ||
| 279 | report_id.c_str(), | ||
| 280 | "duration_us", | ||
| 281 | static_cast<double>(duration)); | ||
| 282 | |||
| 283 | report("embread_stages", | ||
| 284 | report_id.c_str(), | ||
| 285 | "request_size", | ||
| 286 | static_cast<double>(keys.Size())); | ||
| 287 | |||
| 288 | std::string unique_id = | ||
| 289 | "embread_debug|" + std::to_string(static_cast<uint64_t>(start_us)); | ||
| 290 | FlameGraphData fg_data = { | ||
| 291 | "cache_ps::GetParameterRun2Completion", | ||
| 292 | start_us, | ||
| 293 | 3, // level | ||
| 294 | static_cast<double>(duration), | ||
| 295 | static_cast<double>(duration)}; | ||
| 296 | report_flame_graph("emb_read_flame_map", unique_id.c_str(), fg_data); | ||
| 297 | #endif | ||
| 298 | ✗ | return true; | |
| 299 | ✗ | } | |
| 300 | |||
| 301 | 48 | bool GetParameterFlat( | |
| 302 | base::ConstArray<uint64_t> keys, | ||
| 303 | float* values, | ||
| 304 | int64_t num_rows, | ||
| 305 | int64_t embedding_dim, | ||
| 306 | int tid, | ||
| 307 | FlatGetProfile* profile = nullptr) { | ||
| 308 | 48 | if (values == nullptr) { | |
| 309 | ✗ | LOG(ERROR) << "GetParameterFlat values pointer is null"; | |
| 310 | ✗ | return false; | |
| 311 | } | ||
| 312 | 48 | if (num_rows < 0 || embedding_dim <= 0) { | |
| 313 | ✗ | LOG(ERROR) << "GetParameterFlat invalid shape rows=" << num_rows | |
| 314 | ✗ | << " dim=" << embedding_dim; | |
| 315 | ✗ | return false; | |
| 316 | } | ||
| 317 | 48 | if (keys.Size() != static_cast<size_t>(num_rows)) { | |
| 318 | ✗ | LOG(ERROR) << "GetParameterFlat keys size mismatch " << keys.Size() | |
| 319 | ✗ | << " vs " << num_rows; | |
| 320 | ✗ | return false; | |
| 321 | } | ||
| 322 | |||
| 323 | 48 | const auto batch_get_start = std::chrono::steady_clock::now(); | |
| 324 | 48 | BaseKV::BatchGetFlatStats flat_stats; | |
| 325 | 48 | BaseKV::BatchGetFlatStats* flat_stats_ptr = | |
| 326 | 48 | profile != nullptr ? &flat_stats : nullptr; | |
| 327 | 48 | const bool flat_ok = base_kv_->BatchGetFlat( | |
| 328 | keys, values, num_rows, embedding_dim, tid, flat_stats_ptr); | ||
| 329 | 48 | if (flat_ok) { | |
| 330 | 46 | if (profile != nullptr) { | |
| 331 | ✗ | profile->batch_get_ns = static_cast<std::uint64_t>( | |
| 332 | ✗ | std::chrono::duration_cast< std::chrono::nanoseconds>( | |
| 333 | ✗ | std::chrono::steady_clock::now() - batch_get_start) | |
| 334 | ✗ | .count()); | |
| 335 | ✗ | profile->rows = static_cast<std::uint64_t>(num_rows); | |
| 336 | ✗ | profile->value_bytes = | |
| 337 | ✗ | static_cast<std::uint64_t>(num_rows) * | |
| 338 | ✗ | static_cast<std::uint64_t>(embedding_dim) * sizeof(float); | |
| 339 | ✗ | profile->zero_fill_ns = flat_stats.zero_fill_ns; | |
| 340 | ✗ | profile->index_lookup_ns = flat_stats.index_lookup_ns; | |
| 341 | ✗ | profile->row_copy_ns = flat_stats.row_copy_ns; | |
| 342 | ✗ | profile->missing_rows = flat_stats.missing_rows; | |
| 343 | } | ||
| 344 | 46 | recstore::ReportLocalShmStageMetric( | |
| 345 | "cache_ps_get_batch_get_us", | ||
| 346 | recstore::LocalShmElapsedUs(batch_get_start)); | ||
| 347 | 46 | recstore::ReportLocalShmStageMetric("cache_ps_get_copy_us", 0); | |
| 348 | 46 | return true; | |
| 349 | } | ||
| 350 | |||
| 351 | 2 | std::vector<base::ConstArray<float>> value_slices; | |
| 352 | 2 | value_slices.reserve(static_cast<std::size_t>(num_rows)); | |
| 353 | 2 | base_kv_->BatchGet(keys, &value_slices, tid); | |
| 354 | 2 | if (profile != nullptr) { | |
| 355 | ✗ | profile->batch_get_ns = static_cast<std::uint64_t>( | |
| 356 | ✗ | std::chrono::duration_cast< std::chrono::nanoseconds>( | |
| 357 | ✗ | std::chrono::steady_clock::now() - batch_get_start) | |
| 358 | ✗ | .count()); | |
| 359 | ✗ | profile->rows = static_cast<std::uint64_t>(num_rows); | |
| 360 | ✗ | profile->value_bytes = | |
| 361 | ✗ | static_cast<std::uint64_t>(num_rows) * | |
| 362 | ✗ | static_cast<std::uint64_t>(embedding_dim) * sizeof(float); | |
| 363 | } | ||
| 364 | 2 | recstore::ReportLocalShmStageMetric( | |
| 365 | "cache_ps_get_batch_get_us", | ||
| 366 | recstore::LocalShmElapsedUs(batch_get_start)); | ||
| 367 | 2 | if (value_slices.size() != static_cast<size_t>(num_rows)) { | |
| 368 | ✗ | LOG(ERROR) << "GetParameterFlat BatchGet returned " << value_slices.size() | |
| 369 | ✗ | << " rows, expected " << num_rows; | |
| 370 | ✗ | return false; | |
| 371 | } | ||
| 372 | |||
| 373 | 2 | for (int64_t row = 0; row < num_rows; ++row) { | |
| 374 | 2 | const auto& slice = value_slices[static_cast<size_t>(row)]; | |
| 375 | 4 | if (slice.Size() != 0 && | |
| 376 | 2 | static_cast<int64_t>(slice.Size()) != embedding_dim) { | |
| 377 | 4 | LOG(ERROR) << "GetParameterFlat embedding_dim mismatch at row=" << row | |
| 378 | 2 | << " key=" << keys[static_cast<size_t>(row)] << " expected=" | |
| 379 | 2 | << embedding_dim << " actual=" << slice.Size(); | |
| 380 | 2 | return false; | |
| 381 | } | ||
| 382 | } | ||
| 383 | |||
| 384 | ✗ | const auto row_copy_start = std::chrono::steady_clock::now(); | |
| 385 | ✗ | std::uint64_t missing_zero_fill_ns = 0; | |
| 386 | ✗ | for (int64_t row = 0; row < num_rows; ++row) { | |
| 387 | ✗ | const auto& slice = value_slices[static_cast<size_t>(row)]; | |
| 388 | ✗ | if (slice.Size() > 0) { | |
| 389 | ✗ | std::memcpy(values + row * embedding_dim, | |
| 390 | ✗ | slice.Data(), | |
| 391 | ✗ | static_cast<std::size_t>(embedding_dim) * sizeof(float)); | |
| 392 | } else { | ||
| 393 | ✗ | const auto missing_zero_start = std::chrono::steady_clock::now(); | |
| 394 | ✗ | std::memset(values + row * embedding_dim, | |
| 395 | 0, | ||
| 396 | ✗ | static_cast<std::size_t>(embedding_dim) * sizeof(float)); | |
| 397 | ✗ | if (profile != nullptr) { | |
| 398 | ✗ | missing_zero_fill_ns += static_cast<std::uint64_t>( | |
| 399 | ✗ | std::chrono::duration_cast< std::chrono::nanoseconds>( | |
| 400 | ✗ | std::chrono::steady_clock::now() - missing_zero_start) | |
| 401 | ✗ | .count()); | |
| 402 | ✗ | ++profile->missing_rows; | |
| 403 | } | ||
| 404 | } | ||
| 405 | } | ||
| 406 | ✗ | if (profile != nullptr) { | |
| 407 | ✗ | profile->zero_fill_ns = missing_zero_fill_ns; | |
| 408 | ✗ | profile->row_copy_ns = static_cast<std::uint64_t>( | |
| 409 | ✗ | std::chrono::duration_cast< std::chrono::nanoseconds>( | |
| 410 | ✗ | std::chrono::steady_clock::now() - row_copy_start) | |
| 411 | ✗ | .count()); | |
| 412 | } | ||
| 413 | ✗ | recstore::ReportLocalShmStageMetric( | |
| 414 | "cache_ps_get_copy_us", recstore::LocalShmElapsedUs(row_copy_start)); | ||
| 415 | ✗ | return true; | |
| 416 | 2 | } | |
| 417 | |||
| 418 | ✗ | bool ProbeParameterIndex(base::ConstArray<uint64_t> keys, | |
| 419 | int tid, | ||
| 420 | FlatGetProfile* profile = nullptr) { | ||
| 421 | ✗ | const auto batch_get_start = std::chrono::steady_clock::now(); | |
| 422 | ✗ | BaseKV::BatchGetFlatStats flat_stats; | |
| 423 | ✗ | BaseKV::BatchGetFlatStats* flat_stats_ptr = | |
| 424 | ✗ | profile != nullptr ? &flat_stats : nullptr; | |
| 425 | ✗ | const bool ok = base_kv_->BatchGetIndexOnly(keys, tid, flat_stats_ptr); | |
| 426 | ✗ | if (profile != nullptr) { | |
| 427 | ✗ | profile->batch_get_ns = static_cast<std::uint64_t>( | |
| 428 | ✗ | std::chrono::duration_cast< std::chrono::nanoseconds>( | |
| 429 | ✗ | std::chrono::steady_clock::now() - batch_get_start) | |
| 430 | ✗ | .count()); | |
| 431 | ✗ | profile->rows = static_cast<std::uint64_t>(keys.Size()); | |
| 432 | ✗ | profile->value_bytes = 0; | |
| 433 | ✗ | profile->missing_rows = flat_stats.missing_rows; | |
| 434 | } | ||
| 435 | ✗ | return ok; | |
| 436 | } | ||
| 437 | |||
| 438 | ✗ | bool GetParameterDirectFixedRows( | |
| 439 | base::ConstArray<uint64_t> keys, | ||
| 440 | int64_t num_rows, | ||
| 441 | int64_t embedding_dim, | ||
| 442 | int tid, | ||
| 443 | std::vector<DirectFixedRow>* rows, | ||
| 444 | FlatGetProfile* profile = nullptr) { | ||
| 445 | ✗ | const auto batch_get_start = std::chrono::steady_clock::now(); | |
| 446 | ✗ | BaseKV::BatchGetFlatStats flat_stats; | |
| 447 | ✗ | BaseKV::BatchGetFlatStats* flat_stats_ptr = | |
| 448 | ✗ | profile != nullptr ? &flat_stats : nullptr; | |
| 449 | ✗ | const bool ok = base_kv_->BatchGetDirectFixedRows( | |
| 450 | keys, num_rows, embedding_dim, tid, rows, flat_stats_ptr); | ||
| 451 | ✗ | if (profile != nullptr) { | |
| 452 | ✗ | profile->batch_get_ns = static_cast<std::uint64_t>( | |
| 453 | ✗ | std::chrono::duration_cast< std::chrono::nanoseconds>( | |
| 454 | ✗ | std::chrono::steady_clock::now() - batch_get_start) | |
| 455 | ✗ | .count()); | |
| 456 | ✗ | profile->rows = static_cast<std::uint64_t>(num_rows); | |
| 457 | ✗ | profile->value_bytes = | |
| 458 | ✗ | static_cast<std::uint64_t>(num_rows) * | |
| 459 | ✗ | static_cast<std::uint64_t>(embedding_dim) * sizeof(float); | |
| 460 | ✗ | profile->zero_fill_ns = flat_stats.zero_fill_ns; | |
| 461 | ✗ | profile->index_lookup_ns = flat_stats.index_lookup_ns; | |
| 462 | ✗ | profile->row_copy_ns = 0; | |
| 463 | ✗ | profile->missing_rows = flat_stats.missing_rows; | |
| 464 | } | ||
| 465 | ✗ | return ok; | |
| 466 | } | ||
| 467 | |||
| 468 | ✗ | RDMABackingRegion GetRDMABackingRegion() const { | |
| 469 | ✗ | return base_kv_->GetRDMABackingRegion(); | |
| 470 | } | ||
| 471 | |||
| 472 | bool GetParameterRun2Completion( | ||
| 473 | coroutine<void>::push_type& sink, | ||
| 474 | base::ConstArray<uint64_t> keys, | ||
| 475 | std::vector<ParameterPack>& pack, | ||
| 476 | int tid) { | ||
| 477 | std::vector<base::ConstArray<float>> values; | ||
| 478 | |||
| 479 | base_kv_->BatchGet(sink, keys, &values, tid); | ||
| 480 | |||
| 481 | for (int i = 0; i < keys.Size(); i++) { | ||
| 482 | pack.emplace_back(keys[i], values[i].Size(), values[i].Data()); | ||
| 483 | } | ||
| 484 | return true; | ||
| 485 | } | ||
| 486 | |||
| 487 | /// optimizer interface | ||
| 488 | |||
| 489 | 8 | int64_t SparseGradientCommit( | |
| 490 | const std::string& table, std::string payload, unsigned tid) { | ||
| 491 | 8 | const auto h = recstore::DecodeSparseGradientCommit(payload); | |
| 492 | 8 | const auto* sgd = dynamic_cast<SGD*>(optimizer_.get()); | |
| 493 | 8 | if (!sgd || sgd->LearningRate() != h.learning_rate) { | |
| 494 | ✗ | throw std::runtime_error("Sparse gradient commit requires matching " | |
| 495 | ✗ | "client/PS SGD learning rates"); | |
| 496 | } | ||
| 497 | 8 | const auto config = table_configs_.find(table); | |
| 498 | 16 | if (config == table_configs_.end() || | |
| 499 | 8 | config->second.embedding_dim != h.dim) { | |
| 500 | ✗ | throw std::runtime_error( | |
| 501 | ✗ | "Sparse gradient commit table or dimension mismatch"); | |
| 502 | } | ||
| 503 | 24 | return sparse_gradient_commits_.Exchange( | |
| 504 | table, | ||
| 505 | 8 | std::move(payload), | |
| 506 | 2 | [&](const std::vector<uint64_t>& keys, | |
| 507 | const std::vector<float>& grads) { | ||
| 508 | 2 | if (!UpdateParameterFlat( | |
| 509 | 2 | table, | |
| 510 | 4 | base::ConstArray<uint64_t>(keys.data(), keys.size()), | |
| 511 | grads.data(), | ||
| 512 | 2 | keys.size(), | |
| 513 | 2 | h.dim, | |
| 514 | 2 | tid)) { | |
| 515 | ✗ | throw std::runtime_error("Sparse gradient commit optimizer failed"); | |
| 516 | } | ||
| 517 | 18 | }); | |
| 518 | } | ||
| 519 | |||
| 520 | 28 | int InitTable(const std::string& table_name, | |
| 521 | uint64_t num_embeddings, | ||
| 522 | uint64_t embedding_dim, | ||
| 523 | uint64_t table_id = 0) { | ||
| 524 | 28 | std::lock_guard<std::mutex> lock(checkpoint_mu_); | |
| 525 | 28 | EmbeddingTableConfig config{num_embeddings, embedding_dim, table_id}; | |
| 526 | 28 | const auto existing = table_configs_.find(table_name); | |
| 527 | 28 | if (existing != table_configs_.end()) { | |
| 528 | ✗ | const auto& old = existing->second; | |
| 529 | ✗ | if (old.num_embeddings == num_embeddings && | |
| 530 | ✗ | old.embedding_dim == embedding_dim && old.table_id == table_id) { | |
| 531 | ✗ | return optimizer_->Init({table_name}, config, base_kv_.get()); | |
| 532 | } | ||
| 533 | ✗ | LOG(ERROR) << "Embedding table config mismatch for '" << table_name | |
| 534 | ✗ | << "': existing=[" << old.num_embeddings << ", " | |
| 535 | ✗ | << old.embedding_dim << "] requested=[" << num_embeddings | |
| 536 | ✗ | << ", " << embedding_dim << "]"; | |
| 537 | ✗ | return -1; | |
| 538 | } | ||
| 539 | 56 | const int tag = optimizer_->Init({table_name}, config, base_kv_.get()); | |
| 540 | 28 | table_configs_.emplace(table_name, config); | |
| 541 | 28 | return tag; | |
| 542 | 28 | } | |
| 543 | |||
| 544 | ✗ | bool UpdateParameter(const std::string& table_name, | |
| 545 | const ParameterCompressReader* reader, | ||
| 546 | unsigned tid) { | ||
| 547 | ✗ | if (!optimizer_) { | |
| 548 | ✗ | LOG(ERROR) << "Optimizer not initialized. Please call InitTable first."; | |
| 549 | ✗ | return false; | |
| 550 | } | ||
| 551 | |||
| 552 | ✗ | optimizer_->Update(table_name, reader, tid); | |
| 553 | ✗ | MarkCheckpointDirty(); | |
| 554 | ✗ | return true; | |
| 555 | } | ||
| 556 | |||
| 557 | 8 | bool UpdateParameterFlat( | |
| 558 | const std::string& table_name, | ||
| 559 | const base::ConstArray<uint64_t>& keys, | ||
| 560 | const float* grads, | ||
| 561 | int64_t num_rows, | ||
| 562 | int64_t embedding_dim, | ||
| 563 | unsigned tid) { | ||
| 564 | 8 | if (grads == nullptr) { | |
| 565 | ✗ | LOG(ERROR) << "UpdateParameterFlat grads pointer is null"; | |
| 566 | ✗ | return false; | |
| 567 | } | ||
| 568 | 8 | if (num_rows < 0 || embedding_dim <= 0) { | |
| 569 | ✗ | LOG(ERROR) << "UpdateParameterFlat invalid shape rows=" << num_rows | |
| 570 | ✗ | << " dim=" << embedding_dim; | |
| 571 | ✗ | return false; | |
| 572 | } | ||
| 573 | 8 | if (keys.Size() != static_cast<size_t>(num_rows)) { | |
| 574 | ✗ | LOG(ERROR) << "UpdateParameterFlat keys size mismatch " << keys.Size() | |
| 575 | ✗ | << " vs " << num_rows; | |
| 576 | ✗ | return false; | |
| 577 | } | ||
| 578 | 8 | if (!optimizer_) { | |
| 579 | ✗ | LOG(ERROR) << "Optimizer not initialized. Please call InitTable first."; | |
| 580 | ✗ | return false; | |
| 581 | } | ||
| 582 | 8 | optimizer_->UpdateFlat( | |
| 583 | table_name, keys, grads, num_rows, embedding_dim, tid); | ||
| 584 | 8 | MarkCheckpointDirty(); | |
| 585 | 8 | return true; | |
| 586 | } | ||
| 587 | |||
| 588 | private: | ||
| 589 | recstore::SparseGradientCommitCoordinator sparse_gradient_commits_; | ||
| 590 | 52 | void MarkCheckpointDirty() { | |
| 591 | 52 | if (!checkpoint_dirty_.load(std::memory_order_relaxed)) { | |
| 592 | 26 | checkpoint_dirty_.store(true, std::memory_order_relaxed); | |
| 593 | } | ||
| 594 | 52 | } | |
| 595 | |||
| 596 | ✗ | static std::string CheckpointIdentity(const std::string& metadata) { | |
| 597 | ✗ | const json parsed = json::parse(metadata); | |
| 598 | ✗ | if (!parsed.is_object() || !parsed.contains("identity") || | |
| 599 | ✗ | !parsed["identity"].is_object() || !parsed.contains("checkpoint_id") || | |
| 600 | ✗ | !parsed["checkpoint_id"].is_string() || | |
| 601 | ✗ | parsed["checkpoint_id"].get<std::string>().empty()) { | |
| 602 | ✗ | throw std::invalid_argument( | |
| 603 | ✗ | "checkpoint metadata requires identity and non-empty checkpoint_id"); | |
| 604 | } | ||
| 605 | ✗ | return parsed.dump(); | |
| 606 | ✗ | } | |
| 607 | |||
| 608 | std::unique_ptr<BaseKV> base_kv_; | ||
| 609 | std::unique_ptr<Optimizer> optimizer_; | ||
| 610 | std::mutex checkpoint_mu_; | ||
| 611 | std::unordered_map<std::string, EmbeddingTableConfig> table_configs_; | ||
| 612 | std::string active_checkpoint_identity_; | ||
| 613 | std::atomic<bool> checkpoint_dirty_{false}; | ||
| 614 | std::atomic<bool> stopFlag_{false}; | ||
| 615 | }; | ||
| 616 |