GC_DP_FIX

Author: UESTC-AHao

ucm/integration/vllm/ucm_connector.py

@@ -312,6 +312,7 @@ def _create_store(
 
         name = self.connector_configs[0]["ucm_connector_name"]
         module_path = self.connector_configs[0].get("ucm_connector_module_path", None)
+        dp_rank = self._vllm_config.parallel_config.rank
         config = copy.deepcopy(self.connector_configs[0]["ucm_connector_config"])
         config.setdefault("share_buffer_enable", self.is_mla)
         if "storage_backends" in config:
@@ -326,6 +327,7 @@ def _create_store(
             config["shard_size"] = kv_cache_layout.shard_size * self.blocks_per_chunk
             config["block_size"] = kv_cache_layout.block_size * self.blocks_per_chunk
             config["local_rank_size"] = self.tp_size if self.is_mla else 1
+            config["posix_gc_enable"] = False
             if cpu_affinity_cores:
                 config["cpu_affinity_cores"] = list(cpu_affinity_cores)
         else:
@@ -336,10 +338,6 @@ def _create_store(
                 * (1 if self.is_mla else self.num_head * 2)
                 * self.blocks_per_chunk
             )
-        # GC only enabled for Scheduler with data_parallel_rank == 0
-        dp_rank = self._vllm_config.parallel_config.data_parallel_rank
-        if self._role == KVConnectorRole.WORKER or dp_rank != 0:
-            config["posix_gc_enable"] = False
 
         logger.info(f"create {name} with config: {config}")
         return UcmConnectorFactoryV1.create_connector(name, config, module_path)

ucm/store/posix/cc/global_config.h

@@ -50,7 +50,7 @@ struct Config {
     size_t posixGcCheckIntervalSec{30};
     size_t posixCapacityGb{0};
     double posixGcTriggerThresholdRatio{0.7};
-    size_t posixGcMaxRecycleCountPerShard{1000};
+    size_t posixGcMaxRecycleCountPerShard{5000};
     double posixGcShardSampleRatio{0.1};
 };
 

ucm/store/posix/cc/space_manager.cc

@@ -29,12 +29,15 @@ namespace UC::PosixStore {
 
 Status SpaceManager::Setup(const Config& config)
 {
+    hotnessTrackerEnable_ = config.deviceId == -1;
     gcEnable_ = config.posixGcEnable && config.posixCapacityGb > 0;
     auto s = layout_.Setup(config);
     if (s.Failure()) [[unlikely]] { return s; }
-    if (gcEnable_) {
+    if (hotnessTrackerEnable_) {
         s = hotnessTracker_.Setup(&layout_);
         if (s.Failure()) [[unlikely]] { return s; }
+    }
+    if (gcEnable_) {
         s = gcMgr_.Setup(&layout_, config);
         if (s.Failure()) [[unlikely]] { return s; }
     }
@@ -129,7 +132,7 @@ void SpaceManager::OnLookupPrefix(PrefixLookupContext& ctx)
             }
             break;
         }
-        if (gcEnable_) { hotnessTracker_.Touch(*(ctx.blocks + i)); }
+        if (hotnessTrackerEnable_) { hotnessTracker_.Touch(*(ctx.blocks + i)); }
     }
     ctx.waiter->Done();
 }

ucm/store/posix/cc/space_manager.h

@@ -49,6 +49,7 @@ class SpaceManager {
     ThreadPool<PrefixLookupContext> prefixLookupSrv_;
     HotnessTracker hotnessTracker_;
     ShardGarbageCollector gcMgr_;
+    bool hotnessTrackerEnable_{false};
     bool gcEnable_{false};
 
 public: