HiCache: Add @rank_consensus to various functions (#37425)

Co-authored-by: Zhangheng <hzh0425@apache.org>

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -569,6 +569,10 @@ class UnifiedRadixCache(BasePrefixCache):
     def is_chunk_cache(self) -> bool:
         return self.disable
 
+    @rank_consensus(
+        same_params=["len(params.key)"],
+        same_results=["result.prefix_len"],
+    )
     def insert(self, params: InsertParams) -> InsertResult:
         if self.disable:
             return InsertResult(prefix_len=0)
@@ -588,9 +592,11 @@ class UnifiedRadixCache(BasePrefixCache):
             # Drain still-pending actions so frees reach the allocator on abort.
             self._apply_cache_actions(self.tree_core.end_insert())
 
+    @rank_consensus(same_params=True, same_results=True)
     def evict(self, params: EvictParams) -> EvictResult:
         return self._evict(params)
 
+    @rank_consensus(same_params=True, same_results=True)
     def evict_for_alloc(self, params: EvictParams) -> EvictResult:
         """Evict until the requested component allocations become feasible.
 
@@ -942,6 +948,7 @@ class UnifiedRadixCache(BasePrefixCache):
             return DecLockRefResult()
         return self.tree_core.dec_host_lock_ref(node_id, params)
 
+    @rank_consensus(same_params=["req.rid", "is_insert", "kv_len_to_handle"])
     def cache_finished_req(
         self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int, **kwargs
     ) -> None:
@@ -1079,6 +1086,7 @@ class UnifiedRadixCache(BasePrefixCache):
             ):
                 self.session_refs.register_session_ref(req)
 
+    @rank_consensus(same_params=["req.rid", "chunked"])
     def cache_unfinished_req(self, req: Req, chunked: bool = False, **kwargs) -> None:
         if self.session.try_cache_unfinished_req(req, chunked=chunked, **kwargs):
             return
@@ -1797,6 +1805,7 @@ class UnifiedRadixCache(BasePrefixCache):
             )
         return transfers
 
+    @rank_consensus
     def write_backup_storage(self, node_id: NodeId) -> None:
         if not self.enable_storage or self.cache_controller is None:
             return
@@ -1891,6 +1900,7 @@ class UnifiedRadixCache(BasePrefixCache):
         storage_hit_count -= storage_hit_count % self.page_size
         return storage_hit_count
 
+    @rank_consensus(same_params=["req_id", "len(new_input_tokens)"])
     def prefetch_from_storage(
         self,
         req_id: str,
@@ -3086,6 +3096,7 @@ class UnifiedRadixCache(BasePrefixCache):
 
     # ---- HiCache: Scheduler Entry Points ----
 
+    @rank_consensus(same_params=["params.host_hit_length"])
     def init_load_back(
         self,
         params: InitLoadBackParams,

python/sglang/test/server_fixtures/default_fixture.py

@@ -45,6 +45,8 @@ class DefaultServerBase(CustomTestCase):
     base_url = DEFAULT_URL_FOR_TEST
     timeout = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
     other_args: list[str] = []
+    # Extra env vars passed to the launched server subprocess.
+    server_env: dict = None
 
     # For OpenAI API settings
     api_key = "sk-123456"
@@ -55,12 +57,16 @@ class DefaultServerBase(CustomTestCase):
 
         # Set OpenAI API key and base URL environment variables.
         # Needed for lmm-evals to work.
+        kwargs = {}
+        if cls.server_env:
+            kwargs["env"] = cls.server_env
         with openai_api_env(cls.api_key):
             cls.process = popen_launch_server(
                 cls.model,
                 cls.base_url,
                 timeout=cls.timeout,
                 other_args=cls.other_args,
+                **kwargs,
             )
 
     @classmethod

python/sglang/test/test_utils.py

@@ -666,6 +666,7 @@ def unified_radix_tree_server_env(
     return {
         **os.environ,
         **extra_env,
+        "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
         "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
         "SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND": tree_core_backend,
     }