A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 1, 20:00 UTC
Add objectstore as a secondary tier to multi-tier kv cache offloading (#41968)
Signed-off-by: Effi Ofer <effi.ofer@gmail.com>
Effi Ofer committed
Jun 5, 2026 at 18:05 UTC
6a894574bf7b78fe639cb0c5ddf3eed911ea9011
6 files changed
+659
-2
tests/v1/kv_offload/tiering/test_obj_tier.py
new
+365
@@ -0,0 +1,365 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Mock-based unit tests for ObjectStoreSecondaryTierManager.
5
+
6
+These tests replace the NIXL backend with an in-memory mock so they run
7
+without S3 credentials or a live object store. They verify the manager's
8
+state machine: job submission, transfer completion polling, and lookup.
9
+"""
10
+
11
+import uuid
12
+from collections.abc import Callable
13
+from types import SimpleNamespace
14
+from unittest.mock import MagicMock, patch
15
+
16
+import numpy as np
17
+import torch
18
+
19
+from vllm.v1.kv_offload.base import OffloadKey, ReqContext, make_offload_key
20
+from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult
21
+from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManager
22
+
23
+# ---------------------------------------------------------------------------
24
+# Shared stubs
25
+# ---------------------------------------------------------------------------
26
+
27
+
28
+def _make_vllm_config():
29
+ return SimpleNamespace(
30
+ model_config=SimpleNamespace(model="test/model"),
31
+ cache_config=SimpleNamespace(block_size=16, cache_dtype="float16"),
32
+ parallel_config=SimpleNamespace(
33
+ tensor_parallel_size=1,
34
+ pipeline_parallel_size=1,
35
+ prefill_context_parallel_size=1,
36
+ decode_context_parallel_size=1,
37
+ rank=0,
38
+ ),
39
+ )
40
+
41
+
42
+_OFFLOADING_SPEC = SimpleNamespace(
43
+ vllm_config=_make_vllm_config(),
44
+ kv_cache_config=SimpleNamespace(kv_cache_groups=[]),
45
+)
46
+
47
+_STORE_CONFIG = {
48
+ "bucket": "mock-bucket",
49
+ "endpoint_override": "mock:9000",
50
+ "access_key": "mock-access",
51
+ "secret_key": "mock-secret",
52
+}
53
+
54
+_BLOCK_ELEMENTS = 256
55
+_DTYPE = torch.float32
56
+_RUN_PREFIX = f"test/{uuid.uuid4().hex[:8]}"
57
+_CTX = ReqContext(req_id="test-req")
58
+
59
+
60
+def key(n: int) -> OffloadKey:
61
+ return make_offload_key(n.to_bytes(8, "big"), 0)
62
+
63
+
64
+def make_job(
65
+ job_id: int,
66
+ keys: list[OffloadKey],
67
+ block_ids: list[int] | None = None,
68
+) -> JobMetadata:
69
+ if block_ids is None:
70
+ block_ids = list(range(len(keys)))
71
+ return JobMetadata(
72
+ job_id=job_id,
73
+ keys=keys,
74
+ block_ids=np.array(block_ids, dtype=np.int64),
75
+ is_promotion=False,
76
+ req_context=_CTX,
77
+ )
78
+
79
+
80
+# ---------------------------------------------------------------------------
81
+# Mock NIXL agent
82
+# ---------------------------------------------------------------------------
83
+
84
+
85
+class MockNixlAgent:
86
+ """In-memory NIXL agent. Tracks stored object keys and simulates async
87
+ transfers: transfer() returns PROC, check_xfer_state() returns DONE and
88
+ commits the write to the in-memory key set.
89
+
90
+ The four methods overridden by tests (register_memory, make_prepped_xfer,
91
+ check_xfer_state, query_memory) are stored as Callable instance attributes
92
+ so mypy allows reassignment in tests.
93
+ """
94
+
95
+ # Callable attributes — tests may reassign these on instances.
96
+ register_memory: Callable
97
+ make_prepped_xfer: Callable
98
+ check_xfer_state: Callable
99
+ query_memory: Callable
100
+
101
+ def __init__(self):
102
+ self._stored_obj_keys: set[str] = set()
103
+ # handle_id -> (op, [obj_keys])
104
+ self._pending: dict[int, tuple[str, list[str]]] = {}
105
+ self._handle_counter = 0
106
+ self._last_obj_keys: list[str] = []
107
+ # Bind default implementations as instance attributes.
108
+ self.register_memory = self._register_memory
109
+ self.make_prepped_xfer = self._make_prepped_xfer
110
+ self.check_xfer_state = self._check_xfer_state
111
+ self.query_memory = self._query_memory
112
+
113
+ def create_backend(self, backend_type, params):
114
+ pass
115
+
116
+ def _register_memory(self, descs, mem_type=None, backends=None):
117
+ mock = MagicMock()
118
+ mock.trim.return_value = MagicMock()
119
+ # Capture obj_keys from OBJ 4-tuples: (addr, len, dev_id, obj_key)
120
+ if mem_type == "OBJ" and descs:
121
+ self._last_obj_keys = [d[3] for d in descs if d[3]]
122
+ return mock
123
+
124
+ def deregister_memory(self, desc):
125
+ pass
126
+
127
+ def prep_xfer_dlist(self, agent_name, descs, mem_type=None, backends=None):
128
+ return MagicMock()
129
+
130
+ def _make_prepped_xfer(
131
+ self,
132
+ op,
133
+ local_handle,
134
+ local_indices,
135
+ remote_handle,
136
+ remote_indices,
137
+ notif_msg=b"",
138
+ backends=None,
139
+ skip_desc_merge=False,
140
+ ):
141
+ handle = MagicMock()
142
+ handle._id = self._handle_counter
143
+ self._pending[self._handle_counter] = (op, list(self._last_obj_keys))
144
+ self._handle_counter += 1
145
+ return handle
146
+
147
+ def transfer(self, handle):
148
+ return "PROC"
149
+
150
+ def _check_xfer_state(self, handle):
151
+ entry = self._pending.pop(handle._id, None)
152
+ if entry:
153
+ op, obj_keys = entry
154
+ if op == "WRITE":
155
+ self._stored_obj_keys.update(obj_keys)
156
+ return "DONE"
157
+
158
+ def release_xfer_handle(self, handle):
159
+ pass
160
+
161
+ def release_dlist_handle(self, handle):
162
+ pass
163
+
164
+ def _query_memory(self, queries, mem_type, agent_name):
165
+ return [object() if q[3] in self._stored_obj_keys else None for q in queries]
166
+
167
+
168
+# ---------------------------------------------------------------------------
169
+# Fixture
170
+# ---------------------------------------------------------------------------
171
+
172
+
173
+def _make_tier(
174
+ num_blocks: int = 4,
175
+) -> tuple[ObjectStoreSecondaryTierManager, MockNixlAgent]:
176
+ """Create a tier backed by a fresh MockNixlAgent."""
177
+ mock_agent = MockNixlAgent()
178
+ tensor = torch.zeros((num_blocks, _BLOCK_ELEMENTS), dtype=_DTYPE)
179
+ view = memoryview(tensor.numpy())
180
+ with (
181
+ patch("vllm.v1.kv_offload.tiering.obj.manager.nixl_agent_config"),
182
+ patch(
183
+ "vllm.v1.kv_offload.tiering.obj.manager.nixl_agent",
184
+ return_value=mock_agent,
185
+ ),
186
+ ):
187
+ tier = ObjectStoreSecondaryTierManager(
188
+ offloading_spec=_OFFLOADING_SPEC,
189
+ primary_kv_view=view,
190
+ tier_type="obj",
191
+ store_config=_STORE_CONFIG,
192
+ prefix=_RUN_PREFIX,
193
+ )
194
+ return tier, mock_agent
195
+
196
+
197
+def drain(
198
+ tier: ObjectStoreSecondaryTierManager, max_rounds: int = 20
199
+) -> list[JobResult]:
200
+ """Poll get_finished_jobs() until all in-flight jobs resolve."""
201
+ results: list[JobResult] = []
202
+ for _ in range(max_rounds):
203
+ results.extend(tier.get_finished_jobs())
204
+ if not tier._transfers:
205
+ break
206
+ return results
207
+
208
+
209
+# ---------------------------------------------------------------------------
210
+# Tests
211
+# ---------------------------------------------------------------------------
212
+
213
+
214
+class TestMockObjTierBasic:
215
+ def setup_method(self):
216
+ self.tier, self.agent = _make_tier(num_blocks=4)
217
+
218
+ def test_lookup_empty_tier(self):
219
+ assert self.tier.lookup(key(1), _CTX) is False
220
+
221
+ def test_store_and_lookup(self):
222
+ self.tier.submit_store(make_job(1, [key(1)], [0]))
223
+ results = drain(self.tier)
224
+ assert len(results) == 1
225
+ assert results[0].success
226
+ assert self.tier.lookup(key(1), _CTX) is True
227
+
228
+ def test_lookup_unrelated_key_returns_false(self):
229
+ self.tier.submit_store(make_job(1, [key(1)], [0]))
230
+ drain(self.tier)
231
+ assert self.tier.lookup(key(999), _CTX) is False
232
+
233
+ def test_store_then_load_roundtrip(self):
234
+ self.tier.submit_store(make_job(1, [key(1), key(2)], [0, 1]))
235
+ results = drain(self.tier)
236
+ assert results[0].success
237
+
238
+ self.tier.submit_load(make_job(2, [key(1), key(2)], [0, 1]))
239
+ results = drain(self.tier)
240
+ assert len(results) == 1
241
+ assert results[0].success
242
+
243
+ def test_multiple_jobs_tracked_independently(self):
244
+ self.tier.submit_store(make_job(1, [key(1)], [0]))
245
+ self.tier.submit_store(make_job(2, [key(2)], [1]))
246
+ results = drain(self.tier)
247
+ assert len(results) == 2
248
+ assert all(r.success for r in results)
249
+
250
+ def test_failed_transfer_reported(self):
251
+ self.agent.check_xfer_state = lambda h: "ERR"
252
+ self.tier.submit_store(make_job(1, [key(1)], [0]))
253
+ results = drain(self.tier)
254
+ assert len(results) == 1
255
+ assert not results[0].success
256
+
257
+ def test_pending_transfer_not_returned_until_done(self):
258
+ # First poll returns PROC; second poll returns DONE.
259
+ call_count = [0]
260
+ original = self.agent.check_xfer_state
261
+
262
+ def delayed(h):
263
+ call_count[0] += 1
264
+ return "PROC" if call_count[0] == 1 else original(h)
265
+
266
+ self.agent.check_xfer_state = delayed
267
+
268
+ self.tier.submit_store(make_job(1, [key(1)], [0]))
269
+ assert list(self.tier.get_finished_jobs()) == []
270
+ results = list(self.tier.get_finished_jobs())
271
+ assert len(results) == 1
272
+ assert results[0].success
273
+
274
+
275
+class TestMockObjTierMultiBlock:
276
+ def test_store_multiple_blocks(self):
277
+ tier, _ = _make_tier(num_blocks=8)
278
+ keys = [key(i) for i in range(8)]
279
+ tier.submit_store(make_job(1, keys, list(range(8))))
280
+ results = drain(tier)
281
+ assert len(results) == 1
282
+ assert results[0].success
283
+ assert all(tier.lookup(k, _CTX) for k in keys)
284
+
285
+ def test_partial_block_lookup(self):
286
+ tier, _ = _make_tier(num_blocks=4)
287
+ tier.submit_store(make_job(1, [key(0), key(1)], [0, 1]))
288
+ drain(tier)
289
+ assert tier.lookup(key(0), _CTX) is True
290
+ assert tier.lookup(key(1), _CTX) is True
291
+ assert tier.lookup(key(2), _CTX) is False
292
+
293
+
294
+class TestMockObjTierFailures:
295
+ def test_lookup_exception_returns_false(self):
296
+ tier, agent = _make_tier(num_blocks=4)
297
+ agent.query_memory = lambda *a, **k: (_ for _ in ()).throw(
298
+ RuntimeError("backend error")
299
+ )
300
+ assert tier.lookup(key(1), _CTX) is False
301
+
302
+ def test_submit_store_register_memory_failure_reported_in_get_finished(self):
303
+ tier, agent = _make_tier(num_blocks=4)
304
+ agent.register_memory = lambda *a, **k: None
305
+ tier.submit_store(make_job(1, [key(1)], [0]))
306
+ results = list(tier.get_finished_jobs())
307
+ assert len(results) == 1
308
+ assert results[0].job_id == 1
309
+ assert not results[0].success
310
+
311
+ def test_submit_load_register_memory_failure_reported_in_get_finished(self):
312
+ tier, agent = _make_tier(num_blocks=4)
313
+ agent.register_memory = lambda *a, **k: None
314
+ tier.submit_load(make_job(2, [key(1)], [0]))
315
+ results = list(tier.get_finished_jobs())
316
+ assert len(results) == 1
317
+ assert results[0].job_id == 2
318
+ assert not results[0].success
319
+
320
+ def test_submit_store_make_prepped_xfer_failure_reported_in_get_finished(self):
321
+ tier, agent = _make_tier(num_blocks=4)
322
+ agent.make_prepped_xfer = lambda *a, **k: None
323
+ tier.submit_store(make_job(3, [key(1)], [0]))
324
+ results = list(tier.get_finished_jobs())
325
+ assert len(results) == 1
326
+ assert results[0].job_id == 3
327
+ assert not results[0].success
328
+
329
+ def test_failure_and_success_both_returned_by_get_finished(self):
330
+ # One job fails at submission, another succeeds in flight.
331
+ tier, agent = _make_tier(num_blocks=4)
332
+ original_register = agent.register_memory
333
+ call_count = [0]
334
+
335
+ def register_once_fail(*a, **k):
336
+ call_count[0] += 1
337
+ return None if call_count[0] == 1 else original_register(*a, **k)
338
+
339
+ agent.register_memory = register_once_fail
340
+
341
+ tier.submit_store(make_job(1, [key(1)], [0])) # fails immediately
342
+ tier.submit_store(make_job(2, [key(2)], [1])) # succeeds
343
+ results = drain(tier)
344
+ assert len(results) == 2
345
+ by_id = {r.job_id: r for r in results}
346
+ assert not by_id[1].success
347
+ assert by_id[2].success
348
+
349
+
350
+class TestMockObjTierShutdown:
351
+ def test_shutdown_clears_in_flight_transfers(self):
352
+ tier, agent = _make_tier(num_blocks=4)
353
+ # Keep transfer in flight by never completing it
354
+ agent.check_xfer_state = lambda h: "PROC"
355
+ tier.submit_store(make_job(1, [key(1)], [0]))
356
+ assert len(tier._transfers) == 1
357
+ tier.shutdown()
358
+ assert len(tier._transfers) == 0
359
+ assert tier._dram_prepped_handle is None
360
+ assert tier._primary_reg is None
361
+
362
+ def test_shutdown_idempotent(self):
363
+ tier, _ = _make_tier(num_blocks=4)
364
+ tier.shutdown()
365
+ tier.shutdown() # must not raise
vllm/v1/kv_offload/tiering/factory.py
+6
@@ -63,3 +63,9 @@ SecondaryTierFactory.register_tier(
63
"vllm.v1.kv_offload.tiering.fs.manager",
64
"FileSystemTierManager",
65
)
66
+
67
+SecondaryTierFactory.register_tier(
68
+ "obj",
69
+ "vllm.v1.kv_offload.tiering.obj.manager",
70
+ "ObjectStoreSecondaryTierManager",
71
+)
vllm/v1/kv_offload/tiering/obj/__init__.py
vllm/v1/kv_offload/tiering/obj/config.py
new
+30
@@ -0,0 +1,30 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Connection configuration for the object store secondary tier."""
4
+
5
+from dataclasses import dataclass
6
+
7
+
8
+@dataclass
9
+class ObjStoreConfig:
10
+ """Connection parameters for an object store backend."""
11
+
12
+ bucket: str
13
+ endpoint_override: str
14
+ access_key: str
15
+ secret_key: str
16
+ scheme: str = "http"
17
+ ca_bundle: str = ""
18
+
19
+ def to_nixl_params(self) -> dict[str, str]:
20
+ """Build the NIXL backend params dict."""
21
+ params: dict[str, str] = {
22
+ "bucket": self.bucket,
23
+ "endpoint_override": self.endpoint_override,
24
+ "scheme": self.scheme,
25
+ "access_key": self.access_key,
26
+ "secret_key": self.secret_key,
27
+ }
28
+ if self.ca_bundle:
29
+ params["ca_bundle"] = self.ca_bundle
30
+ return params
vllm/v1/kv_offload/tiering/obj/manager.py
new
+256
@@ -0,0 +1,256 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Object store secondary tier implementation."""
4
+
5
+import ctypes
6
+from collections.abc import Iterable
7
+from typing import TYPE_CHECKING, NamedTuple
8
+
9
+from vllm.distributed.nixl_utils import NixlWrapper as nixl_agent
10
+from vllm.distributed.nixl_utils import nixl_agent_config
11
+from vllm.logger import init_logger
12
+from vllm.v1.kv_offload.base import OffloadKey, ReqContext
13
+from vllm.v1.kv_offload.file_mapper import FileMapper
14
+from vllm.v1.kv_offload.tiering.base import (
15
+ JobMetadata,
16
+ JobResult,
17
+ RequestOffloadingContext,
18
+ SecondaryTierManager,
19
+)
20
+from vllm.v1.kv_offload.tiering.obj.config import ObjStoreConfig
21
+
22
+if TYPE_CHECKING:
23
+ from nixl._api import nixl_prepped_dlist_handle, nixl_xfer_handle
24
+
25
+ from vllm.v1.kv_offload.base import OffloadingSpec
26
+
27
+logger = init_logger(__name__)
28
+
29
+NIXL_WRITE = "WRITE"
30
+NIXL_READ = "READ"
31
+NIXL_PROC = "PROC"
32
+NIXL_DONE = "DONE"
33
+
34
+# Device ID for CPU DRAM descriptors. DRAM is not a multi-device resource so
35
+# the device ID is always 0.
36
+NIXL_DEV_ID: int = 0
37
+
38
+# Fields for NIXL OBJ descriptors: (addr, len, dev_id, obj_key).
39
+# For existence probes addr and len are placeholders — no data is read.
40
+# dev_id=0 is reserved for probes; transfers start from 1.
41
+_PROBE_ADDR: int = 0
42
+_PROBE_LEN: int = 1
43
+_PROBE_DEV_ID: int = 0
44
+
45
+
46
+class TransferEntry(NamedTuple):
47
+ xfer_handle: "nixl_xfer_handle"
48
+ files_desc: object
49
+ obj_handle: "nixl_prepped_dlist_handle"
50
+
51
+
52
+class ObjectStoreSecondaryTierManager(SecondaryTierManager):
53
+ """Secondary tier that offloads KV cache blocks to an S3-compatible store.
54
+
55
+ Handles CPU DRAM <-> S3 transfers only. GPU <-> CPU is managed by the
56
+ primary tier. Object keys are formed as ``{prefix}/{hash_shard}/{hash}.bin``.
57
+ """
58
+
59
+ def __init__(
60
+ self,
61
+ offloading_spec: "OffloadingSpec",
62
+ primary_kv_view: memoryview,
63
+ tier_type: str,
64
+ store_config: dict,
65
+ prefix: str = "",
66
+ io_threads: int = 4,
67
+ ):
68
+ super().__init__(offloading_spec, primary_kv_view, tier_type)
69
+ agent_config = nixl_agent_config(backends=[])
70
+ self._agent = nixl_agent("ObjAgent", agent_config)
71
+ obj_config = ObjStoreConfig(**store_config)
72
+ params = {**obj_config.to_nixl_params(), "num_threads": str(io_threads)}
73
+ self._agent.create_backend("OBJ", params)
74
+ self._transfers: dict[int, TransferEntry] = {}
75
+ self._failed_jobs: list[JobResult] = []
76
+ self._primary_reg = None
77
+ self._block_size_bytes: int = 0
78
+ root_dir = f"{prefix}/" if prefix else ""
79
+ self._file_mapper = FileMapper.from_offloading_spec(root_dir, offloading_spec)
80
+ self._next_obj_dev_id: int = 1 # dev_id=0 is reserved for _exists() probes
81
+
82
+ self._probe_connectivity()
83
+
84
+ base_addr = ctypes.addressof(ctypes.c_char.from_buffer(primary_kv_view))
85
+ assert primary_kv_view.strides is not None
86
+ stride = primary_kv_view.strides[0]
87
+ self._primary_reg = self._agent.register_memory(
88
+ [(base_addr, primary_kv_view.nbytes, NIXL_DEV_ID, "")], "DRAM"
89
+ )
90
+ self._block_size_bytes = stride
91
+ all_blocks = [
92
+ (base_addr + i * stride, stride, NIXL_DEV_ID)
93
+ for i in range(len(primary_kv_view))
94
+ ]
95
+ # NIXL_INIT_AGENT marks this as the local side; make_prepped_xfer requires
96
+ # local_xfer_side tagged with NIXL_INIT_AGENT and remote_xfer_side tagged
97
+ # with the peer agent name ("ObjAgent").
98
+ self._dram_prepped_handle: nixl_prepped_dlist_handle = (
99
+ self._agent.prep_xfer_dlist("NIXL_INIT_AGENT", all_blocks, "DRAM")
100
+ )
101
+
102
+ def _probe_connectivity(self) -> None:
103
+ """Verify object store connectivity at startup via a NIXL lookup probe.
104
+
105
+ Performs a single exists() check against a synthetic key that will
106
+ never exist. A True/False result confirms the bucket is reachable;
107
+ an exception indicates misconfigured obj store params and raises RuntimeError.
108
+ """
109
+ probe_key = "__nixl_probe__/connectivity_test"
110
+ try:
111
+ self._exists(probe_key)
112
+ logger.info("Object store tier connectivity probe succeeded")
113
+ except Exception as e:
114
+ raise RuntimeError(
115
+ f"Object store tier connectivity probe failed — check bucket, "
116
+ f"endpoint_override, access_key, secret_key, and scheme. "
117
+ f"Error: {e}"
118
+ ) from e
119
+
120
+ def _exists(self, obj_key: str) -> bool:
121
+ results = self._agent.query_memory(
122
+ [(_PROBE_ADDR, _PROBE_LEN, _PROBE_DEV_ID, obj_key)], "OBJ", "OBJ"
123
+ )
124
+ return results[0] is not None
125
+
126
+ def _submit_transfer(
127
+ self,
128
+ job_id: int,
129
+ block_ids: Iterable[int],
130
+ obj_keys: Iterable[str],
131
+ op: str,
132
+ ) -> None:
133
+ """Submit an async transfer. op is 'WRITE' (store) or 'READ' (load)."""
134
+ block_ids_list = [int(bid) for bid in block_ids]
135
+ # The OBJ backend maps devId -> obj_key. All descriptors must have
136
+ # unique devIds or later registrations overwrite earlier ones.
137
+ nixl_files = [
138
+ (0, self._block_size_bytes, dev_id, key)
139
+ for dev_id, key in enumerate(obj_keys, self._next_obj_dev_id)
140
+ ]
141
+ self._next_obj_dev_id += len(nixl_files)
142
+
143
+ files_desc = self._agent.register_memory(nixl_files, "OBJ")
144
+ if files_desc is None:
145
+ logger.warning("register_memory (OBJ) failed for job %d", job_id)
146
+ self._failed_jobs.append(JobResult(job_id=job_id, success=False))
147
+ return
148
+
149
+ obj_handle = self._agent.prep_xfer_dlist("ObjAgent", files_desc.trim())
150
+ if not obj_handle:
151
+ logger.warning("prep_xfer_dlist (OBJ) failed for job %d", job_id)
152
+ self._agent.deregister_memory(files_desc)
153
+ self._failed_jobs.append(JobResult(job_id=job_id, success=False))
154
+ return
155
+
156
+ xfer_handle = self._agent.make_prepped_xfer(
157
+ op,
158
+ self._dram_prepped_handle,
159
+ block_ids_list,
160
+ obj_handle,
161
+ list(range(len(nixl_files))),
162
+ )
163
+ if not xfer_handle:
164
+ logger.warning("make_prepped_xfer failed for job %d", job_id)
165
+ self._agent.release_dlist_handle(obj_handle)
166
+ self._agent.deregister_memory(files_desc)
167
+ self._failed_jobs.append(JobResult(job_id=job_id, success=False))
168
+ return
169
+
170
+ state = self._agent.transfer(xfer_handle)
171
+ if state == "ERR":
172
+ logger.warning("agent.transfer failed for job %d", job_id)
173
+ self._agent.release_dlist_handle(obj_handle)
174
+ self._agent.deregister_memory(files_desc)
175
+ self._agent.release_xfer_handle(xfer_handle)
176
+ self._failed_jobs.append(JobResult(job_id=job_id, success=False))
177
+ return
178
+
179
+ self._transfers[job_id] = TransferEntry(xfer_handle, files_desc, obj_handle)
180
+
181
+ def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
182
+ try:
183
+ return self._exists(self._file_mapper.get_file_name(key))
184
+ except Exception as e:
185
+ logger.warning("lookup failed for key %s: %s", key, e)
186
+ return False
187
+
188
+ def submit_store(self, job_metadata: JobMetadata) -> None:
189
+ obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys)
190
+ self._submit_transfer(
191
+ job_metadata.job_id, job_metadata.block_ids, obj_keys, NIXL_WRITE
192
+ )
193
+
194
+ def submit_load(self, job_metadata: JobMetadata) -> None:
195
+ obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys)
196
+ self._submit_transfer(
197
+ job_metadata.job_id, job_metadata.block_ids, obj_keys, NIXL_READ
198
+ )
199
+
200
+ def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
201
+ return RequestOffloadingContext()
202
+
203
+ def get_finished_jobs(self) -> Iterable[JobResult]:
204
+ """Poll in-flight transfers; return completed (job_id, success) pairs."""
205
+ results: list[JobResult] = self._failed_jobs
206
+ self._failed_jobs = []
207
+ for job_id, entry in list(self._transfers.items()):
208
+ try:
209
+ state = self._agent.check_xfer_state(entry.xfer_handle)
210
+ except Exception as exc:
211
+ success = False
212
+ logger.warning("check_xfer_state raised for job %d: %s", job_id, exc)
213
+ else:
214
+ if state == NIXL_PROC:
215
+ continue
216
+ elif state == NIXL_DONE:
217
+ success = True
218
+ else:
219
+ success = False
220
+ logger.warning("transfer failed job=%d state=%s", job_id, state)
221
+ del self._transfers[job_id]
222
+ self._agent.release_xfer_handle(entry.xfer_handle)
223
+ self._agent.release_dlist_handle(entry.obj_handle)
224
+ self._agent.deregister_memory(entry.files_desc)
225
+ results.append(JobResult(job_id=job_id, success=success))
226
+ return results
227
+
228
+ def shutdown(self) -> None:
229
+ for job_id, entry in self._transfers.items():
230
+ try:
231
+ self._agent.release_xfer_handle(entry.xfer_handle)
232
+ except Exception as exc:
233
+ logger.warning("release_xfer_handle failed for job %d: %s", job_id, exc)
234
+ try:
235
+ self._agent.release_dlist_handle(entry.obj_handle)
236
+ except Exception as exc:
237
+ logger.warning(
238
+ "release_dlist_handle failed for job %d: %s", job_id, exc
239
+ )
240
+ try:
241
+ self._agent.deregister_memory(entry.files_desc)
242
+ except Exception as exc:
243
+ logger.warning("deregister_memory failed for job %d: %s", job_id, exc)
244
+ self._transfers.clear()
245
+ if self._dram_prepped_handle is not None:
246
+ try:
247
+ self._agent.release_dlist_handle(self._dram_prepped_handle)
248
+ except Exception as exc:
249
+ logger.warning("failed to release DRAM prepped handle: %s", exc)
250
+ self._dram_prepped_handle = None
251
+ if self._primary_reg is not None:
252
+ try:
253
+ self._agent.deregister_memory(self._primary_reg)
254
+ except Exception as exc:
255
+ logger.warning("failed to deregister primary buffer: %s", exc)
256
+ self._primary_reg = None
vllm/v1/kv_offload/tiering/spec.py
+2
-2
@@ -131,8 +131,8 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
131
)
132
except Exception as e:
133
logger.error(
134
- "Failed to create secondary tier from config %s: %s",
135
- tier_config,
134
+ "Failed to create secondary tier from config index %i: %s",
135
+ i,
136
e,
137
)
138
raise