Fix DramKV race: hold rlock during inplace update writes (#5807)

Summary:
Pull Request resolved: #5807

X-link: https://github.com/facebookresearch/FBGEMM/pull/2733

Previously the rlock was released after collecting hit/miss info, then block pointers were written to without the lock held. This allowed eviction to invalidate block pointers during writes. Fix by keeping the rlock held through the entire read-and-write phase using scoped lifetime.

Pull Request resolved: #5536

Test Plan:
Imported from GitHub, without a `Test Plan:` line.

Split from D95983752

Reviewed By: r-barnes

Differential Revision: D98427441

Pulled By: q10

fbshipit-source-id: f30a460b6b71c4963577135e742ef8c92c077105

Author: cyyever

fbgemm_gpu/src/dram_kv_embedding_cache/dram_kv_inference_embedding.h

@@ -255,47 +255,47 @@ class DramKVInferenceEmbedding
                       auto indices_data_ptr = indices.data_ptr<index_t>();
                       auto weights_data_ptr = weights.data_ptr<weight_type>();
                       {
-                        // 1st step, collect hit/miss per inplace update chunk
                         // [tensor_offset, weight_addr]
                         std::vector<std::tuple<int64_t, weight_type*>> hit_info;
                         // [id, tensor_offset]
                         std::vector<std::tuple<int64_t, int64_t>> miss_info;
                         hit_info.reserve(indexes.size() / 2);
                         miss_info.reserve(indexes.size() / 10);
-                        auto rlmap = kv_store_.by(shard_id).rlock();
-                        for (const auto& idx : indexes) {
-                          auto id = int64_t(indices_data_ptr[idx]);
-                          auto it = rlmap->find(id);
-                          if (it != rlmap->end()) {
-                            hit_info.emplace_back(idx, it->second);
-                          } else {
-                            miss_info.emplace_back(id, idx);
+                        {
+                          // 1st step, collect hit/miss per inplace update chunk
+                          auto rlmap = kv_store_.by(shard_id).rlock();
+                          for (const auto& idx : indexes) {
+                            auto id = int64_t(indices_data_ptr[idx]);
+                            auto it = rlmap->find(id);
+                            if (it != rlmap->end()) {
+                              hit_info.emplace_back(idx, it->second);
+                            } else {
+                              miss_info.emplace_back(id, idx);
+                            }
                           }
-                        }
-                        rlmap.unlock();
-                        hit_cnt = hit_info.size();
-                        miss_cnt = miss_info.size();
-                        // 2nd step, no lock on update hits, it is possible that
-                        // inference read is accessing a weight being updated,
-                        // we assume it is fine for now, will iterate on it if
-                        // we find QE regress during inplace update
-                        for (auto& [tensor_offset, block] : hit_info) {
-                          auto* data_ptr =
-                              InferenceFixedBlockPool::data_ptr<weight_type>(
-                                  block);
-                          std::copy(
-                              weights_data_ptr + tensor_offset * stride,
-                              weights_data_ptr + (tensor_offset + 1) * stride,
-                              data_ptr);
-                          // update provided ts for existing blocks
-                          if (feature_evict_config_.has_value() &&
-                              feature_evict_config_.value()->trigger_mode_ !=
-                                  EvictTriggerMode::DISABLED &&
-                              feature_evict_ && inplace_update_ts.has_value()) {
-                            InferenceFixedBlockPool::set_timestamp(
-                                block, inplace_update_ts.value());
+                          hit_cnt = hit_info.size();
+                          miss_cnt = miss_info.size();
+                          // 2nd step, update hits while holding rlock to
+                          // prevent eviction from invalidating block pointers
+                          for (auto& [tensor_offset, block] : hit_info) {
+                            auto* data_ptr =
+                                InferenceFixedBlockPool::data_ptr<weight_type>(
+                                    block);
+                            std::copy(
+                                weights_data_ptr + tensor_offset * stride,
+                                weights_data_ptr + (tensor_offset + 1) * stride,
+                                data_ptr);
+                            // update provided ts for existing blocks
+                            if (feature_evict_config_.has_value() &&
+                                feature_evict_config_.value()->trigger_mode_ !=
+                                    EvictTriggerMode::DISABLED &&
+                                feature_evict_ &&
+                                inplace_update_ts.has_value()) {
+                              InferenceFixedBlockPool::set_timestamp(
+                                  block, inplace_update_ts.value());
+                            }
                           }
-                        }
+                        } // rlmap released here
 
                         // 3rd step, update misses in fixed block pool, we only
                         // need mempool lock at this stage to avoid race