From 3066ba8167eaf188872cf0ba409a5d23752a78f5 Mon Sep 17 00:00:00 2001 From: chenkaiyue Date: Tue, 28 Apr 2026 09:50:43 +0800 Subject: [PATCH] fix(hicache): add retry logic for MooncakeStore warmup (#17195) Co-authored-by: Teng Ma Co-authored-by: Zhiqiang Xie --- .../storage/mooncake_store/mooncake_store.py | 25 +++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index 9d65a7265..0e8d162dc 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -389,13 +389,15 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): ) logger.info("Mooncake store setup successfully.") + self.local_rank = ( + storage_config.tp_rank if storage_config is not None else 0 + ) self.warmup() logger.info("Mooncake store warmup successfully.") self.enable_storage_metrics = False if storage_config is not None: self.is_mla_backend = storage_config.is_mla_model - self.local_rank = storage_config.tp_rank self.pp_rank = storage_config.pp_rank self.pp_size = storage_config.pp_size self.attn_cp_rank = storage_config.attn_cp_rank @@ -483,7 +485,26 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): def warmup(self): warmup_key = "sglang_mooncake_store_warmup_key" + uuid.uuid4().hex warmup_value = bytes(4 * 1024) # 4 KB - assert self.store.put(warmup_key, warmup_value) == 0 + + # Retry logic to handle Transfer Engine startup race condition + max_retries = 10 + retry_delay = 1.0 # seconds + + for attempt in range(max_retries): + ret = self.store.put(warmup_key, warmup_value) + if ret == 0: + break + logger.warning( + f"[TP{self.local_rank}] Warmup put failed (attempt {attempt + 1}/{max_retries}), " + f"ret={ret}, retrying in {retry_delay}s..." + ) + time.sleep(retry_delay) + else: + raise RuntimeError( + f"[TP{self.local_rank}] Warmup put failed after {max_retries} attempts, " + "Transfer Engine might not be ready" + ) + assert self.store.is_exist(warmup_key) == 1 assert self.store.get(warmup_key) == warmup_value