update test

Author: ArronHZG

fbgemm_gpu/src/dram_kv_embedding_cache/dram_kv_embedding_cache.h

@@ -85,7 +85,8 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
         num_shards_(num_shards),
         weight_ttl_in_hours_(weight_ttl_in_hours),
         block_size_(FixedBlockPool::calculate_block_size<weight_type>(max_D)),
-        block_alignment_(FixedBlockPool::calculate_block_alignment<weight_type>()),
+        block_alignment_(
+            FixedBlockPool::calculate_block_alignment<weight_type>()),
         kv_store_(SynchronizedShardedMap<int64_t, weight_type*>(
             num_shards_,
             block_size_,
@@ -262,8 +263,11 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
                             feature_evict_->update_feature_statistics(block);
                           }
                           auto* data_ptr = FixedBlockPool::data_ptr<weight_type>(block);
-                          std::copy(weights[id_index].template data_ptr<weight_type>(),
-                                    weights[id_index].template data_ptr<weight_type>() + weights[id_index].numel(),
+                          std::copy(weights[id_index]
+                                        .template data_ptr<weight_type>(),
+                                    weights[id_index]
+                                            .template data_ptr<weight_type>() +
+                                        weights[id_index].numel(),
                                     data_ptr);
                         }
                       }
@@ -366,7 +370,9 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
                             continue;
                           }
                           // use mempool
-                          const auto* data_ptr = FixedBlockPool::data_ptr<weight_type>(cached_iter->second);
+                          const auto* data_ptr =
+                              FixedBlockPool::data_ptr<weight_type>(
+                                  cached_iter->second);
                           std::copy(
                               data_ptr + width_offset,
                               data_ptr + width_offset + row_width,

fbgemm_gpu/src/dram_kv_embedding_cache/fixed_block_pool.h

@@ -236,7 +236,7 @@ class FixedBlockPool : public std::pmr::memory_resource {
   const std::size_t block_alignment_; // Block alignment requirement
   const std::size_t blocks_per_chunk_; // Number of blocks per chunk
   std::pmr::memory_resource* upstream_; // Upstream memory resource
-  std::pmr::vector<chunk_info> chunks_{1024}; // Records of all allocated chunks
+  std::pmr::vector<ChunkInfo> chunks_{1024}; // Records of all allocated chunks
   void* free_list_ = nullptr; // Free block list head pointer
 };
 } // namespace kv_mem

fbgemm_gpu/test/dram_kv_embedding_cache/feature_evict_test.cpp

@@ -18,45 +18,69 @@ size_t BLOCK_ALIGNMENT = FixedBlockPool::calculate_block_alignment<float>();
 TEST(FeatureEvictTest, CounterBasedEviction) {
   static constexpr int NUM_SHARDS = 8;
   auto executor_ = std::make_unique<folly::CPUThreadPoolExecutor>(4);
-  auto kv_store_ = std::make_unique<SynchronizedShardedMap<int64_t, float*>>(NUM_SHARDS, BLOCK_SIZE, BLOCK_ALIGNMENT);
+  auto kv_store_ = std::make_unique<SynchronizedShardedMap<int64_t, float*>>(
+      NUM_SHARDS, BLOCK_SIZE, BLOCK_ALIGNMENT);
 
   // Insert test data
   for (int i = 0; i < 1000; ++i) {
     int shard_id = i % NUM_SHARDS;
     auto wlock = kv_store_->by(shard_id).wlock();
     auto* pool = kv_store_->pool_by(shard_id);
     auto* block = pool->allocate_t<float>();
+    ASSERT_NE(block, nullptr);
     FixedBlockPool::set_key(block, i);
-    FixedBlockPool::set_count(block, 1);  // Initial score
+    FixedBlockPool::set_count(block, 1);
     FixedBlockPool::set_used(block, true);
-    wlock->insert({i, block});
+    auto result = wlock->insert({i, block});
+    ASSERT_TRUE(result.second);
   }
 
   for (int i = 1000; i < 2000; ++i) {
     int shard_id = i % NUM_SHARDS;
     auto wlock = kv_store_->by(shard_id).wlock();
     auto* pool = kv_store_->pool_by(shard_id);
     auto* block = pool->allocate_t<float>();
+    ASSERT_NE(block, nullptr);
     FixedBlockPool::set_key(block, i);
-    FixedBlockPool::set_count(block, 2);  // Initial score
+    FixedBlockPool::set_count(block, 2);
     FixedBlockPool::set_used(block, true);
-    wlock->insert({i, block});
+    auto result = wlock->insert({i, block});
+    ASSERT_TRUE(result.second);
+  }
+
+  size_t count_1_blocks = 0, count_2_blocks = 0;
+  for (int shard_id = 0; shard_id < NUM_SHARDS; ++shard_id) {
+    auto rlock = kv_store_->by(shard_id).rlock();
+    for (const auto& [key, block] : *rlock) {
+      uint32_t count = FixedBlockPool::get_count(block);
+      if (count == 1)
+        count_1_blocks++;
+      else if (count == 2)
+        count_2_blocks++;
+      ASSERT_TRUE(FixedBlockPool::get_used(block));
+    }
   }
+  ASSERT_EQ(count_1_blocks, 1000);
+  ASSERT_EQ(count_2_blocks, 1000);
 
   std::unique_ptr<FeatureEvict<float>> feature_evict;
   int evict_trigger_mode = 2;
   int evict_trigger_strategy = 1;
   uint32_t count_threshold = 1;
   float count_decay_rate = 0.5;
-  // feature evict config
-  FeatureEvictConfig feature_evict_config;
-  feature_evict_config.trigger_mode = static_cast<EvictTriggerMode>(evict_trigger_mode);
-  feature_evict_config.trigger_strategy = static_cast<EvictTriggerStrategy>(evict_trigger_strategy);
+
+  FeatureEvictConfig feature_evict_config{};
+  feature_evict_config.trigger_mode =
+      static_cast<EvictTriggerMode>(evict_trigger_mode);
+  feature_evict_config.trigger_strategy =
+      static_cast<EvictTriggerStrategy>(evict_trigger_strategy);
   feature_evict_config.count_threshold = count_threshold;
   feature_evict_config.count_decay_rate = count_decay_rate;
 
   if (feature_evict_config.trigger_mode != EvictTriggerMode::DISABLED) {
-    feature_evict = create_feature_evict(feature_evict_config, executor_.get(),*kv_store_.get(), 4);
+    feature_evict = create_feature_evict(
+        feature_evict_config, executor_.get(), *kv_store_.get(), 4);
+    ASSERT_NE(feature_evict, nullptr);
   }
 
   // Initial validation
@@ -69,69 +93,101 @@ TEST(FeatureEvictTest, CounterBasedEviction) {
 
   // Perform eviction
   feature_evict->trigger_evict();
+  ASSERT_TRUE(feature_evict->is_evicting());
 
   // Validate eviction process
-  while (feature_evict->is_evicting()) {
+  int max_iterations = 1000;
+  int iterations = 0;
+  while (feature_evict->is_evicting() && iterations < max_iterations) {
     feature_evict->resume();
     std::this_thread::sleep_for(std::chrono::microseconds(5));
     feature_evict->pause();
+    iterations++;
   }
+  ASSERT_LT(iterations, max_iterations);
+  ASSERT_FALSE(feature_evict->is_evicting());
 
   // Validate results
   size_t remaining = 0;
+  size_t remaining_count_1 = 0;
   for (int shard_id = 0; shard_id < NUM_SHARDS; ++shard_id) {
     auto rlock = kv_store_->by(shard_id).rlock();
     remaining += rlock->size();
     // Validate score decay
     for (const auto& [key, block] : *rlock) {
-      ASSERT_EQ(FixedBlockPool::get_count(block), 1);
+      uint32_t count = FixedBlockPool::get_count(block);
+      ASSERT_EQ(count, 1);
+      if (count == 1) remaining_count_1++;
     }
   }
   std::cout << "remaining: " << remaining << std::endl;
   ASSERT_EQ(remaining, 1000);
+  ASSERT_EQ(remaining_count_1, 1000);
 }
 
 TEST(FeatureEvictTest, TimeBasedEviction) {
   static constexpr int NUM_SHARDS = 8;
   auto executor_ = std::make_unique<folly::CPUThreadPoolExecutor>(4);
-  auto kv_store_ = std::make_unique<SynchronizedShardedMap<int64_t, float*>>(NUM_SHARDS, BLOCK_SIZE, BLOCK_ALIGNMENT);
+  auto kv_store_ = std::make_unique<SynchronizedShardedMap<int64_t, float*>>(
+      NUM_SHARDS, BLOCK_SIZE, BLOCK_ALIGNMENT);
+
+  auto start_time = std::chrono::steady_clock::now();
 
   // Insert test data
   for (int i = 0; i < 1000; ++i) {
     int shard_id = i % NUM_SHARDS;
     auto wlock = kv_store_->by(shard_id).wlock();
     auto* pool = kv_store_->pool_by(shard_id);
     auto* block = pool->allocate_t<float>();
+    ASSERT_NE(block, nullptr);
     FixedBlockPool::set_key(block, i);
-    FixedBlockPool::update_timestamp(block);  // Initial score
+    FixedBlockPool::update_timestamp(block);
     FixedBlockPool::set_used(block, true);
-    wlock->insert({i, block});
+    auto result = wlock->insert({i, block});
+    ASSERT_TRUE(result.second);
   }
+
   std::this_thread::sleep_for(std::chrono::seconds(5));
+  auto mid_time = std::chrono::steady_clock::now();
 
   for (int i = 1000; i < 2000; ++i) {
     int shard_id = i % NUM_SHARDS;
     auto wlock = kv_store_->by(shard_id).wlock();
     auto* pool = kv_store_->pool_by(shard_id);
     auto* block = pool->allocate_t<float>();
+    ASSERT_NE(block, nullptr);
     FixedBlockPool::set_key(block, i);
-    FixedBlockPool::update_timestamp(block); // Initial score
+    FixedBlockPool::update_timestamp(block);
     FixedBlockPool::set_used(block, true);
-    wlock->insert({i, block});
+    auto result = wlock->insert({i, block});
+    ASSERT_TRUE(result.second);
   }
 
+  auto current_time = std::chrono::steady_clock::now();
+  auto old_blocks_age = std::chrono::duration_cast<std::chrono::seconds>(
+                            current_time - start_time)
+                            .count();
+  auto new_blocks_age =
+      std::chrono::duration_cast<std::chrono::seconds>(current_time - mid_time)
+          .count();
+  ASSERT_GT(old_blocks_age, new_blocks_age);  // 确保时间差异存在
+
   std::unique_ptr<FeatureEvict<float>> feature_evict;
   int evict_trigger_mode = 2;
   int evict_trigger_strategy = 0;
   uint32_t ttl = 4;
-  // feature evict config
+
   FeatureEvictConfig feature_evict_config;
-  feature_evict_config.trigger_mode = static_cast<EvictTriggerMode>(evict_trigger_mode);
-  feature_evict_config.trigger_strategy = static_cast<EvictTriggerStrategy>(evict_trigger_strategy);
+  feature_evict_config.trigger_mode =
+      static_cast<EvictTriggerMode>(evict_trigger_mode);
+  feature_evict_config.trigger_strategy =
+      static_cast<EvictTriggerStrategy>(evict_trigger_strategy);
   feature_evict_config.ttl = ttl;
 
   if (feature_evict_config.trigger_mode != EvictTriggerMode::DISABLED) {
-    feature_evict = create_feature_evict(feature_evict_config, executor_.get(),*kv_store_.get(), 4);
+    feature_evict = create_feature_evict(
+        feature_evict_config, executor_.get(), *kv_store_.get(), 4);
+    ASSERT_NE(feature_evict, nullptr);
   }
 
   // Initial validation
@@ -144,28 +200,40 @@ TEST(FeatureEvictTest, TimeBasedEviction) {
 
   // Perform eviction
   feature_evict->trigger_evict();
+  ASSERT_TRUE(feature_evict->is_evicting());
 
   // Validate eviction process
-  while (feature_evict->is_evicting()) {
+  int max_iterations = 1000;
+  int iterations = 0;
+  while (feature_evict->is_evicting() && iterations < max_iterations) {
     feature_evict->resume();
     std::this_thread::sleep_for(std::chrono::microseconds(5));
     feature_evict->pause();
+    iterations++;
   }
+  ASSERT_LT(iterations, max_iterations);
+  ASSERT_FALSE(feature_evict->is_evicting());
 
   // Validate results
   size_t remaining = 0;
+  size_t newer_blocks = 0;
   for (int shard_id = 0; shard_id < NUM_SHARDS; ++shard_id) {
     auto rlock = kv_store_->by(shard_id).rlock();
     remaining += rlock->size();
+    for (const auto& [key, block] : *rlock) {
+      if (key >= 1000) newer_blocks++;
+      ASSERT_TRUE(FixedBlockPool::get_used(block));
+    }
   }
   std::cout << "remaining: " << remaining << std::endl;
   ASSERT_EQ(remaining, 1000);
+  ASSERT_GT(newer_blocks, 800);
 }
-
 TEST(FeatureEvictTest, TimeCounterBasedEviction) {
   static constexpr int NUM_SHARDS = 8;
   auto executor_ = std::make_unique<folly::CPUThreadPoolExecutor>(4);
-  auto kv_store_ = std::make_unique<SynchronizedShardedMap<int64_t, float*>>(NUM_SHARDS, BLOCK_SIZE, BLOCK_ALIGNMENT);
+  auto kv_store_ = std::make_unique<SynchronizedShardedMap<int64_t, float*>>(
+      NUM_SHARDS, BLOCK_SIZE, BLOCK_ALIGNMENT);
 
   // Insert test data
   for (int i = 0; i < 500; ++i) {
@@ -198,7 +266,7 @@ TEST(FeatureEvictTest, TimeCounterBasedEviction) {
     auto* pool = kv_store_->pool_by(shard_id);
     auto* block = pool->allocate_t<float>();
     FixedBlockPool::set_key(block, i);
-    FixedBlockPool::update_timestamp(block); // Initial score
+    FixedBlockPool::update_timestamp(block);  // Initial score
     FixedBlockPool::set_count(block, 2);
     FixedBlockPool::set_used(block, true);
     wlock->insert({i, block});
@@ -213,14 +281,17 @@ TEST(FeatureEvictTest, TimeCounterBasedEviction) {
 
   // feature evict config
   FeatureEvictConfig feature_evict_config;
-  feature_evict_config.trigger_mode = static_cast<EvictTriggerMode>(evict_trigger_mode);
-  feature_evict_config.trigger_strategy = static_cast<EvictTriggerStrategy>(evict_trigger_strategy);
+  feature_evict_config.trigger_mode =
+      static_cast<EvictTriggerMode>(evict_trigger_mode);
+  feature_evict_config.trigger_strategy =
+      static_cast<EvictTriggerStrategy>(evict_trigger_strategy);
   feature_evict_config.ttl = ttl;
   feature_evict_config.count_threshold = count_threshold;
   feature_evict_config.count_decay_rate = count_decay_rate;
 
   if (feature_evict_config.trigger_mode != EvictTriggerMode::DISABLED) {
-    feature_evict = create_feature_evict(feature_evict_config, executor_.get(),*kv_store_.get(), 4);
+    feature_evict = create_feature_evict(
+        feature_evict_config, executor_.get(), *kv_store_.get(), 4);
   }
 
   // Initial validation
@@ -254,7 +325,8 @@ TEST(FeatureEvictTest, TimeCounterBasedEviction) {
 TEST(FeatureEvictTest, L2WeightBasedEviction) {
   static constexpr int NUM_SHARDS = 8;
   auto executor_ = std::make_unique<folly::CPUThreadPoolExecutor>(4);
-  auto kv_store_ = std::make_unique<SynchronizedShardedMap<int64_t, float*>>(NUM_SHARDS, BLOCK_SIZE, BLOCK_ALIGNMENT);
+  auto kv_store_ = std::make_unique<SynchronizedShardedMap<int64_t, float*>>(
+      NUM_SHARDS, BLOCK_SIZE, BLOCK_ALIGNMENT);
   int dim = 4;
   std::vector<float> weight1(dim, 1.0);
   // Insert test data
@@ -288,12 +360,15 @@ TEST(FeatureEvictTest, L2WeightBasedEviction) {
   double l2_weight_threshold = 3.0;
   // feature evict config
   FeatureEvictConfig feature_evict_config;
-  feature_evict_config.trigger_mode = static_cast<EvictTriggerMode>(evict_trigger_mode);
-  feature_evict_config.trigger_strategy = static_cast<EvictTriggerStrategy>(evict_trigger_strategy);
+  feature_evict_config.trigger_mode =
+      static_cast<EvictTriggerMode>(evict_trigger_mode);
+  feature_evict_config.trigger_strategy =
+      static_cast<EvictTriggerStrategy>(evict_trigger_strategy);
   feature_evict_config.l2_weight_threshold = l2_weight_threshold;
 
   if (feature_evict_config.trigger_mode != EvictTriggerMode::DISABLED) {
-    feature_evict = create_feature_evict(feature_evict_config, executor_.get(),*kv_store_.get(), dim);
+    feature_evict = create_feature_evict(
+        feature_evict_config, executor_.get(), *kv_store_.get(), dim);
   }
 
   // Initial validation
@@ -327,7 +402,8 @@ TEST(FeatureEvictTest, L2WeightBasedEviction) {
 TEST(FeatureEvictTest, PerformanceTest) {
   static constexpr int NUM_SHARDS = 1;
   // Test configurations
-  const std::vector<int> test_sizes = {100'000, 500'000, 1'000'000, 5'000'000, 10'000'000};
+  const std::vector<int> test_sizes = {
+      100'000, 500'000, 1'000'000, 5'000'000, 10'000'000};
 
   fmt::print("\nPerformance Test Results:\n");
   fmt::print("{:<15} {:<15} {:<15}\n", "Size", "Time(ms)", "Items/ms");
@@ -336,8 +412,8 @@ TEST(FeatureEvictTest, PerformanceTest) {
   for (const auto& size : test_sizes) {
     // Create executor and store for each test size
     auto executor = std::make_unique<folly::CPUThreadPoolExecutor>(8);
-    auto kv_store =
-        std::make_unique<SynchronizedShardedMap<int64_t, float*>>(NUM_SHARDS, BLOCK_SIZE, BLOCK_ALIGNMENT, 1000);
+    auto kv_store = std::make_unique<SynchronizedShardedMap<int64_t, float*>>(
+        NUM_SHARDS, BLOCK_SIZE, BLOCK_ALIGNMENT, 1000);
 
     // Insert test data with different initial scores
     for (int i = 0; i < size; ++i) {
@@ -346,7 +422,8 @@ TEST(FeatureEvictTest, PerformanceTest) {
       auto* pool = kv_store->pool_by(shard_id);
       auto* block = pool->allocate_t<float>();
       FixedBlockPool::set_key(block, i);
-      FixedBlockPool::set_count(block, (i % 2) ? 1 : 2);  // Alternate between scores
+      FixedBlockPool::set_count(block,
+                                (i % 2) ? 1 : 2);  // Alternate between scores
       FixedBlockPool::set_used(block, true);
       wlock->insert({i, block});
     }
@@ -365,17 +442,22 @@ TEST(FeatureEvictTest, PerformanceTest) {
     }
 
     auto end_time = std::chrono::high_resolution_clock::now();
-    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count();
+    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(
+                        end_time - start_time)
+                        .count();
 
     std::size_t current_size = 0;
     for (int shard_id = 0; shard_id < NUM_SHARDS; ++shard_id) {
       auto wlock = kv_store->by(shard_id).wlock();
       current_size += wlock->size();
     }
-    double eviction_rate = static_cast<double>(size - current_size) / static_cast<double>(size);
+    double eviction_rate =
+        static_cast<double>(size - current_size) / static_cast<double>(size);
 
     // Print results
     fmt::print("{:<15d} {:<15d} {:<15.2f}\n", size, duration, eviction_rate);
+    ASSERT_LT(current_size, size);
+    ASSERT_GT(eviction_rate, 0.0);
   }
 }
 }  // namespace kv_mem