chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[BUG FIX] add fail check when get fail in case wait complete block (#9971)

Browse Source 
Co-authored-by: mashisong <mashisong@bytedance.com>
Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
This commit is contained in:
Shisong Ma
2025-09-08 09:34:04 +08:00
committed by GitHub
parent b0fcbb74d0
commit 33467c05a4
2 changed files with 21 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/sglang/srt/managers/cache_controller.py
View File

@@ -207,26 +207,25 @@ class PrefetchOperation(StorageOperation):
):
self.request_id = request_id
self._done_flag = False
self._lock = threading.Lock()
self._terminated_flag = False
self.start_time = time.monotonic()
super().__init__(host_indices, token_ids, last_hash)
def increment(self, num_tokens: int):
with self._lock:
if self._done_flag:
if self._terminated_flag:
return False
self.completed_tokens += num_tokens
return True
def mark_done(self):
def mark_terminate(self):
with self._lock:
self._done_flag = True
self._terminated_flag = True
def is_done(self) -> bool:
return self._done_flag
def is_terminated(self) -> bool:
return self._terminated_flag
class HiCacheController:
@@ -628,7 +627,7 @@ class HiCacheController:
return operation
def terminate_prefetch(self, operation):
operation.mark_done()
operation.mark_terminate()
return operation.completed_tokens, operation.hash_value
def append_host_mem_release(self, host_indices: torch.Tensor):
@@ -709,6 +708,7 @@ class HiCacheController:
operation.completed_tokens
!= prev_completed_tokens + len(batch_hashes) * self.page_size
):
operation.mark_terminate()
break # Some operations fail or operation terminated by controller
# release pre-allocated memory
self.append_host_mem_release(