@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 2d814a0082

[kv_offload] Emit tier-owned BlockStored events from FS/OBJ secondary tiers (#47923)

Signed-off-by: Change72 <changg@nvidia.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>

Chang Guo committed Jul 9, 2026 at 20:17 UTC 2d814a00820daec7082599bea75ae1d0959a346c
6 files changed +433 -9
docs/features/kv_offloading_usage.md
+28
@@ -81,6 +81,8 @@ vllm serve <model> \
81
82 Each entry in `secondary_tiers` is a dict with a required `type` field plus tier-specific fields.
83
84 +The filesystem and object-store tiers can publish hash-only `BlockStored` KV events for blocks they successfully store, tagged with a stable per-tier `medium` (`FS` for the filesystem tier, `OBJ` for the object-store tier). Set `enable_kv_events: true` in the tier's entry to opt in; events are published only when KV cache events are also enabled globally via `--kv-events-config`.
85 +
86 ### Filesystem (FS)
87
88 The filesystem tier (`type: "fs"`) writes blocks to a directory on local storage.
@@ -91,6 +93,7 @@ The filesystem tier (`type: "fs"`) writes blocks to a directory on local storage
93 | `root_dir` | yes | — | Base directory; vLLM creates subdirectories beneath it (see [On-Disk Layout](#on-disk-layout)). |
94 | `n_read_threads` | no | `16` | Read-priority I/O threads (load path). |
95 | `n_write_threads` | no | `16` | Write-priority I/O threads (store path). |
96 +| `enable_kv_events` | no | `false` | Publish `BlockStored` KV events (medium `FS`) for successfully stored blocks. Requires KV cache events to be enabled globally. |
97
98 Each thread group prefers its own queue but pulls from the other when its primary queue is empty, so a write-heavy or read-heavy burst won't leave the off-priority queue waiting. Size the totals to your storage's effective concurrency.
99
@@ -120,6 +123,31 @@ To enable KV cache sharing between multiple vLLM instances using the same `root_
123 PYTHONHASHSEED=0 vllm serve ...
124 ```
125
126 +### Object Store (OBJ)
127 +
128 +The object-store tier (`type: "obj"`) offloads blocks to an S3-compatible object store through the NIXL OBJ backend.
129 +
130 +| Key | Required | Default | Notes |
131 +| --- | --- | --- | --- |
132 +| `type` | yes | — | Must be `obj`. |
133 +| `store_config` | yes | — | Object store connection parameters (see below). |
134 +| `prefix` | no | `""` | Key prefix prepended to all object keys. |
135 +| `io_threads` | no | `4` | Number of NIXL OBJ backend I/O threads. |
136 +| `enable_kv_events` | no | `false` | Publish `BlockStored` KV events (medium `OBJ`) for successfully stored blocks. Requires KV cache events to be enabled globally. |
137 +
138 +`store_config` fields:
139 +
140 +| Key | Required | Default | Notes |
141 +| --- | --- | --- | --- |
142 +| `bucket` | yes | — | Bucket name. |
143 +| `endpoint_override` | yes | — | Object store endpoint host; the URL scheme is set separately via `scheme`. |
144 +| `scheme` | no | `http` | `http` or `https`. |
145 +| `access_key`, `secret_key`, `session_token` | no | `""` | Explicit credentials. When left empty, the NIXL OBJ plugin falls back to the AWS SDK default credential provider chain (IAM roles, environment variables, credential files), which enables workload-identity auth on Kubernetes. |
146 +| `region` | no | `""` | Bucket region, if the endpoint requires one. |
147 +| `ca_bundle` | no | `""` | CA bundle path for TLS verification. |
148 +
149 +Object keys follow the same run-configuration digest scheme as the filesystem tier (see [On-Disk Layout](#on-disk-layout)) and are stored under the optional `prefix`. The [Cross-Process Sharing](#cross-process-sharing) requirement (`PYTHONHASHSEED`) applies to shared buckets as well, so instances sharing a bucket produce identical keys for identical content. At startup the tier probes object store connectivity and fails fast with a configuration error if the bucket is unreachable.
150 +
151 ### P2P (Including P/D)
152
153 The P2P tier (`type: "p2p"`) shares completed KV blocks between vLLM instances over RDMA via NIXL. Each instance binds a control socket on `host:port` and exchanges blocks directly with peers — no shared filesystem required.
tests/v1/kv_offload/tiering/test_fs_tier.py
+196
@@ -18,8 +18,10 @@ import numpy as np
18 import pytest
19 import torch
20
21 +from vllm.distributed.kv_events import MEDIUM_FS
22 from vllm.v1.kv_offload.base import (
23 LookupResult,
24 + OffloadingEvent,
25 OffloadKey,
26 ReqContext,
27 ScheduleEndContext,
@@ -58,6 +60,16 @@ _MOCK_OFFLOADING_SPEC.kv_cache_config = _MOCK_KV_CACHE_CONFIG
60 _MOCK_OFFLOADING_SPEC.block_size_factor = 1
61
62
63 +def _make_offloading_spec(enable_kv_cache_events: bool) -> MagicMock:
64 + """Mock spec with an explicit global KV events flag."""
65 + spec = MagicMock()
66 + spec.vllm_config = _MOCK_VLLM_CONFIG
67 + spec.kv_cache_config = _MOCK_KV_CACHE_CONFIG
68 + spec.block_size_factor = 1
69 + spec.kv_events_config.enable_kv_cache_events = enable_kv_cache_events
70 + return spec
71 +
72 +
73 def key(n: int) -> OffloadKey:
74 return make_offload_key(n.to_bytes(8, "big"), 0)
75
@@ -164,6 +176,23 @@ def fs_tier(tmp_path):
176 tier.shutdown()
177
178
179 +@pytest.fixture
180 +def fs_tier_with_events(tmp_path):
181 + tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS)
182 + mock_view = memoryview(tensor.numpy())
183 + tier = FileSystemTierManager(
184 + offloading_spec=_make_offloading_spec(enable_kv_cache_events=True),
185 + primary_kv_view=mock_view,
186 + tier_type="fs",
187 + root_dir=str(tmp_path),
188 + n_read_threads=4,
189 + n_write_threads=4,
190 + enable_kv_events=True,
191 + )
192 + yield tier
193 + tier.shutdown()
194 +
195 +
196 # ---------------------------------------------------------------------------
197 # Tests
198 # ---------------------------------------------------------------------------
@@ -396,3 +425,170 @@ def test_batch_lookup_dispatch(fs_tier, monkeypatch, use_c_ext):
425
426 results = lookup_and_wait(tier, [key(1), key(2)])
427 assert results == [LookupResult.HIT, LookupResult.MISS]
428 +
429 +
430 +# ---------------------------------------------------------------------------
431 +# KV events
432 +# ---------------------------------------------------------------------------
433 +
434 +
435 +def test_successful_store_emits_stored_event(fs_tier_with_events):
436 + """A completed store job emits one stored event with the job's keys."""
437 + tier = fs_tier_with_events
438 + keys = [key(1), key(2)]
439 + tier.submit_store(make_job(1, keys, [0, 1]))
440 + assert all(r.success for r in drain(tier))
441 +
442 + events = list(tier.take_events())
443 + assert len(events) == 1
444 + assert events[0].keys == keys
445 + # Literal medium pins the wire contract, not just the constant choice.
446 + assert events[0].medium == "FS"
447 + assert not events[0].removed
448 + # take_events drains the buffer.
449 + assert list(tier.take_events()) == []
450 +
451 +
452 +def test_load_job_emits_no_event(fs_tier_with_events):
453 + tier = fs_tier_with_events
454 + tier.submit_store(make_job(1, [key(1)], [0]))
455 + results = drain(tier)
456 + assert len(results) == 1
457 + assert results[0].success
458 + list(tier.take_events())
459 +
460 + tier.submit_load(make_job(2, [key(1)], [1], is_promotion=True))
461 + results = drain(tier)
462 + assert len(results) == 1
463 + assert results[0].success
464 + assert list(tier.take_events()) == []
465 +
466 +
467 +def test_mixed_job_results_emit_event_only_for_successful_job(
468 + fs_tier_with_events, monkeypatch
469 +):
470 + """With a failed and a successful store job in flight, exactly one event
471 + is emitted and its keys belong to the successful job."""
472 + import vllm.v1.kv_offload.tiering.fs.manager as mgr_mod
473 +
474 + tier = fs_tier_with_events
475 + failing_path = tier.file_mapper.get_file_name(key(1))
476 + original_store_block = mgr_mod.store_block
477 +
478 + def flaky_store_block(dest_path, *args, **kwargs):
479 + if dest_path == failing_path:
480 + raise OSError("injected store failure")
481 + return original_store_block(dest_path, *args, **kwargs)
482 +
483 + monkeypatch.setattr(mgr_mod, "store_block", flaky_store_block)
484 +
485 + tier.submit_store(make_job(1, [key(1)], [0]))
486 + tier.submit_store(make_job(2, [key(2)], [1]))
487 + results = drain(tier)
488 + assert len(results) == 2
489 + by_id = {r.job_id: r for r in results}
490 + assert not by_id[1].success
491 + assert by_id[2].success
492 +
493 + events = list(tier.take_events())
494 + assert len(events) == 1
495 + assert events[0].keys == [key(2)]
496 +
497 +
498 +def test_partially_failed_store_emits_no_event(fs_tier_with_events, monkeypatch):
499 + """A store job with any failed block emits no event for the whole job."""
500 + import vllm.v1.kv_offload.tiering.fs.manager as mgr_mod
501 +
502 + tier = fs_tier_with_events
503 + failing_path = tier.file_mapper.get_file_name(key(2))
504 + original_store_block = mgr_mod.store_block
505 +
506 + def flaky_store_block(dest_path, *args, **kwargs):
507 + if dest_path == failing_path:
508 + raise OSError("injected store failure")
509 + return original_store_block(dest_path, *args, **kwargs)
510 +
511 + monkeypatch.setattr(mgr_mod, "store_block", flaky_store_block)
512 +
513 + tier.submit_store(make_job(1, [key(1), key(2)], [0, 1]))
514 + results = drain(tier)
515 + assert len(results) == 1
516 + assert not results[0].success
517 + assert list(tier.take_events()) == []
518 + assert tier._store_job_keys == {}
519 +
520 +
521 +def test_events_disabled_by_default(fs_tier):
522 + tier, _ = fs_tier
523 + tier.submit_store(make_job(1, [key(1)], [0]))
524 + results = drain(tier)
525 + assert len(results) == 1
526 + assert results[0].success
527 + assert tier.events is None
528 + assert tier._store_job_keys == {}
529 + assert list(tier.take_events()) == []
530 +
531 +
532 +def test_events_require_global_kv_events_flag(tmp_path):
533 + """Tier-level opt-in alone is not enough; the global flag gates events."""
534 + tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS)
535 + tier = FileSystemTierManager(
536 + offloading_spec=_make_offloading_spec(enable_kv_cache_events=False),
537 + primary_kv_view=memoryview(tensor.numpy()),
538 + tier_type="fs",
539 + root_dir=str(tmp_path),
540 + enable_kv_events=True,
541 + )
542 + try:
543 + assert tier.events is None
544 + tier.submit_store(make_job(1, [key(1)], [0]))
545 + results = drain(tier)
546 + assert len(results) == 1
547 + assert results[0].success
548 + assert list(tier.take_events()) == []
549 + assert tier._store_job_keys == {}
550 + finally:
551 + tier.shutdown()
552 +
553 +
554 +def test_cascade_store_emits_fs_event_through_tiering_manager(tmp_path):
555 + """A GPU->CPU->fs cascade surfaces the tier-owned FS stored event via the
556 + TieringOffloadingManager's aggregated take_events()."""
557 + from vllm.v1.kv_offload.tiering.manager import (
558 + CPUPrimaryTierOffloadingManager,
559 + TieringOffloadingManager,
560 + )
561 +
562 + tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS)
563 + view = memoryview(tensor.numpy())
564 + mock_region = MagicMock()
565 + mock_region.create_kv_memoryview.return_value = view
566 + primary = CPUPrimaryTierOffloadingManager(num_blocks=4, mmap_region=mock_region)
567 + tier = FileSystemTierManager(
568 + offloading_spec=_make_offloading_spec(enable_kv_cache_events=True),
569 + primary_kv_view=primary.get_kv_memoryview(),
570 + tier_type="fs",
571 + root_dir=str(tmp_path),
572 + enable_kv_events=True,
573 + )
574 + manager = TieringOffloadingManager(primary_tier=primary, secondary_tiers=[tier])
575 + try:
576 + keys = [key(1), key(2)]
577 + manager.on_new_request(_CTX)
578 + assert manager.prepare_store(keys, _CTX) is not None
579 + manager.complete_store(keys, _CTX) # cascades to the fs tier
580 +
581 + events: list[OffloadingEvent] = []
582 + ctx = ScheduleEndContext(new_req_ids=[], preempted_req_ids=())
583 + deadline = time.monotonic() + 5.0
584 + while time.monotonic() < deadline and not events:
585 + manager.on_schedule_end(ctx)
586 + events.extend(manager.take_events())
587 + time.sleep(0.01)
588 +
589 + fs_events = [e for e in events if e.medium == MEDIUM_FS]
590 + assert len(fs_events) == 1
591 + assert set(fs_events[0].keys) == set(keys)
592 + assert not fs_events[0].removed
593 + finally:
594 + tier.shutdown()
tests/v1/kv_offload/tiering/test_obj_tier.py
+107 -1
@@ -179,8 +179,19 @@ class MockNixlAgent:
179 # ---------------------------------------------------------------------------
180
181
182 +def _make_events_spec(enable_kv_cache_events: bool) -> SimpleNamespace:
183 + """Offloading spec stub with an explicit global KV events flag."""
184 + return SimpleNamespace(
185 + vllm_config=_make_vllm_config(),
186 + kv_cache_config=SimpleNamespace(kv_cache_groups=[]),
187 + kv_events_config=SimpleNamespace(enable_kv_cache_events=enable_kv_cache_events),
188 + )
189 +
190 +
191 def _make_tier(
192 num_blocks: int = 4,
193 + offloading_spec: SimpleNamespace = _OFFLOADING_SPEC,
194 + **tier_kwargs,
195 ) -> tuple[ObjectStoreSecondaryTierManager, MockNixlAgent]:
196 """Create a tier backed by a fresh MockNixlAgent."""
197 mock_agent = MockNixlAgent()
@@ -194,11 +205,12 @@ def _make_tier(
205 ),
206 ):
207 tier = ObjectStoreSecondaryTierManager(
197 - offloading_spec=_OFFLOADING_SPEC,
208 + offloading_spec=offloading_spec,
209 primary_kv_view=view,
210 tier_type="obj",
211 store_config=_STORE_CONFIG,
212 prefix=_RUN_PREFIX,
213 + **tier_kwargs,
214 )
215 return tier, mock_agent
216
@@ -421,6 +433,100 @@ class TestMockObjTierShutdown:
433 tier.shutdown() # must not raise
434
435
436 +class TestObjTierKVEvents:
437 + def setup_method(self):
438 + self.tier, self.agent = _make_tier(
439 + offloading_spec=_make_events_spec(enable_kv_cache_events=True),
440 + enable_kv_events=True,
441 + )
442 +
443 + def test_successful_store_emits_stored_event(self):
444 + """A completed store transfer emits one stored event with the job's keys."""
445 + keys = [key(1), key(2)]
446 + self.tier.submit_store(make_job(1, keys, [0, 1]))
447 + assert all(r.success for r in drain(self.tier))
448 +
449 + events = list(self.tier.take_events())
450 + assert len(events) == 1
451 + assert events[0].keys == keys
452 + # Literal medium pins the wire contract, not just the constant choice.
453 + assert events[0].medium == "OBJ"
454 + assert not events[0].removed
455 + # take_events drains the buffer.
456 + assert list(self.tier.take_events()) == []
457 +
458 + def test_mixed_job_results_emit_event_only_for_successful_job(self):
459 + """With a failed and a successful store job resolving in the same
460 + poll, exactly one event is emitted and its keys belong to the
461 + successful job."""
462 + original = self.agent.check_xfer_state
463 + self.agent.check_xfer_state = lambda h: "ERR" if h._id == 0 else original(h)
464 + self.tier.submit_store(make_job(1, [key(1)], [0])) # handle 0: fails
465 + self.tier.submit_store(make_job(2, [key(2)], [1])) # handle 1: succeeds
466 + results = drain(self.tier)
467 + by_id = {r.job_id: r for r in results}
468 + assert not by_id[1].success
469 + assert by_id[2].success
470 +
471 + events = list(self.tier.take_events())
472 + assert len(events) == 1
473 + assert events[0].keys == [key(2)]
474 + assert self.tier._store_job_keys == {}
475 +
476 + def test_load_job_emits_no_event(self):
477 + self.tier.submit_store(make_job(1, [key(1)], [0]))
478 + results = drain(self.tier)
479 + assert len(results) == 1
480 + assert results[0].success
481 + list(self.tier.take_events())
482 +
483 + self.tier.submit_load(make_job(2, [key(1)], [0]))
484 + results = drain(self.tier)
485 + assert len(results) == 1
486 + assert results[0].success
487 + assert list(self.tier.take_events()) == []
488 +
489 + def test_failed_transfer_emits_no_event(self):
490 + self.agent.check_xfer_state = lambda h: "ERR"
491 + self.tier.submit_store(make_job(1, [key(1)], [0]))
492 + results = drain(self.tier)
493 + assert not results[0].success
494 + assert list(self.tier.take_events()) == []
495 + assert self.tier._store_job_keys == {}
496 +
497 + def test_submission_failure_emits_no_event(self):
498 + self.agent.make_prepped_xfer = lambda *a, **k: None
499 + self.tier.submit_store(make_job(1, [key(1)], [0]))
500 + results = list(self.tier.get_finished_jobs())
501 + assert not results[0].success
502 + assert list(self.tier.take_events()) == []
503 + assert self.tier._store_job_keys == {}
504 +
505 + def test_events_disabled_by_default(self):
506 + tier, _ = _make_tier()
507 + tier.submit_store(make_job(1, [key(1)], [0]))
508 + results = drain(tier)
509 + assert len(results) == 1
510 + assert results[0].success
511 + assert tier.events is None
512 + assert tier._store_job_keys == {}
513 + assert list(tier.take_events()) == []
514 +
515 + def test_events_require_global_kv_events_flag(self):
516 + """Tier-level opt-in alone is not enough; the global flag gates events."""
517 + tier, _ = _make_tier(
518 + offloading_spec=_make_events_spec(enable_kv_cache_events=False),
519 + enable_kv_events=True,
520 + )
521 + tier.submit_store(make_job(1, [key(1)], [0]))
522 + results = drain(tier)
523 + assert len(results) == 1
524 + assert results[0].success
525 + assert tier.events is None
526 + assert tier._store_job_keys == {}
527 + assert list(tier.take_events()) == []
528 +
529 +
530 class TestObjStoreConfig:
531 def test_explicit_credentials_included(self):
532 cfg = ObjStoreConfig(
vllm/distributed/kv_events.py
+2
@@ -44,6 +44,8 @@ class KVCacheEvent(
44
45 MEDIUM_GPU = "GPU"
46 MEDIUM_CPU = "CPU"
47 +MEDIUM_FS = "FS"
48 +MEDIUM_OBJ = "OBJ"
49
50
51 class BlockStored(KVCacheEvent):
vllm/v1/kv_offload/tiering/fs/manager.py
+47 -6
@@ -19,7 +19,7 @@ import functools
19 import json
20 import os
21 from collections.abc import Iterable
22 -from typing import TYPE_CHECKING
22 +from typing import TYPE_CHECKING, ClassVar
23
24 try:
25 from vllm.fs_io_C import batch_lookup as batch_lookup_C
@@ -30,11 +30,18 @@ except ImportError:
30
31 from typing_extensions import override
32
33 +from vllm.distributed.kv_events import MEDIUM_FS
34 from vllm.logger import init_logger
34 -from vllm.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext
35 +from vllm.v1.kv_offload.base import (
36 + LookupResult,
37 + OffloadingEvent,
38 + OffloadKey,
39 + ReqContext,
40 +)
41 from vllm.v1.kv_offload.file_mapper import FileMapper
42 from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager
43 from vllm.v1.kv_offload.tiering.base import (
44 + JobId,
45 JobMetadata,
46 JobResult,
47 RequestOffloadingContext,
@@ -92,6 +99,8 @@ class FileSystemTierManager(SecondaryTierManager):
99 content.
100 """
101
102 + medium: ClassVar[str] = MEDIUM_FS
103 +
104 def __init__(
105 self,
106 offloading_spec: "OffloadingSpec",
@@ -100,6 +109,7 @@ class FileSystemTierManager(SecondaryTierManager):
109 root_dir: str,
110 n_read_threads: int = 16,
111 n_write_threads: int = 16,
112 + enable_kv_events: bool = False,
113 ):
114 """
115 Args:
@@ -110,9 +120,26 @@ class FileSystemTierManager(SecondaryTierManager):
120 root_dir: Root directory for block files.
121 n_read_threads: Number of read-priority I/O threads.
122 n_write_threads: Number of write-priority I/O threads.
123 + enable_kv_events: Emit BlockStored KV events for blocks
124 + successfully stored to this tier. Effective only when KV
125 + cache events are enabled globally (kv_events_config).
126 """
127 super().__init__(offloading_spec, primary_kv_view, tier_type)
128
129 + self.events: list[OffloadingEvent] | None = None
130 + if enable_kv_events:
131 + if offloading_spec.kv_events_config.enable_kv_cache_events:
132 + self.events = []
133 + else:
134 + logger.warning(
135 + "enable_kv_events is set on secondary tier '%s' but KV "
136 + "cache events are disabled globally; the tier will not "
137 + "emit events.",
138 + tier_type,
139 + )
140 + # Keys of in-flight store jobs, tracked only when events are enabled.
141 + self._store_job_keys: dict[JobId, list[OffloadKey]] = {}
142 +
143 # Extract block size from primary view
144 assert primary_kv_view.strides is not None, (
145 "primary_kv_view.strides cannot be None"
@@ -157,6 +184,8 @@ class FileSystemTierManager(SecondaryTierManager):
184
185 @override
186 def submit_store(self, job_metadata: JobMetadata) -> None:
187 + if self.events is not None:
188 + self._store_job_keys[job_metadata.job_id] = list(job_metadata.keys)
189 tasks = (
190 functools.partial(
191 store_block,
@@ -188,10 +217,22 @@ class FileSystemTierManager(SecondaryTierManager):
217 """
218 Collect completed jobs from the finished-jobs queue.
219 """
191 - return (
192 - JobResult(job_id=job_id, success=success)
193 - for job_id, success in self._pool.get_finished()
194 - )
220 + results = []
221 + for job_id, success in self._pool.get_finished():
222 + if self.events is not None:
223 + keys = self._store_job_keys.pop(job_id, None)
224 + if success and keys:
225 + self.events.append(
226 + OffloadingEvent(keys=keys, medium=self.medium, removed=False)
227 + )
228 + results.append(JobResult(job_id=job_id, success=success))
229 + return results
230 +
231 + @override
232 + def take_events(self) -> Iterable[OffloadingEvent]:
233 + if self.events is not None:
234 + yield from self.events
235 + self.events.clear()
236
237 @override
238 def drain_jobs(self) -> None:
vllm/v1/kv_offload/tiering/obj/manager.py
+53 -2
@@ -5,15 +5,22 @@
5 import ctypes
6 import time
7 from collections.abc import Iterable
8 -from typing import TYPE_CHECKING, NamedTuple
8 +from typing import TYPE_CHECKING, ClassVar, NamedTuple
9
10 +from vllm.distributed.kv_events import MEDIUM_OBJ
11 from vllm.distributed.nixl_utils import NixlWrapper as nixl_agent
12 from vllm.distributed.nixl_utils import nixl_agent_config
13 from vllm.logger import init_logger
13 -from vllm.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext
14 +from vllm.v1.kv_offload.base import (
15 + LookupResult,
16 + OffloadingEvent,
17 + OffloadKey,
18 + ReqContext,
19 +)
20 from vllm.v1.kv_offload.file_mapper import FileMapper
21 from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager
22 from vllm.v1.kv_offload.tiering.base import (
23 + JobId,
24 JobMetadata,
25 JobResult,
26 RequestOffloadingContext,
@@ -90,6 +97,8 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
97 primary tier. Object keys are formed as ``{prefix}/{hash_shard}/{hash}.bin``.
98 """
99
100 + medium: ClassVar[str] = MEDIUM_OBJ
101 +
102 def __init__(
103 self,
104 offloading_spec: "OffloadingSpec",
@@ -98,8 +107,36 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
107 store_config: dict,
108 prefix: str = "",
109 io_threads: int = 4,
110 + enable_kv_events: bool = False,
111 ):
112 + """
113 + Args:
114 + offloading_spec: Offloading configuration.
115 + primary_kv_view: Memoryview of the primary tier's CPU KV cache.
116 + tier_type: Tier type identifier, set by SecondaryTierFactory.
117 + store_config: Object store connection parameters (see ObjStoreConfig).
118 + prefix: Key prefix prepended to all object keys.
119 + io_threads: Number of NIXL I/O threads.
120 + enable_kv_events: Emit BlockStored KV events for blocks
121 + successfully stored to this tier. Effective only when KV
122 + cache events are enabled globally (kv_events_config).
123 + """
124 super().__init__(offloading_spec, primary_kv_view, tier_type)
125 +
126 + self.events: list[OffloadingEvent] | None = None
127 + if enable_kv_events:
128 + if offloading_spec.kv_events_config.enable_kv_cache_events:
129 + self.events = []
130 + else:
131 + logger.warning(
132 + "enable_kv_events is set on secondary tier '%s' but KV "
133 + "cache events are disabled globally; the tier will not "
134 + "emit events.",
135 + tier_type,
136 + )
137 + # Keys of in-flight store jobs, tracked only when events are enabled.
138 + self._store_job_keys: dict[JobId, list[OffloadKey]] = {}
139 +
140 agent_config = nixl_agent_config(backends=[])
141 self._agent = nixl_agent("ObjAgent", agent_config)
142 obj_config = ObjStoreConfig(**store_config)
@@ -231,6 +268,8 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
268 return LookupResult.HIT if result else LookupResult.MISS
269
270 def submit_store(self, job_metadata: JobMetadata) -> None:
271 + if self.events is not None:
272 + self._store_job_keys[job_metadata.job_id] = list(job_metadata.keys)
273 obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys)
274 self._submit_transfer(
275 job_metadata.job_id, job_metadata.block_ids, obj_keys, NIXL_WRITE
@@ -279,8 +318,20 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
318 self._poll_active_transfers()
319 results = self._pending_results
320 self._pending_results = []
321 + if self.events is not None:
322 + for result in results:
323 + keys = self._store_job_keys.pop(result.job_id, None)
324 + if result.success and keys:
325 + self.events.append(
326 + OffloadingEvent(keys=keys, medium=self.medium, removed=False)
327 + )
328 return results
329
330 + def take_events(self) -> Iterable[OffloadingEvent]:
331 + if self.events is not None:
332 + yield from self.events
333 + self.events.clear()
334 +
335 def drain_jobs(self) -> None:
336 """Block until every submitted transfer has completed or failed.
337