[KV Offload] Pass request context (vllm-project#39185)

Signed-off-by: omerpaz95 <omerpaz95@gmail.com>

Author: omerpaz95

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

@@ -32,8 +32,8 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
     # 3 blocks, store just the middle block (skip first and last)
     # blocks = [0, 1, 2], [3, 4, 5], [6, 7, 8]
     runner.new_request(token_ids=[0] * offloaded_block_size * 3)
-    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(
-        list(keys)[1:2]
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output(list(keys)[1:2])
     )
     runner.run(decoded_tokens=[0])
 
@@ -45,18 +45,22 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
     runner.manager.prepare_store.assert_not_called()
 
     # +1 token -> single block, fail prepare_store
-    runner.manager.prepare_store.side_effect = lambda keys: None
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: None
     runner.run(decoded_tokens=[0])
     runner.manager.prepare_store.assert_called()
 
     # 1 more block (+ token for async scheduling)
     # now set block_hashes_to_store = []
-    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output([])
+    )
     runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))
 
     # 1 more block (+ token for kicking off offloading)
     # now check touch was called with all 6 blocks
-    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output(keys)
+    )
     runner.run(
         decoded_tokens=[0] * (offloaded_block_size + 1),
         expected_stored_gpu_block_indexes=(15, 16, 17),
@@ -89,21 +93,27 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
     runner.new_request(
         token_ids=[0] * gpu_block_size + [1] * (offloaded_block_size - gpu_block_size)
     )
-    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output([])
+    )
     runner.run(decoded_tokens=[EOS_TOKEN_ID])
     runner.manager.lookup.assert_not_called()
 
     # single block lookup with no hits
     runner.new_request(token_ids=[1] * offloaded_block_size)
-    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output([])
+    )
     runner.run(decoded_tokens=[EOS_TOKEN_ID])
     runner.manager.lookup.assert_called()
     assert len(list(runner.manager.lookup.call_args.args[0])) == 1
 
     # single block lookup with a hit
     runner.scheduler.reset_prefix_cache()
     runner.new_request(token_ids=[0] * offloaded_block_size)
-    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output([])
+    )
     runner.manager.lookup.return_value = 1
     runner.run(
         decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2)
@@ -113,7 +123,9 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
     runner.new_request(
         token_ids=[0] * offloaded_block_size * 2 + [1] * offloaded_block_size
     )
-    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output([])
+    )
     runner.manager.lookup.return_value = 1
     runner.run(
         decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5)
@@ -164,14 +176,18 @@ def test_request_preemption(request_runner, async_scheduling: bool):
     # 2 blocks, store all, without flushing
     # blocks = [0, 1, 2], [3, 4, 5]
     runner.new_request(token_ids=[0] * offloaded_block_size * 2)
-    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output(keys)
+    )
     runner.run(
         decoded_tokens=[0],
         complete_transfers=False,
     )
 
     # decode 2 more blocks - 1 gpu block, storing [6, 7, 8] (no flush)
-    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output(keys)
+    )
     runner.run(
         decoded_tokens=[0] * (2 * offloaded_block_size - gpu_block_size),
         complete_transfers=False,
@@ -195,7 +211,9 @@ def test_request_preemption(request_runner, async_scheduling: bool):
     # request should now return from preemption
     # re-load [0, ..., 8] from the CPU and store [9, 10, 11]
     runner.manager.lookup.return_value = 3
-    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output(keys)
+    )
     runner.run(
         decoded_tokens=[0] * gpu_block_size,
         expected_loaded_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
@@ -222,7 +240,9 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
 
     # store 1 blocks
     runner.new_request(token_ids=[0] * offloaded_block_size)
-    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output(keys)
+    )
     runner.run(
         decoded_tokens=[EOS_TOKEN_ID],
         expected_stored_gpu_block_indexes=(0, 1, 2),
@@ -253,7 +273,9 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
     assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
 
     # complete transfers
-    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output([])
+    )
     runner.run(
         decoded_tokens=[EOS_TOKEN_ID],
         expected_loaded_gpu_block_indexes=(0, 1, 2),
@@ -278,7 +300,9 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
 
     # store 1 blocks
     runner.new_request(token_ids=[0] * offloaded_block_size)
-    runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
+    runner.manager.prepare_store.side_effect = (
+        lambda keys, req_context: generate_store_output(keys)
+    )
     runner.run(
         decoded_tokens=[EOS_TOKEN_ID],
         expected_stored_gpu_block_indexes=(0, 1, 2),

tests/v1/kv_connector/unit/offloading_connector/utils.py

@@ -115,7 +115,7 @@ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
 
         self.manager = MagicMock(spec=OffloadingManager)
         self.manager.lookup.return_value = 0
-        self.manager.prepare_load = lambda keys: MockLoadStoreSpec(keys)
+        self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys)
         self.handler = MockOffloadingHandler()
 
     def get_manager(self) -> OffloadingManager: