[wip]fix duplicate create (#220)

* fix duplicate dump

* move create to update_state

* fix idx bug

* modify log

* fix commit bug

Author: qyh111

ucm/integration/vllm/uc_connector.py

@@ -39,6 +39,7 @@
 from vllm.distributed.parallel_state import get_world_group
 from vllm.v1.core.kv_cache_utils import hash_request_tokens
 from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.request import Request, RequestStatus
 
 from ucm.logger import init_logger
 from ucm.store.base import Task
@@ -48,7 +49,6 @@
     from vllm.attention.backends.abstract import AttentionMetadata
     from vllm.forward_context import ForwardContext
     from vllm.v1.core.kv_cache_manager import KVCacheBlocks
-    from vllm.v1.request import Request
 
 logger = init_logger(__name__)
 
@@ -533,12 +533,11 @@ def get_num_new_matched_tokens(
             the number of tokens that can be loaded from the
             external KV cache beyond what is already computed.
         """
-        # When the request is preempt req, need to commit succeed dumped blocks
-        # to avoid duplicate invoking create/commit funcs. Only preempt reqs
-        # whose succeed_dumped_blocks is non-empty need this check.
-        if hasattr(request, "succeed_dumped_blocks") and request.succeed_dumped_blocks:
-            self.connector.commit(request.succeed_dumped_blocks, True)
-            request.succeed_dumped_blocks.clear()
+        logger.info(f"get_num_new_matched_tokens request {request.request_id}.")
+
+        if request.status == RequestStatus.PREEMPTED:
+            logger.info(f"Handle preempted request {request.request_id}.")
+            self.request_finished(request, [])
 
         def md5(input) -> int:
             input_bytes = pickle.dumps(input, protocol=pickle.HIGHEST_PROTOCOL)
@@ -598,17 +597,6 @@ def md5(input) -> int:
                 self._need_load_reqs[request.request_id] = []
                 return num_lookup_hits * self.block_size, True
 
-        # Create blocks for the remaining (unmatched) blocks
-        if num_lookup_hits < len(remain_hashes):
-            remaining_hashes = remain_hashes[num_lookup_hits:]
-            create_results = self.connector.create(remaining_hashes)
-            logger.info(f"\ncreate_results on storage: {create_results}\n")
-            for j, ret in enumerate(create_results):
-                idx = num_lookup_hits + j
-                block_operations[start_position + idx] = (
-                    BlockOperation.DUMP if ret == 0 else BlockOperation.NONE
-                )
-
         # When all the tokens are cached in ssd or hbm,
         # we need to recompute the last token. This if condition will be removed
         # once vLLM's scheduler provides a better solution in the future.
@@ -638,6 +626,23 @@ def update_state_after_alloc(
             )
             self._need_load_reqs[request.request_id] = local_block_ids
 
+        request_block_info = self.request_block_infos.get(request.request_id, None)
+        if request_block_info:
+            start_position = request_block_info.start_position
+            block_operations = request_block_info.block_operations
+            block_hashes = request_block_info.block_hashes
+            start_create_pos = start_position + num_external_tokens // self.block_size
+            remaining_hashes = block_hashes[start_create_pos:]
+            if remaining_hashes:
+                create_results = self.connector.create(remaining_hashes)
+                if any(ret != 0 for ret in create_results):
+                    logger.warning(f"\ncreate_results on storage: {create_results}\n")
+                for j, ret in enumerate(create_results):
+                    idx = start_create_pos + j
+                    block_operations[idx] = (
+                        BlockOperation.DUMP if ret == 0 else BlockOperation.NONE
+                    )
+
     def build_connector_meta(
         self, scheduler_output: SchedulerOutput
     ) -> KVConnectorMetadata:
@@ -733,7 +738,6 @@ def request_finished(
     ) -> tuple[bool, Optional[dict[str, Any]]]:
         block_info = self.request_block_infos.pop(request.request_id, None)
         if hasattr(request, "succeed_dumped_blocks") and request.succeed_dumped_blocks:
-            logger.debug(f"commit {request.succeed_dumped_blocks} to True.")
             self.connector.commit(request.succeed_dumped_blocks, True)
         if block_info is not None:
             cancel_blocks = [
@@ -744,8 +748,8 @@ def request_finished(
                 and block_info.block_hashes[i] not in request.succeed_dumped_blocks
             ]
             if cancel_blocks:
-                logger.warning(f"commit {cancel_blocks} to False.")
                 self.connector.commit(cancel_blocks, False)
+        request.succeed_dumped_blocks.clear()
         return False, None
 
     def _extract_blocks(

ucm/store/ucm_nfs_store.py

@@ -80,11 +80,6 @@ def create(self, block_ids: List[str]) -> List[int]:
         rets = ucmnfsstore.AllocBatch(block_ids)
         if rets and all(ret == 0 for ret in rets):
             logger.debug("Succeed in allocating kv cache space.")
-        else:
-            failed_blocks = [block_ids[i] for i, ret in enumerate(rets) if ret != 0]
-            logger.warning(
-                f"Failed to allocate kv cache space for blocks: {failed_blocks}."
-            )
         return rets
 
     def lookup(self, block_ids: List[str]) -> List[bool]:
@@ -184,6 +179,7 @@ def commit(self, block_ids: List[str], is_success: bool = True) -> None:
             block_ids (List[str]): vLLM block hash.
             is_success(bool): if False, we need release block
         """
+        logger.debug(f"commit {block_ids} to {is_success}.")
         ucmnfsstore.CommitBatch(block_ids, is_success)
 
     def check(self, task: Task) -> int: