update unit tests for new version (taco-project#79)

* update test cache engine

* update test cache engine accel

* remove some tests

Author: zhuofan1123

flexkv/cache/mempool.py

@@ -20,7 +20,7 @@ def __init__(
         self._free_ids_offset = 0
 
     def reset(self) -> None:
-        self._free_mask.fill_(True)
+        self._free_mask.fill(True)
         self._num_free = self.num_total_blocks
         self._free_ids = self._free_mask.nonzero()[0]
         self._free_ids_offset = 0

flexkv/cache/transfer_pattern.py

@@ -29,259 +29,6 @@ def add_virtal_op_for_mutiple_finished_ops(
             graph.add_dependency(op.op_id, op_id)
         return graph, op.op_id
 
-def create_read_graph_cpu_storage(
-    gpu_blocks: torch.Tensor,
-    cpu_blocks: torch.Tensor,
-    ssd_blocks: torch.Tensor,
-    gpu_device_id: int = 0,
-    layer_num: int = 1,
-    graph: Optional[TransferOpGraph] = None,
-)->Tuple[TransferOpGraph, List[int]]:
-    """
-    Create a read transfer graph with (REMOTE_STORAGE / SSD)->CPU->GPU operations
-    ssd_blocks: the blocks of ssd that are used as a lower-level storage backend,
-    including ssd or remote storage. This can be empty, which means cpu-only kvcache.
-    Returns:
-        graph: TransferOpGraph
-        ops_to_be_tracked: List[int]: a list of transfer ops that can indicate
-        the completion of some key operations
-    """
-    assert len(gpu_blocks) == len(cpu_blocks)
-    if graph is None:
-        graph = TransferOpGraph()
-    assert len(gpu_blocks) > 0
-    if len(ssd_blocks) == 0:
-        op = TransferOp(
-            graph_id = graph.graph_id,
-            transfer_type = TransferType.H2D,
-            src_block_ids = cpu_blocks,
-            dst_block_ids = gpu_blocks,
-            layer_id = 0,
-            layer_granularity = layer_num,
-        )
-        graph.add_transfer_op(op)
-        return graph, [op.op_id]
-    elif len(ssd_blocks) < len(cpu_blocks):
-        task_end_ops_ids = []
-        if len(ssd_blocks) > 0:
-            op1 = TransferOp(
-                graph_id = graph.graph_id,
-                transfer_type = TransferType.DISK2H,
-                src_block_ids = ssd_blocks,
-                dst_block_ids = cpu_blocks[-len(ssd_blocks):],
-                layer_id = 0,
-                layer_granularity = layer_num
-            )
-            graph.add_transfer_op(op1)
-            op2 = TransferOp(
-                graph_id = graph.graph_id,
-                transfer_type = TransferType.H2D,
-                src_block_ids = cpu_blocks[-len(ssd_blocks):],
-                dst_block_ids = gpu_blocks[-len(ssd_blocks):],
-                layer_id = 0,
-                layer_granularity = layer_num
-            )
-            graph.add_transfer_op(op2)
-            graph.add_dependency(op2.op_id, op1.op_id)
-            task_end_ops_ids.append(op2.op_id)
-        op3 = TransferOp(
-            graph_id = graph.graph_id,
-            transfer_type = TransferType.H2D,
-            src_block_ids = cpu_blocks[:len(cpu_blocks) - len(ssd_blocks)],
-            dst_block_ids = gpu_blocks[:len(cpu_blocks) - len(ssd_blocks)],
-            layer_id = 0,
-            layer_granularity = layer_num
-        )
-        graph.add_transfer_op(op3)
-        task_end_ops_ids.append(op3.op_id)
-        return graph, task_end_ops_ids
-    else:
-        op1 = TransferOp(
-            graph_id = graph.graph_id,
-            transfer_type=TransferType.DISK2H,
-            src_block_ids=ssd_blocks,
-            dst_block_ids=cpu_blocks,
-            layer_id = 0,
-            layer_granularity = layer_num
-        )
-        graph.add_transfer_op(op1)
-        op2 = TransferOp(
-            graph_id = graph.graph_id,
-            transfer_type=TransferType.H2D,
-            src_block_ids=cpu_blocks,
-            dst_block_ids=gpu_blocks,
-            layer_id = 0,
-            layer_granularity = layer_num
-        )
-        graph.add_transfer_op(op2)
-        graph.add_dependency(op2.op_id, op1.op_id)
-        return graph, [op2.op_id]
-
-def create_read_graph_cpu_ssd_remote(
-    gpu_blocks: torch.Tensor,
-    cpu_blocks: torch.Tensor,
-    ssd_blocks: torch.Tensor,
-    remote_blocks: torch.Tensor,
-    gpu_device_id: int = 0,
-    layer_num: int = 1,
-    write_back_to_ssd: bool = True,
-)->Tuple[TransferOpGraph, List[int]]:
-    """
-    Create a read transfer graph with (REMOTE_STORAGE + SSD)->CPU->GPU operations
-    Returns:
-        graph: TransferOpGraph
-        finished_ops_ids: List[int]: a list of transfer ops that can indicate
-        the completion of each layer or each layer for each tp rank
-    """
-    graph = TransferOpGraph()
-    finished_ops_ids: List[int] = []
-    if len(remote_blocks) == 0:
-        graph, finished_ops_ids = create_read_graph_cpu_storage(gpu_blocks=gpu_blocks,
-                                                            cpu_blocks=cpu_blocks,
-                                                            ssd_blocks=ssd_blocks,
-                                                            gpu_device_id=gpu_device_id,
-                                                            layer_num=layer_num,
-                                                            graph=graph)
-        if len(finished_ops_ids) > 0:
-            graph, finished_ops_ids = add_virtal_op_for_mutiple_finished_ops(graph, finished_ops_ids)
-        assert len(finished_ops_ids) > 0
-        return graph, finished_ops_ids
-    else:
-        if len(remote_blocks) < len(gpu_blocks):
-            graph, finished_ops_ids = create_read_graph_cpu_storage(gpu_blocks=gpu_blocks[:-len(remote_blocks)],
-                                                                cpu_blocks=cpu_blocks[:-len(remote_blocks)],
-                                                                ssd_blocks=ssd_blocks[:-len(remote_blocks)],
-                                                                gpu_device_id=gpu_device_id,
-                                                                layer_num=layer_num,
-                                                                graph=graph)
-    op_r2h = TransferOp(
-        graph_id = graph.graph_id,
-        transfer_type = TransferType.REMOTE2H,
-        src_block_ids = remote_blocks,
-        dst_block_ids = cpu_blocks[-len(remote_blocks):],
-        layer_id = 0,
-        layer_granularity = layer_num
-    )
-    graph.add_transfer_op(op_r2h)
-    op_h2d = TransferOp(
-        graph_id = graph.graph_id,
-        transfer_type = TransferType.H2D,
-        src_block_ids = cpu_blocks[-len(remote_blocks):],
-        dst_block_ids = gpu_blocks[-len(remote_blocks):],
-        layer_id = 0,
-        layer_granularity = layer_num
-    )
-    graph.add_transfer_op(op_h2d)
-    graph.add_dependency(op_h2d.op_id, op_r2h.op_id)
-    if write_back_to_ssd:
-        op_h2disk = TransferOp(
-            graph_id = graph.graph_id,
-            transfer_type = TransferType.H2DISK,
-            src_block_ids = cpu_blocks[-len(remote_blocks):],
-            dst_block_ids = ssd_blocks[-len(remote_blocks):],
-            layer_id = 0,
-            layer_granularity = layer_num
-        )
-        graph.add_transfer_op(op_h2disk)
-        graph.add_dependency(op_h2disk.op_id, op_r2h.op_id)
-    finished_ops_ids.append(op_h2d.op_id)
-    if len(finished_ops_ids) > 0:
-        graph, finished_ops_ids = add_virtal_op_for_mutiple_finished_ops(graph, finished_ops_ids)
-    return graph, finished_ops_ids
-
-def create_write_graph_cpu_storage(
-    gpu_blocks: torch.Tensor,
-    cpu_blocks: torch.Tensor,
-    ssd_blocks: torch.Tensor,
-    gpu_device_id: int = 0,
-    layer_num: int = 1,
-    graph: Optional[TransferOpGraph] = None,
-)->Tuple[TransferOpGraph, List[int]]:
-    """
-    Create a write transfer graph with CPU->REMOTE_STORAGE / SSD operations
-    ssd_blocks: the blocks of ssd that are used as a lower-level storage backend,
-    including ssd or remote storage. This can be empty, which means cpu-only kvcache.
-    Write op granularity is larger: gpu->cpu is put into the same op.
-    Returns:
-        graph: TransferOpGraph
-        layer_wise_ops: List[int]: a list of transfer ops that can indicate
-        the completion of each layer or each layer for each tp rank
-    """
-    if graph is None:
-        graph = TransferOpGraph()
-    op_d2h = TransferOp(
-        graph_id = graph.graph_id,
-        transfer_type = TransferType.D2H,
-        src_block_ids = gpu_blocks,
-        dst_block_ids = cpu_blocks[-len(gpu_blocks):],
-        layer_id = 0,
-        layer_granularity = layer_num
-    )
-    graph.add_transfer_op(op_d2h)
-    if len(ssd_blocks) == 0:
-        return graph, [op_d2h.op_id]
-    else:
-        op_h2disk = TransferOp(
-            graph_id = graph.graph_id,
-            transfer_type = TransferType.H2DISK,
-            src_block_ids = cpu_blocks[-len(ssd_blocks):],
-            dst_block_ids = ssd_blocks,
-            layer_id = 0,
-            layer_granularity = layer_num
-        )
-        graph.add_transfer_op(op_h2disk)
-        graph.add_dependency(op_h2disk.op_id, op_d2h.op_id)
-        return graph, [op_d2h.op_id]
-
-def create_write_graph_cpu_ssd_remote(
-    gpu_blocks: torch.Tensor,
-    cpu_blocks: torch.Tensor,
-    ssd_blocks: torch.Tensor,
-    remote_blocks: torch.Tensor,
-    gpu_device_id: int = 0,
-    layer_num: int = 1,
-)->Tuple[TransferOpGraph, List[int]]:
-    """
-    Create a write transfer graph with CPU->REMOTE_STORAGE + SSD operations
-    Returns:
-        graph: TransferOpGraph
-        layer_wise_ops: List[int]: a list of transfer ops that can indicate
-        the completion of each layer or each layer for each tp rank
-    """
-    graph = TransferOpGraph()
-    op_d2h = TransferOp(
-        graph_id = graph.graph_id,
-        transfer_type = TransferType.D2H,
-        src_block_ids = gpu_blocks,
-        dst_block_ids = cpu_blocks[-len(gpu_blocks):],
-        layer_id = 0,
-        layer_granularity = layer_num
-    )
-    graph.add_transfer_op(op_d2h)
-    if len(ssd_blocks) != 0:
-        op_h2disk = TransferOp(
-            graph_id = graph.graph_id,
-            transfer_type = TransferType.H2DISK,
-            src_block_ids = cpu_blocks[-len(gpu_blocks):],
-            dst_block_ids = ssd_blocks,
-            layer_id = 0,
-            layer_granularity = layer_num
-        )
-        graph.add_transfer_op(op_h2disk)
-        graph.add_dependency(op_h2disk.op_id, op_d2h.op_id)
-    if len(remote_blocks) != 0:
-        op_h2remote = TransferOp(
-            graph_id = graph.graph_id,
-            transfer_type = TransferType.H2REMOTE,
-            src_block_ids = cpu_blocks[-len(remote_blocks):],
-            dst_block_ids = remote_blocks,
-            layer_id = 0,
-            layer_granularity = layer_num
-        )
-        graph.add_transfer_op(op_h2remote)
-        graph.add_dependency(op_h2remote.op_id, op_d2h.op_id)
-    return graph, [op_d2h.op_id]
-
 def convert_read_graph_to_layer_wise_graph(
     transfer_graph: TransferOpGraph,
     finished_ops_ids: List[int],

tests/conftest.py

@@ -4,9 +4,3 @@
 """
 # Import fixtures from test_utils so pytest can discover them
 from test_utils import model_config, cache_config, test_config
-
-import multiprocessing as mp
-
-# Set the start method for multiprocessing to 'spawn'
-# This ensures consistent behavior across different platforms
-mp.set_start_method("spawn", force=True)