diff --git a/flexkv/kvtask.py b/flexkv/kvtask.py index 486f2bd2c..07cdcbcce 100644 --- a/flexkv/kvtask.py +++ b/flexkv/kvtask.py @@ -7,7 +7,6 @@ from typing import Callable import multiprocessing as mp import copy -from expiring_dict import ExpiringDict import nvtx import numpy as np @@ -99,6 +98,8 @@ class KVTask: prefetch_has_swa_remote: bool = False prefetch_namespace: Optional[List[str]] = None prefetch_swa_aware: bool = False + cancel_requested: bool = False + prefetch_key: Optional[int] = None def is_completed(self) -> bool: return self.status in [TaskStatus.COMPLETED, TaskStatus.CANCELLED, TaskStatus.FAILED] @@ -111,6 +112,8 @@ def shed_heavy_resources(self) -> None: self.slot_mapping = None self.token_mask = None self.callback = None + self.op_callback_dict = {} + self.swa_slot_mapping = None TASK_STATUS_TO_RESPONSE_STATUS = { TaskStatus.COMPLETED: KVResponseStatus.SUCCESS, @@ -214,10 +217,12 @@ def __init__(self, self.transfer_handles.append(remote_handle) remote_handle._handle.send_config_to_remotes() - self.tasks: ExpiringDict[int, KVTask] = ExpiringDict(ttl=1800) # 30 minutes + # Active transfers own blocks and callbacks until graph completion. + # A TTL must never evict that ownership while DMA is still in flight. + self.tasks: Dict[int, KVTask] = {} # hash(token_ids) -> task_id - self.prefetch_tasks: ExpiringDict[int, int] = ExpiringDict(ttl=1800) # 30 minutes + self.prefetch_tasks: Dict[int, int] = {} # removed with the owning task self._gen_prefetch_key = lambda token_ids, namespace: hash_token(token_ids, namespace) self.graph_to_task: Dict[int, int] = {} @@ -424,7 +429,9 @@ def create_prefetch_task(self, prefetch_namespace=namespace, prefetch_swa_aware=swa_aware) - self.prefetch_tasks[self._gen_prefetch_key(token_ids, namespace)] = task_id + key = self._gen_prefetch_key(token_ids, namespace) + self.tasks[task_id].prefetch_key = key + self.prefetch_tasks[key] = task_id self.graph_to_task[graph.graph_id] = task_id self._log_task_created(self.tasks[task_id]) @@ -536,14 +543,6 @@ def _update_tasks(self, timeout: float = 0.001) -> None: completed_op.xfer_ms, completed_op.e2e_ms, ) - if task.status == TaskStatus.CANCELLED and task.callback is None: - # Cache was reset while this task was in flight: reset_cache() - # cleared its callbacks and freed the radix nodes / mempool blocks - flexkv_logger.warning( - f"task {task_id}: transfer op {completed_op.op_id} completed " - "after reset_cache(), callback no longer exists and will be skipped." - ) - continue has_callback = completed_op.op_id in task.op_callback_dict if has_callback: try: @@ -737,7 +736,8 @@ def _fail_task(self, task_id: int) -> None: flexkv_logger.error(f"[KVTaskEngine] task {task_id} FAILED: a transfer " f"op of graph {task.graph.graph_id} failed") self._abort_task_plans(task) - task.status = TaskStatus.FAILED + task.status = TaskStatus.CANCELLED if task.cancel_requested else TaskStatus.FAILED + self._log_task_terminal(task, task.status) task.task_end_op_finished = True self.graph_to_task.pop(task.graph.graph_id, None) task.shed_heavy_resources() @@ -745,22 +745,19 @@ def _fail_task(self, task_id: int) -> None: self._release_task(task_id) def _cancel_task(self, task_id: int) -> None: - if task_id not in self.tasks: + task = self.tasks.get(task_id) + if task is None: + return + task.request_returned = True + if task.status == TaskStatus.RUNNING: + # Cancellation cannot stop submitted DMA. Keep graph mappings, + # plan pins and allocations until every transfer handle drains. + task.cancel_requested = True return - task = self.tasks[task_id] if not task.is_completed(): - # A task whose graph never launched still holds everything its - # plan acquired at create time: locked radix nodes and staging - # blocks that only a completion callback could mount on the tree. - # Dropping the task without aborting leaks all of it -- the staging - # blocks become unreachable (mempool exhaustion) and the pinned - # nodes stay unevictable. Abort rolls those back; RUNNING tasks - # keep the old behavior (their graph is in flight and completion - # callbacks will still fire). - if task.status in (TaskStatus.UNREADY, TaskStatus.READY): - self._abort_task_plans(task) + self._abort_task_plans(task) task.status = TaskStatus.CANCELLED - self._log_task_terminal(task, TaskStatus.CANCELLED) + self._log_task_terminal(task, task.status) self._release_task(task_id) def check_completed(self, task_id: int, completely: bool = False) -> bool: @@ -777,7 +774,7 @@ def check_completed(self, task_id: int, completely: bool = False) -> bool: # A partial-capable backend may finish the data-path sink after already # reporting failed blocks. Wait for graph completion so cleanup runs and # the caller observes FAILED instead of an early RUNNING-as-success result. - if task.transfer_failed: + if task.transfer_failed or task.cancel_requested: return task.is_completed() # For tasks with callback (e.g., PUT tasks that need to call insert_and_publish), # we must wait until _mark_completed is called (i.e., is_completed() returns True) @@ -846,13 +843,17 @@ def _release_task(self, task_id: int) -> None: task = self.tasks[task_id] if task.graph is not None: self.graph_to_task.pop(task.graph.graph_id, None) + if task.prefetch_key is not None and self.prefetch_tasks.get(task.prefetch_key) == task_id: + self.prefetch_tasks.pop(task.prefetch_key, None) self.tasks.pop(task_id, None) def _mark_completed(self, task_id: int) -> None: task = self.tasks[task_id] if task.is_completed(): return - if task.callback: + if task.cancel_requested: + self._abort_task_plans(task) + elif task.callback: callbacks = ( task.callback if isinstance(task.callback, list) else [task.callback] @@ -872,9 +873,10 @@ def _mark_completed(self, task_id: int) -> None: if task.task_type == TaskType.PREFETCH: self._finalize_prefetch_return_mask(task) task.status = ( + TaskStatus.CANCELLED if task.cancel_requested else TaskStatus.FAILED if task.transfer_failed else TaskStatus.COMPLETED) task.task_end_op_finished = True - self._log_task_terminal(task, TaskStatus.COMPLETED) + self._log_task_terminal(task, task.status) self.graph_to_task.pop(task.graph.graph_id, None) task.shed_heavy_resources() if task.request_returned: @@ -1317,7 +1319,9 @@ def prefetch_async(self, if task_id == -1: task_id = self._gen_task_id() nvtx.push_range(f"prefetch match: task_id={task_id}", color=get_nvtx_default_color()) - self.create_prefetch_task(task_id, token_ids, dp_client_id=dp_client_id, namespace=namespace, swa_aware=swa_aware) + self.create_prefetch_task( + task_id, token_ids, dp_client_id=dp_client_id, namespace=namespace, swa_aware=swa_aware + ) self._process_empty_graph(task_id) nvtx.pop_range() # trace prefetch async request @@ -1467,37 +1471,21 @@ def reset_cache(self) -> None: Used after a weight update (e.g. verl RL rollout) so that KV computed against stale weights is never reused. - We do NOT drain or cancel in-flight transfers here. verl issues the - reset at a rollout/weight-update boundary where no new generation - requests are being served, so in practice no task is ongoing. If any - task IS still in flight we only warn: resetting the radix tree + - mempool is cheap and the stale-weight invalidation must not be blocked - on transfer completion. This mirrors vLLM's own reset_encoder_cache / - reset_mm_cache, which likewise only warn on has_unfinished_requests(). + The caller must drain transfers first. Resetting a busy cache can + recycle destinations while DMA still writes to them; clearing Python + callbacks cannot make that safe. """ - ongoing = sum(1 for t in list(self.tasks.values()) if not t.is_completed()) + ongoing = sum(1 for t in self.tasks.values() if not t.is_completed()) if ongoing: - flexkv_logger.warning( - f"reset_cache called while {ongoing} task(s) are still in flight; " - f"resetting anyway. In-flight transfers may target blocks that are " - f"being freed — ensure reset is issued at a quiesced boundary." - ) - - # Invalidate in-flight tasks' callbacks BEFORE dropping the cache. The - # callback / op_callback_dict partials close over the exact radix nodes - # and mempool blocks we are about to free; if a late transfer fired them - # after reset they would unlock or mount onto a deleted node, or recycle a - # block into the freshly-emptied mempool (which raises "already free"). - # Clear them here so the dispatch in _update_tasks has nothing to fire. - # Note: reset_cache() runs on the same thread as the callback dispatch - # (_update_tasks), so no lock is needed. We keep the graph_to_task - # mapping so a late-completing op still resolves to its task and warns. - for task_id, task in list(self.tasks.items()): - if task.is_completed(): - continue # already-fired callbacks are harmless - task.callback = None - task.op_callback_dict = {} - task.status = TaskStatus.CANCELLED + raise RuntimeError(f"reset_cache requires drained transfers; {ongoing} task(s) active") + + # Only terminal response metadata remains after the drain boundary. + self.tasks.clear() + self.prefetch_tasks.clear() + self.graph_to_task.clear() + self.uncompleted_ops.clear() + self.uncompleted_op_results.clear() + self.uncompleted_graphs.clear() # Drop index (radix tree) + mempool on every tier. CRadixTreeIndex (C++) # and the pure-Python RadixTreeIndex both expose reset(); GlobalCacheEngine diff --git a/tests/test_task_ownership.py b/tests/test_task_ownership.py new file mode 100644 index 000000000..12466f254 --- /dev/null +++ b/tests/test_task_ownership.py @@ -0,0 +1,144 @@ +"""CPU task ownership regressions; no DMA or worker process is launched.""" + +from types import SimpleNamespace as NS +from unittest.mock import Mock, patch + +import numpy as np +import pytest + +from flexkv.kvtask import KVTask, KVTaskEngine, KVTaskManager, TaskStatus, TaskType +from flexkv.common.transfer import CompletedOp + + +def engine_with_task(status=TaskStatus.RUNNING, kind=TaskType.GET, request_returned=False): + engine = KVTaskEngine.__new__(KVTaskEngine) + plan = Mock() + task = KVTask( + 7, + kind, + 11, + False, + status, + np.arange(4), + None, + None, + NS(graph_id=13, num_ops=2, _op_map={}), + np.ones(4, dtype=bool), + plan, + {}, + ) + task.request_returned = request_returned + task.prefetch_key = 19 + engine.tasks = {7: task} + engine.prefetch_tasks = {19: 7} + engine.graph_to_task = {13: 7} + engine.uncompleted_ops = {} + engine.uncompleted_op_results = {} + engine.uncompleted_graphs = {} + engine._log_task_terminal = Mock() + engine._finalize_prefetch_return_mask = Mock() + engine.tracer = Mock() + engine.cache_engine = Mock() + return engine, task, plan + + +@pytest.mark.parametrize("kind", list(TaskType)) +@pytest.mark.parametrize("failed", [False, True]) +def test_running_cancel_keeps_ownership_until_graph_drains(kind, failed): + engine, task, plan = engine_with_task(kind=kind) + engine.cancel_tasks([7]) + engine.cancel_tasks([7]) # idempotent while draining + assert task.cancel_requested and task.status == TaskStatus.RUNNING + assert engine.tasks[7] is task and engine.graph_to_task == {13: 7} + assert engine.prefetch_tasks == {19: 7} + assert not plan.mock_calls + task.task_end_op_finished = True + assert not engine.check_completed(7) # early data-path completion is insufficient + if failed: + engine._fail_task(7) + else: + engine._mark_completed(7) + plan.abort.assert_called_once_with() + plan.assert_not_called() + assert task.status == TaskStatus.CANCELLED + assert engine.tasks == engine.graph_to_task == engine.prefetch_tasks == {} + assert task.op_callback_dict == {} and task.graph is None + + +@pytest.mark.parametrize("status", [TaskStatus.UNREADY, TaskStatus.READY]) +def test_unlaunched_cancel_releases_immediately(status): + engine, task, plan = engine_with_task(status=status) + engine._cancel_task(7) + plan.abort.assert_called_once_with() + assert engine.tasks == engine.prefetch_tasks == engine.graph_to_task == {} + + +def test_terminal_response_wait_recycles_all_indexes(): + engine, task, plan = engine_with_task() + engine._mark_completed(7) + assert engine.tasks[7] is task + engine._update_tasks = Mock() + result = engine.try_wait([7]) + assert result[7].status.value == "success" + plan.assert_called_once_with() + assert engine.tasks == engine.prefetch_tasks == engine.graph_to_task == {} + + +def test_releasing_old_prefetch_does_not_remove_new_owner(): + engine, task, plan = engine_with_task(status=TaskStatus.READY) + engine.prefetch_tasks[19] = 8 + engine._cancel_task(7) + assert engine.prefetch_tasks == {19: 8} + + +@pytest.mark.parametrize("status", [TaskStatus.RUNNING, TaskStatus.READY, TaskStatus.UNREADY]) +def test_busy_reset_has_no_side_effects(status): + engine, task, plan = engine_with_task(status=status) + with pytest.raises(RuntimeError, match="requires drained transfers"): + engine.reset_cache() + assert engine.tasks[7] is task + assert engine.graph_to_task == {13: 7} + assert not engine.cache_engine.mock_calls and not plan.mock_calls + + +def test_drained_reset_clears_terminal_metadata(): + engine, task, plan = engine_with_task() + engine._mark_completed(7) + engine.reset_cache() + engine.cache_engine.reset.assert_called_once_with() + assert engine.tasks == engine.prefetch_tasks == engine.graph_to_task == {} + + +def test_constructor_uses_explicit_task_ownership(): + import flexkv.kvtask as module + + config = NS(enable_cpu=True, enable_remote=True, enable_gds=False, enable_kv_sharing=False, enable_nixl=False) + model = NS(use_trtllm_subprocess=False, nnodes=1) + with patch.object(module, "GlobalCacheEngine"), patch.object(module, "TransferManagerHandle"): + manager = KVTaskManager(model, config) + assert type(manager.tasks) is dict and type(manager.prefetch_tasks) is dict + + +@pytest.mark.parametrize("failed", [False, True]) +def test_cancel_waits_for_every_transfer_handle_before_abort(failed): + engine, task, plan = engine_with_task() + writer_commit = Mock() + task.op_callback_dict[11] = writer_commit + engine.required_completed_count = 2 + first = [CompletedOp(13, 11), CompletedOp(13, -1)] + second = [CompletedOp(13, 11), CompletedOp(13, -1, failed=failed)] + engine.transfer_handles = [ + NS(wait=Mock(side_effect=[first, []])), + NS(wait=Mock(side_effect=[[], second])), + ] + engine.cancel_tasks([7]) + engine._update_tasks(timeout=0) + assert engine.tasks[7] is task + plan.abort.assert_not_called() + writer_commit.assert_not_called() + engine._update_tasks(timeout=0) + writer_commit.assert_called_once() # completed writers can safely publish + plan.abort.assert_called_once_with() + plan.assert_not_called() + assert engine.tasks == engine.graph_to_task == engine.prefetch_tasks == {} + assert engine.uncompleted_graphs == engine.uncompleted_ops == {}