ps/rdma/rdma_ps_client_adapter.cc
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #include "ps/rdma/rdma_ps_client_adapter.h" | ||
| 2 | |||
| 3 | #include <algorithm> | ||
| 4 | #include <atomic> | ||
| 5 | #include <chrono> | ||
| 6 | #include <cstddef> | ||
| 7 | #include <cstdint> | ||
| 8 | #include <cstdlib> | ||
| 9 | #include <cstring> | ||
| 10 | #include <deque> | ||
| 11 | #include <fstream> | ||
| 12 | #include <limits> | ||
| 13 | #include <memory> | ||
| 14 | #include <string> | ||
| 15 | #include <stdexcept> | ||
| 16 | #include <thread> | ||
| 17 | #include <utility> | ||
| 18 | |||
| 19 | #include <folly/portability/GFlags.h> | ||
| 20 | #include <folly/init/Init.h> | ||
| 21 | |||
| 22 | #include "framework/common/ps_client_config_adapter.h" | ||
| 23 | #include "optimizer/sparse_tensor.h" | ||
| 24 | #include "ps/base/config.h" | ||
| 25 | #include "ps/base/parameters.h" | ||
| 26 | #include "ps/rdma/rdma_common.h" | ||
| 27 | #include "ps/rdma/rc_options.h" | ||
| 28 | |||
| 29 | DECLARE_int32(global_id); | ||
| 30 | DECLARE_int32(num_server_processes); | ||
| 31 | DECLARE_int32(num_client_processes); | ||
| 32 | DECLARE_int32(value_size); | ||
| 33 | DECLARE_int32(max_kv_num_per_request); | ||
| 34 | DECLARE_int32(rdma_rc_client_id_base); | ||
| 35 | DECLARE_int32(rdma_rc_num_logical_clients); | ||
| 36 | DECLARE_int32(rdma_control_plane_timeout_ms); | ||
| 37 | DECLARE_string(rdma_get_response_mode); | ||
| 38 | DECLARE_string(rdma_transport_mode); | ||
| 39 | DEFINE_string(rdma_transport_mode, "rc_write", "RDMA transport mode: rc_write"); | ||
| 40 | DEFINE_bool(rdma_adapter_skip_prefetch_result_copy, | ||
| 41 | false, | ||
| 42 | "Benchmark-only option to skip copying RDMA prefetch results into " | ||
| 43 | "the GetPrefetchResult output tensor"); | ||
| 44 | |||
| 45 | namespace recstore { | ||
| 46 | |||
| 47 | namespace detail { | ||
| 48 | |||
| 49 | 42 | bool TryParseIntEnv(const char* env_name, int* parsed_value) { | |
| 50 | 42 | const char* value = std::getenv(env_name); | |
| 51 | 42 | if (value == nullptr || *value == '\0') { | |
| 52 | 26 | return false; | |
| 53 | } | ||
| 54 | 16 | char* end = nullptr; | |
| 55 | 16 | const long parsed = std::strtol(value, &end, 10); | |
| 56 | 16 | if (end == value || *end != '\0') { | |
| 57 | ✗ | return false; | |
| 58 | } | ||
| 59 | 16 | *parsed_value = static_cast<int>(parsed); | |
| 60 | 16 | return true; | |
| 61 | } | ||
| 62 | |||
| 63 | } // namespace detail | ||
| 64 | |||
| 65 | namespace { | ||
| 66 | |||
| 67 | ✗ | bool AdapterProfileEnabled() { | |
| 68 | ✗ | const char* value = std::getenv("RECSTORE_RDMA_ADAPTER_PROFILE"); | |
| 69 | ✗ | return value != nullptr && std::string(value) != "0"; | |
| 70 | } | ||
| 71 | |||
| 72 | 14 | void SetIntFlagFromEnv(const char* env_name, int32_t* flag_value) { | |
| 73 | 14 | int parsed = 0; | |
| 74 | 14 | if (detail::TryParseIntEnv(env_name, &parsed)) { | |
| 75 | ✗ | *flag_value = static_cast<int32_t>(parsed); | |
| 76 | } | ||
| 77 | 14 | } | |
| 78 | |||
| 79 | 2 | void ApplyRdmaFlagsFromEnv() { | |
| 80 | 2 | if (const char* value = std::getenv("RECSTORE_RDMA_RC_NAMESPACE")) { | |
| 81 | ✗ | FLAGS_rdma_rc_namespace = value; | |
| 82 | } | ||
| 83 | 2 | if (const char* value = std::getenv("RECSTORE_RDMA_CONTROL_PLANE_HOST")) { | |
| 84 | ✗ | FLAGS_rdma_control_plane_host = value; | |
| 85 | } | ||
| 86 | 2 | if (const char* value = std::getenv("RECSTORE_RDMA_GET_RESPONSE_MODE")) { | |
| 87 | 2 | const std::string mode(value); | |
| 88 | 2 | if (mode != "direct_sg" && mode != "staging_copy") { | |
| 89 | ✗ | throw std::runtime_error( | |
| 90 | ✗ | "RECSTORE_RDMA_GET_RESPONSE_MODE must be direct_sg or staging_copy"); | |
| 91 | } | ||
| 92 | 2 | FLAGS_rdma_get_response_mode = mode; | |
| 93 | 2 | } | |
| 94 | 2 | SetIntFlagFromEnv( | |
| 95 | "RECSTORE_RDMA_CONTROL_PLANE_PORT", &FLAGS_rdma_control_plane_port); | ||
| 96 | 2 | SetIntFlagFromEnv("RECSTORE_RDMA_CONTROL_PLANE_TIMEOUT_MS", | |
| 97 | &FLAGS_rdma_control_plane_timeout_ms); | ||
| 98 | 2 | SetIntFlagFromEnv( | |
| 99 | "RECSTORE_RDMA_WAIT_TIMEOUT_MS", &FLAGS_rdma_wait_timeout_ms); | ||
| 100 | 2 | SetIntFlagFromEnv("RECSTORE_RDMA_RC_QPS_PER_CLIENT_PER_SHARD", | |
| 101 | &FLAGS_rdma_rc_qps_per_client_per_shard); | ||
| 102 | 2 | SetIntFlagFromEnv( | |
| 103 | "RECSTORE_RDMA_RC_SLOTS_PER_QP", &FLAGS_rdma_rc_slots_per_qp); | ||
| 104 | 2 | SetIntFlagFromEnv("RECSTORE_RDMA_RC_SERVER_COROUTINES_PER_THREAD", | |
| 105 | &FLAGS_rdma_rc_server_coroutines_per_thread); | ||
| 106 | 2 | SetIntFlagFromEnv( | |
| 107 | "RECSTORE_RDMA_RC_SERVER_GET_WORKERS", &FLAGS_rdma_rc_server_get_workers); | ||
| 108 | 2 | } | |
| 109 | |||
| 110 | ✗ | std::int64_t NsSince(std::chrono::steady_clock::time_point start, | |
| 111 | std::chrono::steady_clock::time_point end) { | ||
| 112 | ✗ | return std::chrono::duration_cast<std::chrono::nanoseconds>(end - start) | |
| 113 | ✗ | .count(); | |
| 114 | } | ||
| 115 | |||
| 116 | ✗ | int ValueSizeHintFromBaseKvConfig(const json& base_kv_config, | |
| 117 | int fallback_value_size) { | ||
| 118 | ✗ | if (!base_kv_config.is_object()) { | |
| 119 | ✗ | return fallback_value_size; | |
| 120 | } | ||
| 121 | ✗ | if (!base_kv_config.contains("value") || | |
| 122 | ✗ | !base_kv_config["value"].is_object()) { | |
| 123 | ✗ | return fallback_value_size; | |
| 124 | } | ||
| 125 | ✗ | return base_kv_config["value"].value( | |
| 126 | ✗ | "default_value_size_hint", fallback_value_size); | |
| 127 | } | ||
| 128 | |||
| 129 | ✗ | std::vector<std::string> ReadProcessArgv() { | |
| 130 | ✗ | std::ifstream cmdline("/proc/self/cmdline", std::ios::binary); | |
| 131 | ✗ | std::vector<std::string> argv; | |
| 132 | ✗ | if (!cmdline.is_open()) { | |
| 133 | ✗ | return argv; | |
| 134 | } | ||
| 135 | |||
| 136 | ✗ | std::string current; | |
| 137 | ✗ | char ch = '\0'; | |
| 138 | ✗ | while (cmdline.get(ch)) { | |
| 139 | ✗ | if (ch == '\0') { | |
| 140 | ✗ | if (!current.empty()) { | |
| 141 | ✗ | argv.push_back(current); | |
| 142 | ✗ | current.clear(); | |
| 143 | } | ||
| 144 | ✗ | continue; | |
| 145 | } | ||
| 146 | ✗ | current.push_back(ch); | |
| 147 | } | ||
| 148 | ✗ | if (!current.empty()) { | |
| 149 | ✗ | argv.push_back(current); | |
| 150 | } | ||
| 151 | ✗ | return argv; | |
| 152 | ✗ | } | |
| 153 | } // namespace | ||
| 154 | |||
| 155 | EmbeddedRdmaClientIdentity | ||
| 156 | 8 | ResolveEmbeddedRdmaClientIdentity(int num_shards, int configured_num_clients) { | |
| 157 | 8 | if (num_shards <= 0) { | |
| 158 | ✗ | throw std::runtime_error("embedded RDMA num_shards must be positive"); | |
| 159 | } | ||
| 160 | 8 | if (configured_num_clients <= 0) { | |
| 161 | ✗ | throw std::runtime_error( | |
| 162 | ✗ | "embedded RDMA configured_num_clients must be positive"); | |
| 163 | } | ||
| 164 | |||
| 165 | 8 | int client_index = 0; | |
| 166 | 14 | if (!detail::TryParseIntEnv("RECSTORE_RDMA_OS_CLIENT_INDEX", &client_index) && | |
| 167 | 6 | !detail::TryParseIntEnv("RANK", &client_index)) { | |
| 168 | ✗ | detail::TryParseIntEnv("LOCAL_RANK", &client_index); | |
| 169 | } | ||
| 170 | |||
| 171 | 8 | int num_client_processes = 1; | |
| 172 | 8 | if (!detail::TryParseIntEnv( | |
| 173 | 14 | "RECSTORE_RDMA_NUM_CLIENT_PROCESSES", &num_client_processes) && | |
| 174 | 6 | !detail::TryParseIntEnv("WORLD_SIZE", &num_client_processes)) { | |
| 175 | ✗ | detail::TryParseIntEnv("LOCAL_WORLD_SIZE", &num_client_processes); | |
| 176 | } | ||
| 177 | |||
| 178 | 8 | if (client_index < 0) { | |
| 179 | ✗ | throw std::runtime_error("embedded RDMA client index must be non-negative"); | |
| 180 | } | ||
| 181 | 8 | if (num_client_processes <= 0) { | |
| 182 | ✗ | throw std::runtime_error( | |
| 183 | ✗ | "embedded RDMA num_client_processes must be positive"); | |
| 184 | } | ||
| 185 | 8 | if (client_index >= num_client_processes) { | |
| 186 | 2 | throw std::runtime_error( | |
| 187 | 4 | "embedded RDMA client index out of range for num_client_processes"); | |
| 188 | } | ||
| 189 | 6 | if (num_client_processes != configured_num_clients) { | |
| 190 | 2 | throw std::runtime_error("embedded RDMA num_client_processes must match " | |
| 191 | 4 | "rdma_deployment.num_clients"); | |
| 192 | } | ||
| 193 | |||
| 194 | 4 | EmbeddedRdmaClientIdentity identity; | |
| 195 | 4 | identity.client_index = client_index; | |
| 196 | 4 | identity.num_client_processes = num_client_processes; | |
| 197 | 4 | identity.global_id = num_shards + client_index; | |
| 198 | 4 | return identity; | |
| 199 | } | ||
| 200 | |||
| 201 | ✗ | std::vector<RDMAPSClientAdapter::ShardChunk> RDMAPSClientAdapter::BuildChunks( | |
| 202 | base::ConstArray<uint64_t> keys, std::size_t max_keys_per_rpc) const { | ||
| 203 | return shard_routing::BuildChunks( | ||
| 204 | keys, | ||
| 205 | ✗ | deployment_.num_shards, | |
| 206 | ✗ | deployment_.hash_method, | |
| 207 | ✗ | deployment_.shard_to_endpoint_index, | |
| 208 | ✗ | max_keys_per_rpc); | |
| 209 | } | ||
| 210 | |||
| 211 | 2 | void InitializeRdmaProcessRuntime() { | |
| 212 | static std::once_flag init_once; | ||
| 213 | 2 | std::call_once(init_once, []() { | |
| 214 | // Python entrypoints pass application CLI flags that are not gflags. | ||
| 215 | // Passing them to folly::init makes gflags abort before the RDMA client can | ||
| 216 | // start. | ||
| 217 | 8 | std::vector<std::string> argv_strings = {"recstore_rdma_client"}; | |
| 218 | 2 | std::vector<char*> argv_storage; | |
| 219 | 2 | argv_storage.reserve(argv_strings.size() + 1); | |
| 220 | 4 | for (auto& arg : argv_strings) { | |
| 221 | 2 | argv_storage.push_back(arg.data()); | |
| 222 | } | ||
| 223 | 2 | argv_storage.push_back(nullptr); | |
| 224 | |||
| 225 | 2 | int argc = static_cast<int>(argv_strings.size()); | |
| 226 | 2 | char** argv = argv_storage.data(); | |
| 227 | 2 | folly::init(&argc, &argv); | |
| 228 | 2 | ApplyRdmaFlagsFromEnv(); | |
| 229 | 2 | }); | |
| 230 | 2 | } | |
| 231 | |||
| 232 | 6 | RDMAPSClientAdapter::RDMAPSClientAdapter(json config) | |
| 233 | 6 | : BasePSClient(config), config_(std::move(config)) {} | |
| 234 | |||
| 235 | ✗ | void RDMAPSClientAdapter::EnsureClientInitialized() { | |
| 236 | ✗ | std::lock_guard<std::mutex> guard(init_mu_); | |
| 237 | ✗ | if (initialized_) { | |
| 238 | ✗ | return; | |
| 239 | } | ||
| 240 | |||
| 241 | const json cache_ps_cfg = | ||
| 242 | ✗ | config_.contains("cache_ps") ? config_["cache_ps"] : json::object(); | |
| 243 | const json client_cfg = | ||
| 244 | ✗ | config_.contains("client") ? config_["client"] : json::object(); | |
| 245 | ✗ | const json dist_cfg = ResolveFrameworkDistributedClientConfig(config_); | |
| 246 | |||
| 247 | ✗ | deployment_ = ParseResolvedRdmaDeploymentConfig(config_); | |
| 248 | |||
| 249 | ✗ | if (FLAGS_global_id < deployment_.num_shards) { | |
| 250 | ✗ | const auto identity = ResolveEmbeddedRdmaClientIdentity( | |
| 251 | deployment_.num_shards, deployment_.num_clients); | ||
| 252 | ✗ | FLAGS_num_server_processes = deployment_.num_shards; | |
| 253 | ✗ | FLAGS_num_client_processes = identity.num_client_processes; | |
| 254 | ✗ | FLAGS_global_id = identity.global_id; | |
| 255 | ✗ | if (FLAGS_rdma_rc_num_logical_clients < 0) { | |
| 256 | ✗ | FLAGS_rdma_rc_num_logical_clients = identity.num_client_processes; | |
| 257 | } | ||
| 258 | ✗ | if (FLAGS_rdma_rc_client_id_base < 0) { | |
| 259 | ✗ | FLAGS_rdma_rc_client_id_base = identity.client_index; | |
| 260 | } | ||
| 261 | } | ||
| 262 | ✗ | if (FLAGS_num_server_processes != deployment_.num_shards) { | |
| 263 | ✗ | throw std::runtime_error( | |
| 264 | ✗ | "RDMA num_server_processes must match distributed_client.num_shards"); | |
| 265 | } | ||
| 266 | ✗ | if (FLAGS_num_client_processes != deployment_.num_clients) { | |
| 267 | ✗ | throw std::runtime_error( | |
| 268 | ✗ | "RDMA num_client_processes must match rdma_deployment.num_clients"); | |
| 269 | } | ||
| 270 | ✗ | if (FLAGS_global_id < 0 || | |
| 271 | ✗ | FLAGS_global_id >= deployment_.num_shards + deployment_.num_clients) { | |
| 272 | ✗ | throw std::runtime_error( | |
| 273 | ✗ | "RDMA global_id is outside the resolved deployment node range"); | |
| 274 | } | ||
| 275 | ✗ | FLAGS_value_size = | |
| 276 | ✗ | cache_ps_cfg.contains("base_kv_config") | |
| 277 | ✗ | ? ValueSizeHintFromBaseKvConfig( | |
| 278 | ✗ | cache_ps_cfg["base_kv_config"], FLAGS_value_size) | |
| 279 | : FLAGS_value_size; | ||
| 280 | ✗ | FLAGS_max_kv_num_per_request = deployment_.max_keys_per_request; | |
| 281 | ✗ | if (const char* mode = std::getenv("RECSTORE_RDMA_TRANSPORT_MODE")) { | |
| 282 | ✗ | FLAGS_rdma_transport_mode = mode; | |
| 283 | } | ||
| 284 | |||
| 285 | const int logical_client_id = | ||
| 286 | ✗ | config_.value("rdma_logical_client_id", FLAGS_rdma_rc_client_id_base); | |
| 287 | |||
| 288 | ✗ | shard_clients_.clear(); | |
| 289 | ✗ | client_ = nullptr; | |
| 290 | |||
| 291 | ✗ | if (deployment_.num_shards <= 1) { | |
| 292 | ✗ | shard_clients_.push_back(std::make_unique<petps::PetPSClient>( | |
| 293 | ✗ | client_cfg.value("host", std::string("127.0.0.1")), | |
| 294 | ✗ | client_cfg.value("port", 25000), | |
| 295 | ✗ | client_cfg.value("shard", 0), | |
| 296 | logical_client_id)); | ||
| 297 | ✗ | client_ = shard_clients_.front().get(); | |
| 298 | } else { | ||
| 299 | ✗ | for (const auto& endpoint : deployment_.endpoints) { | |
| 300 | ✗ | shard_clients_.push_back(std::make_unique<petps::PetPSClient>( | |
| 301 | ✗ | endpoint.host, endpoint.port, endpoint.shard, logical_client_id)); | |
| 302 | } | ||
| 303 | } | ||
| 304 | |||
| 305 | ✗ | executor_ = std::make_unique<RdmaDistributedExecutor>( | |
| 306 | ✗ | [&]() { | |
| 307 | ✗ | std::vector<petps::PetPSClient*> clients; | |
| 308 | ✗ | clients.reserve(shard_clients_.size()); | |
| 309 | ✗ | for (auto& client : shard_clients_) | |
| 310 | ✗ | clients.push_back(client.get()); | |
| 311 | ✗ | return clients; | |
| 312 | ✗ | }(), | |
| 313 | FLAGS_value_size, | ||
| 314 | ✗ | -1, | |
| 315 | ✗ | -1); | |
| 316 | ✗ | initialized_ = true; | |
| 317 | ✗ | } | |
| 318 | |||
| 319 | ✗ | void RDMAPSClientAdapter::EnsureThreadInitialized() { | |
| 320 | ✗ | EnsureClientInitialized(); | |
| 321 | ✗ | const std::thread::id tid = std::this_thread::get_id(); | |
| 322 | ✗ | std::lock_guard<std::mutex> guard(thread_init_mu_); | |
| 323 | ✗ | if (initialized_threads_.find(tid) != initialized_threads_.end()) { | |
| 324 | ✗ | return; | |
| 325 | } | ||
| 326 | |||
| 327 | ✗ | if (deployment_.num_shards <= 1) { | |
| 328 | ✗ | if (client_ != nullptr) { | |
| 329 | ✗ | client_->InitThread(); | |
| 330 | } | ||
| 331 | } else { | ||
| 332 | ✗ | for (auto& shard_client : shard_clients_) { | |
| 333 | ✗ | shard_client->InitThread(); | |
| 334 | } | ||
| 335 | } | ||
| 336 | |||
| 337 | ✗ | initialized_threads_.insert(tid); | |
| 338 | ✗ | } | |
| 339 | |||
| 340 | ✗ | void RDMAPSClientAdapter::EnsureTableReady(const std::string& table_name, | |
| 341 | int64_t embedding_dim) { | ||
| 342 | ✗ | std::lock_guard<std::mutex> guard(state_mu_); | |
| 343 | ✗ | const auto it = tables_.find(table_name); | |
| 344 | ✗ | if (it == tables_.end()) { | |
| 345 | ✗ | throw std::runtime_error("RDMA table is not initialized: " + table_name); | |
| 346 | } | ||
| 347 | ✗ | if (static_cast<int64_t>(it->second.config.embedding_dim) != embedding_dim) { | |
| 348 | ✗ | throw std::runtime_error( | |
| 349 | ✗ | "RDMA embedding dimension mismatch for table " + table_name); | |
| 350 | } | ||
| 351 | ✗ | } | |
| 352 | |||
| 353 | ✗ | int64_t RDMAPSClientAdapter::DefaultEmbeddingDimOrThrow() const { | |
| 354 | ✗ | std::lock_guard<std::mutex> guard(state_mu_); | |
| 355 | ✗ | if (tables_.empty()) { | |
| 356 | ✗ | throw std::runtime_error( | |
| 357 | ✗ | "RDMA table metadata is empty; call InitEmbeddingTable first"); | |
| 358 | } | ||
| 359 | ✗ | return static_cast<int64_t>(tables_.begin()->second.config.embedding_dim); | |
| 360 | ✗ | } | |
| 361 | |||
| 362 | ✗ | int64_t RDMAPSClientAdapter::EmbeddingDimForKeys( | |
| 363 | base::ConstArray<uint64_t> keys) const { | ||
| 364 | ✗ | if (keys.Size() == 0) { | |
| 365 | ✗ | return DefaultEmbeddingDimOrThrow(); | |
| 366 | } | ||
| 367 | ✗ | const int tag = static_cast<int>(ExtractKeyTag(keys[0])); | |
| 368 | ✗ | std::lock_guard<std::mutex> guard(state_mu_); | |
| 369 | ✗ | const auto it = tag_to_dim_.find(tag); | |
| 370 | ✗ | if (it == tag_to_dim_.end()) { | |
| 371 | ✗ | throw std::runtime_error("unknown RDMA key tag: " + std::to_string(tag)); | |
| 372 | } | ||
| 373 | ✗ | return it->second; | |
| 374 | ✗ | } | |
| 375 | |||
| 376 | ✗ | std::size_t RDMAPSClientAdapter::MaxGetKeysPerRpc(int64_t embedding_dim) const { | |
| 377 | ✗ | const std::size_t value_size = | |
| 378 | ✗ | static_cast<std::size_t>(embedding_dim) * sizeof(float); | |
| 379 | ✗ | const std::size_t response_limited = petps::GetKeysPerRpcByResponseBudget( | |
| 380 | value_size, | ||
| 381 | static_cast<std::size_t>(FLAGS_rdma_rc_mtu_bytes), | ||
| 382 | ✗ | static_cast<std::size_t>(FLAGS_rdma_rc_target_response_mtu)); | |
| 383 | const std::size_t request_limited = | ||
| 384 | ✗ | petps::PutPayloadBudget( | |
| 385 | ✗ | static_cast<std::size_t>(FLAGS_rdma_rc_request_slot_bytes)) / | |
| 386 | ✗ | sizeof(std::uint64_t); | |
| 387 | ✗ | std::size_t limit = static_cast<std::size_t>(FLAGS_max_kv_num_per_request); | |
| 388 | ✗ | if (response_limited > 0) { | |
| 389 | ✗ | limit = std::min(limit, response_limited); | |
| 390 | } | ||
| 391 | ✗ | if (request_limited > 0) { | |
| 392 | ✗ | limit = std::min(limit, request_limited); | |
| 393 | } | ||
| 394 | ✗ | return std::max<std::size_t>(limit, 1); | |
| 395 | } | ||
| 396 | |||
| 397 | ✗ | std::size_t RDMAPSClientAdapter::MaxPutKeysPerRpc(int64_t embedding_dim) const { | |
| 398 | ✗ | const std::size_t payload_budget = petps::PutPayloadBudget( | |
| 399 | static_cast<std::size_t>(FLAGS_rdma_rc_request_slot_bytes)); | ||
| 400 | ✗ | const std::size_t bytes_per_row = | |
| 401 | ✗ | sizeof(ParameterCompressItem) + | |
| 402 | ✗ | static_cast<std::size_t>(embedding_dim) * sizeof(float) + sizeof(int); | |
| 403 | ✗ | std::size_t limit = static_cast<std::size_t>(FLAGS_max_kv_num_per_request); | |
| 404 | ✗ | if (payload_budget > sizeof(int) && bytes_per_row > 0) { | |
| 405 | ✗ | const std::size_t request_limited = | |
| 406 | ✗ | (payload_budget - sizeof(int)) / bytes_per_row; | |
| 407 | ✗ | if (request_limited > 0) { | |
| 408 | ✗ | limit = std::min(limit, request_limited); | |
| 409 | } | ||
| 410 | } | ||
| 411 | ✗ | return std::max<std::size_t>(limit, 1); | |
| 412 | } | ||
| 413 | |||
| 414 | ✗ | std::size_t RDMAPSClientAdapter::MaxInFlightGetRpcs() const { | |
| 415 | const std::size_t qps = static_cast<std::size_t>( | ||
| 416 | ✗ | std::max(FLAGS_rdma_rc_qps_per_client_per_shard, 1)); | |
| 417 | const std::size_t slots = | ||
| 418 | ✗ | static_cast<std::size_t>(std::max(FLAGS_rdma_rc_slots_per_qp, 1)); | |
| 419 | ✗ | return std::max<std::size_t>(qps * slots, 1); | |
| 420 | } | ||
| 421 | |||
| 422 | ✗ | bool RDMAPSClientAdapter::QueryRPCFinished(int rpc_id) { | |
| 423 | ✗ | return executor_->QueryBatchFinished(rpc_id); | |
| 424 | } | ||
| 425 | |||
| 426 | ✗ | void RDMAPSClientAdapter::WaitRPCFinish(int rpc_id) { | |
| 427 | ✗ | executor_->WaitBatch(rpc_id); | |
| 428 | ✗ | } | |
| 429 | |||
| 430 | ✗ | void RDMAPSClientAdapter::RevokeRPCResource(int rpc_id) { | |
| 431 | ✗ | executor_->ReleaseBatch(rpc_id); | |
| 432 | ✗ | } | |
| 433 | |||
| 434 | ✗ | int RDMAPSClientAdapter::SubmitGetParameter( | |
| 435 | base::ConstArray<uint64_t> keys, | ||
| 436 | float* values, | ||
| 437 | bool isAsync, | ||
| 438 | int async_req_id, | ||
| 439 | int64_t embedding_dim, | ||
| 440 | bool write_status_word) { | ||
| 441 | ✗ | EnsureThreadInitialized(); | |
| 442 | ✗ | if (keys.Size() == 0) { | |
| 443 | ✗ | if (write_status_word) { | |
| 444 | ✗ | auto* status = reinterpret_cast<std::int32_t*>(values); | |
| 445 | ✗ | *status = static_cast<std::int32_t>(petps::RpcStatus::kOk); | |
| 446 | } | ||
| 447 | ✗ | return 0; | |
| 448 | } | ||
| 449 | ✗ | const int value_size = | |
| 450 | static_cast<int>(static_cast<std::size_t>(embedding_dim) * sizeof(float)); | ||
| 451 | ✗ | return executor_->SubmitGetParameter( | |
| 452 | ✗ | BuildChunks(keys, MaxGetKeysPerRpc(embedding_dim)), | |
| 453 | ✗ | keys.Size(), | |
| 454 | value_size, | ||
| 455 | values, | ||
| 456 | isAsync, | ||
| 457 | async_req_id, | ||
| 458 | MaxInFlightGetRpcs(), | ||
| 459 | ✗ | write_status_word); | |
| 460 | } | ||
| 461 | |||
| 462 | ✗ | int RDMAPSClientAdapter::GetParameter(const base::ConstArray<uint64_t>& keys, | |
| 463 | base::RecTensor& values) { | ||
| 464 | ✗ | EnsureThreadInitialized(); | |
| 465 | ✗ | std::lock_guard<std::mutex> io_guard(io_mu_); | |
| 466 | ✗ | executor_->DrainBatchSlots(); | |
| 467 | ✗ | if (!IsFloatEmbeddingValues(values, static_cast<int64_t>(keys.Size()))) { | |
| 468 | ✗ | return -1; | |
| 469 | } | ||
| 470 | ✗ | if (keys.Size() == 0) { | |
| 471 | ✗ | return 0; | |
| 472 | } | ||
| 473 | ✗ | const int64_t embedding_dim = values.shape(1); | |
| 474 | ✗ | const int value_size = | |
| 475 | static_cast<int>(static_cast<std::size_t>(embedding_dim) * sizeof(float)); | ||
| 476 | |||
| 477 | // The shard merge writes every requested row, so submit the caller's tensor | ||
| 478 | // as the batch destination instead of round-tripping through a scratch | ||
| 479 | // buffer. `values` has no room for the trailing status word, so read the | ||
| 480 | // batch verdict from WaitBatch. | ||
| 481 | ✗ | const int rpc_id = SubmitGetParameter( | |
| 482 | keys, | ||
| 483 | values.data_as<float>(), | ||
| 484 | /*isAsync=*/false, | ||
| 485 | /*async_req_id=*/0, | ||
| 486 | embedding_dim, | ||
| 487 | /*write_status_word=*/false); | ||
| 488 | ✗ | const bool ok = executor_->WaitBatch(rpc_id); | |
| 489 | ✗ | RevokeRPCResource(rpc_id); | |
| 490 | ✗ | return ok ? 0 : -1; | |
| 491 | ✗ | } | |
| 492 | |||
| 493 | 2 | int RDMAPSClientAdapter::PutParameter(const base::ConstArray<uint64_t>& keys, | |
| 494 | const base::RecTensor& values) { | ||
| 495 | 2 | if (!IsFloatEmbeddingValues(values, static_cast<int64_t>(keys.Size()))) { | |
| 496 | ✗ | return -1; | |
| 497 | } | ||
| 498 | 2 | if (keys.Size() == 0) { | |
| 499 | 2 | return 0; | |
| 500 | } | ||
| 501 | ✗ | EnsureThreadInitialized(); | |
| 502 | ✗ | std::lock_guard<std::mutex> io_guard(io_mu_); | |
| 503 | ✗ | executor_->DrainBatchSlots(); | |
| 504 | ✗ | return executor_->PutParameter( | |
| 505 | ✗ | BuildChunks(keys, MaxPutKeysPerRpc(values.shape(1))), values); | |
| 506 | ✗ | } | |
| 507 | |||
| 508 | ✗ | int64_t RDMAPSClientAdapter::SparseGradientCommit( | |
| 509 | const std::string& table, | ||
| 510 | SparseGradientCommitHeader h, | ||
| 511 | const base::ConstArray<uint64_t>& keys, | ||
| 512 | const base::RecTensor& grads) { | ||
| 513 | ✗ | EnsureClientInitialized(); | |
| 514 | ✗ | EnsureThreadInitialized(); | |
| 515 | ✗ | EnsureTableReady(table, h.dim); | |
| 516 | ✗ | std::lock_guard<std::mutex> io_guard(io_mu_); | |
| 517 | ✗ | executor_->DrainBatchSlots(); | |
| 518 | ✗ | if (h.kind == 0 && (grads.dim() != 2 || grads.shape(1) != h.dim || | |
| 519 | ✗ | !IsFloatEmbeddingValues(grads, keys.Size()))) { | |
| 520 | ✗ | throw std::invalid_argument("Invalid RDMA sparse gradients"); | |
| 521 | } | ||
| 522 | const size_t limit = | ||
| 523 | ✗ | shard_clients_.front()->SparseGradientCommitPayloadLimit(); | |
| 524 | ✗ | const size_t row_bytes = 8 + size_t(h.dim) * 4; | |
| 525 | ✗ | if (limit < sizeof(h) + row_bytes) | |
| 526 | ✗ | throw std::runtime_error("Gradient row exceeds RDMA slot"); | |
| 527 | ✗ | std::vector<std::vector<ShardChunk>> shards(shard_clients_.size()); | |
| 528 | ✗ | if (h.kind == 0) { | |
| 529 | ✗ | for (auto& chunk : BuildChunks(keys, (limit - sizeof(h)) / row_bytes)) { | |
| 530 | ✗ | shards[chunk.client_index].push_back(std::move(chunk)); | |
| 531 | ✗ | } | |
| 532 | } | ||
| 533 | ✗ | size_t rounds = 1; | |
| 534 | ✗ | for (auto& shard : shards) { | |
| 535 | ✗ | if (shard.empty()) | |
| 536 | ✗ | shard.emplace_back(); // Empty-rank marker for every PS. | |
| 537 | ✗ | rounds = std::max(rounds, shard.size()); | |
| 538 | } | ||
| 539 | ✗ | std::vector<int64_t> committed(shards.size(), -1); | |
| 540 | ✗ | std::deque<std::pair<size_t, int>> pending; | |
| 541 | ✗ | auto drain_one = [&] { | |
| 542 | ✗ | const auto [client, rpc] = pending.front(); | |
| 543 | ✗ | pending.pop_front(); | |
| 544 | ✗ | committed[client] = | |
| 545 | ✗ | std::max(committed[client], | |
| 546 | ✗ | shard_clients_[client]->WaitSparseGradientCommit(rpc)); | |
| 547 | ✗ | }; | |
| 548 | try { | ||
| 549 | // Round-robin chunks across shards, with a bounded window of outstanding | ||
| 550 | // RC writes. Gather directly into each packet, without a full shard copy. | ||
| 551 | ✗ | for (size_t round = 0; round < rounds; ++round) { | |
| 552 | ✗ | for (size_t client = 0; client < shards.size(); ++client) { | |
| 553 | ✗ | if (round >= shards[client].size()) | |
| 554 | ✗ | continue; | |
| 555 | ✗ | const auto& chunk = shards[client][round]; | |
| 556 | ✗ | h.chunk = round; | |
| 557 | ✗ | h.chunks = shards[client].size(); | |
| 558 | ✗ | h.rows = chunk.keys.size(); | |
| 559 | ✗ | std::string packet(sizeof(h) + h.rows * row_bytes, '\0'); | |
| 560 | ✗ | std::memcpy(packet.data(), &h, sizeof(h)); | |
| 561 | ✗ | if (h.rows) { | |
| 562 | ✗ | std::memcpy(packet.data() + sizeof(h), chunk.keys.data(), h.rows * 8); | |
| 563 | ✗ | char* dst = packet.data() + sizeof(h) + h.rows * 8; | |
| 564 | ✗ | for (size_t row = 0; row < h.rows; ++row) { | |
| 565 | ✗ | std::memcpy(dst + row * h.dim * 4, | |
| 566 | ✗ | grads.data_as<float>() + chunk.positions[row] * h.dim, | |
| 567 | ✗ | h.dim * 4); | |
| 568 | } | ||
| 569 | } | ||
| 570 | ✗ | pending.emplace_back( | |
| 571 | client, | ||
| 572 | ✗ | shard_clients_[client]->SubmitSparseGradientCommit(table, packet)); | |
| 573 | ✗ | if (pending.size() >= MaxInFlightGetRpcs()) | |
| 574 | ✗ | drain_one(); | |
| 575 | ✗ | } | |
| 576 | } | ||
| 577 | ✗ | while (!pending.empty()) | |
| 578 | ✗ | drain_one(); | |
| 579 | ✗ | } catch (...) { | |
| 580 | ✗ | while (!pending.empty()) { | |
| 581 | try { | ||
| 582 | ✗ | drain_one(); | |
| 583 | ✗ | } catch (...) { | |
| 584 | ✗ | } | |
| 585 | } | ||
| 586 | ✗ | throw; | |
| 587 | ✗ | } | |
| 588 | ✗ | return *std::min_element(committed.begin(), committed.end()); | |
| 589 | ✗ | } | |
| 590 | |||
| 591 | 2 | int RDMAPSClientAdapter::UpdateParameter(const std::string& table_name, | |
| 592 | const base::ConstArray<uint64_t>& keys, | ||
| 593 | const base::RecTensor& grads) { | ||
| 594 | 2 | if (keys.Size() == 0) { | |
| 595 | 2 | return 0; | |
| 596 | } | ||
| 597 | ✗ | return WaitUpdateParameter( | |
| 598 | ✗ | SubmitUpdateParameterAsync(table_name, keys, grads)); | |
| 599 | } | ||
| 600 | |||
| 601 | 2 | uint64_t RDMAPSClientAdapter::SubmitUpdateParameterAsync( | |
| 602 | const std::string& table_name, | ||
| 603 | const base::ConstArray<uint64_t>& keys, | ||
| 604 | const base::RecTensor& grads) { | ||
| 605 | 2 | if (keys.Size() == 0) { | |
| 606 | 2 | throw std::invalid_argument("RDMA update requires at least one key"); | |
| 607 | } | ||
| 608 | ✗ | if (!IsFloatEmbeddingValues(grads, static_cast<int64_t>(keys.Size()))) { | |
| 609 | ✗ | throw std::invalid_argument("RDMA update has invalid rows or gradients"); | |
| 610 | } | ||
| 611 | ✗ | const int64_t embedding_dim = grads.shape(1); | |
| 612 | ✗ | EnsureTableReady(table_name, embedding_dim); | |
| 613 | ✗ | EnsureThreadInitialized(); | |
| 614 | ✗ | std::lock_guard<std::mutex> io_guard(io_mu_); | |
| 615 | ✗ | executor_->DrainBatchSlots(); | |
| 616 | ✗ | return executor_->SubmitUpdateParameterFlat( | |
| 617 | table_name, | ||
| 618 | keys, | ||
| 619 | ✗ | grads.data_as<float>(), | |
| 620 | static_cast<std::size_t>(embedding_dim), | ||
| 621 | ✗ | BuildChunks(keys, MaxPutKeysPerRpc(embedding_dim)), | |
| 622 | ✗ | MaxInFlightGetRpcs()); | |
| 623 | ✗ | } | |
| 624 | |||
| 625 | ✗ | int RDMAPSClientAdapter::WaitUpdateParameter(uint64_t update_id) { | |
| 626 | ✗ | return executor_->WaitUpdateParameterFlat(update_id); | |
| 627 | } | ||
| 628 | |||
| 629 | ✗ | int RDMAPSClientAdapter::InitEmbeddingTable( | |
| 630 | const std::string& table_name, const EmbeddingTableConfig& config) { | ||
| 631 | ✗ | EnsureThreadInitialized(); | |
| 632 | ✗ | std::lock_guard<std::mutex> io_guard(io_mu_); | |
| 633 | ✗ | executor_->DrainBatchSlots(); | |
| 634 | ✗ | int tag = -1; | |
| 635 | ✗ | for (auto& shard_client : shard_clients_) { | |
| 636 | ✗ | const int rc = shard_client->InitEmbeddingTable( | |
| 637 | table_name, | ||
| 638 | ✗ | config.num_embeddings, | |
| 639 | ✗ | config.embedding_dim, | |
| 640 | ✗ | config.table_id); | |
| 641 | ✗ | if (rc < 0) { | |
| 642 | ✗ | return rc; | |
| 643 | } | ||
| 644 | ✗ | if (tag < 0) { | |
| 645 | ✗ | tag = rc; | |
| 646 | ✗ | } else if (tag != rc) { | |
| 647 | ✗ | return -1; | |
| 648 | } | ||
| 649 | } | ||
| 650 | ✗ | std::lock_guard<std::mutex> guard(state_mu_); | |
| 651 | ✗ | const auto [it, inserted] = | |
| 652 | ✗ | tables_.emplace(table_name, TableState{config, tag}); | |
| 653 | ✗ | if (!inserted && (it->second.config.embedding_dim != config.embedding_dim || | |
| 654 | ✗ | it->second.config.num_embeddings != config.num_embeddings || | |
| 655 | ✗ | it->second.tag != tag)) { | |
| 656 | ✗ | return -1; | |
| 657 | } | ||
| 658 | ✗ | tag_to_dim_[tag] = static_cast<int64_t>(config.embedding_dim); | |
| 659 | ✗ | return tag; | |
| 660 | ✗ | } | |
| 661 | |||
| 662 | ✗ | void RDMAPSClientAdapter::Command(PSCommand) { | |
| 663 | ✗ | EnsureThreadInitialized(); | |
| 664 | ✗ | if (deployment_.num_shards <= 1) { | |
| 665 | ✗ | if (client_ == nullptr) { | |
| 666 | ✗ | throw std::runtime_error("RDMA adapter has no initialized client"); | |
| 667 | } | ||
| 668 | ✗ | client_->Barrier("rdma_command", 0); | |
| 669 | ✗ | return; | |
| 670 | } | ||
| 671 | ✗ | if (shard_clients_.empty()) { | |
| 672 | ✗ | throw std::runtime_error("RDMA adapter has no initialized clients"); | |
| 673 | } | ||
| 674 | ✗ | shard_clients_.front()->Barrier("rdma_command", 0); | |
| 675 | } | ||
| 676 | |||
| 677 | uint64_t | ||
| 678 | ✗ | RDMAPSClientAdapter::PrefetchParameter(const base::ConstArray<uint64_t>& keys) { | |
| 679 | ✗ | EnsureThreadInitialized(); | |
| 680 | ✗ | std::lock_guard<std::mutex> io_guard(io_mu_); | |
| 681 | ✗ | if (keys.Size() == 0) { | |
| 682 | ✗ | throw std::invalid_argument("RDMA prefetch requires at least one key"); | |
| 683 | } | ||
| 684 | ✗ | const int64_t embedding_dim = EmbeddingDimForKeys(keys); | |
| 685 | ✗ | const int value_size = | |
| 686 | static_cast<int>(static_cast<std::size_t>(embedding_dim) * sizeof(float)); | ||
| 687 | ✗ | const auto chunks = BuildChunks(keys, MaxGetKeysPerRpc(embedding_dim)); | |
| 688 | ✗ | const std::size_t max_in_flight = MaxInFlightGetRpcs(); | |
| 689 | ✗ | if (!executor_->HasGetSlotCapacity(chunks, max_in_flight)) { | |
| 690 | ✗ | executor_->DrainBatchSlots(); | |
| 691 | } | ||
| 692 | ✗ | return executor_->SubmitPrefetch( | |
| 693 | ✗ | chunks, keys.Size(), embedding_dim, value_size, max_in_flight); | |
| 694 | ✗ | } | |
| 695 | |||
| 696 | ✗ | bool RDMAPSClientAdapter::IsPrefetchDone(uint64_t prefetch_id) { | |
| 697 | ✗ | EnsureThreadInitialized(); | |
| 698 | ✗ | return executor_->IsPrefetchDone(prefetch_id); | |
| 699 | } | ||
| 700 | |||
| 701 | ✗ | void RDMAPSClientAdapter::WaitForPrefetch(uint64_t prefetch_id) { | |
| 702 | ✗ | EnsureThreadInitialized(); | |
| 703 | ✗ | executor_->WaitForPrefetch(prefetch_id); | |
| 704 | ✗ | } | |
| 705 | |||
| 706 | ✗ | bool RDMAPSClientAdapter::GetPrefetchResult(uint64_t prefetch_id, | |
| 707 | base::RecTensor& values) { | ||
| 708 | // Keep borrowed MR payloads live until the consumer has copied them. | ||
| 709 | ✗ | std::lock_guard<std::mutex> io_guard(io_mu_); | |
| 710 | ✗ | const auto result = executor_->ReadPrefetch(prefetch_id); | |
| 711 | ✗ | const bool discard = values.data() == nullptr && values.dim() == 0; | |
| 712 | ✗ | if (!discard && !EnsureEmbeddingOutput(values, result.key_count)) { | |
| 713 | ✗ | executor_->ReleasePrefetch(prefetch_id); | |
| 714 | ✗ | return false; | |
| 715 | } | ||
| 716 | ✗ | if (!discard && values.shape(1) != result.embedding_dim) { | |
| 717 | ✗ | executor_->ReleasePrefetch(prefetch_id); | |
| 718 | ✗ | return false; | |
| 719 | } | ||
| 720 | ✗ | const bool ok = | |
| 721 | ✗ | result.status_code == static_cast<std::int32_t>(petps::RpcStatus::kOk); | |
| 722 | ✗ | if (ok && !discard && !FLAGS_rdma_adapter_skip_prefetch_result_copy && | |
| 723 | ✗ | result.payload != nullptr && result.response_bytes > 0) { | |
| 724 | ✗ | const std::size_t expected_bytes = | |
| 725 | ✗ | static_cast<std::size_t>(result.key_count) * | |
| 726 | ✗ | static_cast<std::size_t>(result.embedding_dim) * sizeof(float); | |
| 727 | ✗ | if (result.response_bytes < expected_bytes) { | |
| 728 | ✗ | executor_->ReleasePrefetch(prefetch_id); | |
| 729 | ✗ | return false; | |
| 730 | } | ||
| 731 | ✗ | std::memcpy(values.data_as<float>(), result.payload, expected_bytes); | |
| 732 | } | ||
| 733 | ✗ | executor_->ReleasePrefetch(prefetch_id); | |
| 734 | ✗ | return ok; | |
| 735 | ✗ | } | |
| 736 | |||
| 737 | } // namespace recstore | ||
| 738 |