From 0255eeaabf66661958ae72b1e85eebcdb39de2a5 Mon Sep 17 00:00:00 2001 From: srinivas_oo7 Date: Wed, 15 Jul 2026 21:21:24 -0500 Subject: [PATCH 01/13] Add tiering offloading metrics Signed-off-by: srinivas_oo7 --- .../unit/offloading_connector/test_metrics.py | 38 +++++- .../tiering/test_tiering_offloading.py | 128 +++++++++++++++++- vllm/v1/kv_offload/tiering/base.py | 12 +- vllm/v1/kv_offload/tiering/example/manager.py | 20 ++- vllm/v1/kv_offload/tiering/fs/manager.py | 32 ++++- vllm/v1/kv_offload/tiering/manager.py | 60 +++++++- vllm/v1/kv_offload/tiering/obj/manager.py | 19 ++- vllm/v1/kv_offload/tiering/p2p/manager.py | 19 +-- vllm/v1/kv_offload/tiering/spec.py | 41 ++++++ 9 files changed, 340 insertions(+), 29 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py b/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py index fb6b785c886d..58bbd4091560 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py @@ -25,6 +25,7 @@ OffloadingHistogramMetadata, ) from vllm.v1.kv_offload.factory import OffloadingSpecFactory +from vllm.v1.kv_offload.tiering.base import TieringOffloadingMetrics LOAD_BYTES = _TransferMetricName.LOAD_BYTES LOAD_TIME = _TransferMetricName.LOAD_TIME @@ -101,10 +102,15 @@ def set(self, value): class _FakeVllmConfig: - def __init__(self, store_threshold: int = 2): - self.kv_transfer_config = SimpleNamespace( - kv_connector_extra_config={"store_threshold": store_threshold} - ) + def __init__( + self, + store_threshold: int = 2, + extra_config: dict[str, Any] | None = None, + ): + config = {"store_threshold": store_threshold} + if extra_config is not None: + config.update(extra_config) + self.kv_transfer_config = SimpleNamespace(kv_connector_extra_config=config) def _spec_cls_with_metric_definitions( @@ -679,6 +685,30 @@ def test_prom_metrics_uses_configured_manager_metrics(): assert STORES_SKIPPED not in prom_metrics._offloading_metric_metadata +def test_prom_metrics_registers_tiering_metrics_from_spec(): + prom_metrics = OffloadPromMetrics( + vllm_config=_FakeVllmConfig( + store_threshold=0, + extra_config={ + "spec_name": "TieringOffloadingSpec", + "secondary_tiers": [], + }, + ), # type: ignore[arg-type] + metric_types={ + Gauge: _FakeMetric, + Counter: _FakeMetric, + Histogram: _FakeMetric, + }, + labelnames=["model_name", "engine"], + per_engine_labelvalues={0: ["model", "0"]}, + ) + + metric = prom_metrics._offloading_metric_defs[ + TieringOffloadingMetrics.BLOCK_QUERIES + ] + assert metric.kwargs["labelnames"] == ["model_name", "engine", "tier"] + + def test_aggregate_into_empty_stats(): """Aggregating non-empty stats into a fresh (empty) stats object works.""" empty = OffloadingConnectorStats() diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 2e9fa224d241..4d354d398cca 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -102,18 +102,26 @@ def build_metric_definitions(cls, extra_config): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.stats: OffloadingConnectorStats | None = None + self.lookup_result = LookupResult.MISS + self.finished_jobs: list[JobResult] = [] + self.submitted_loads: list[JobMetadata] = [] + self.submitted_stores: list[JobMetadata] = [] - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: - return False + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: + return self.lookup_result def submit_store(self, job_metadata: JobMetadata) -> None: + self.submitted_stores.append(job_metadata) return def submit_load(self, job_metadata: JobMetadata) -> None: + self.submitted_loads.append(job_metadata) return def get_finished_jobs(self) -> Iterable[JobResult]: - return () + results = self.finished_jobs + self.finished_jobs = [] + return results def drain_jobs(self) -> None: return @@ -141,6 +149,10 @@ def test_tiering_spec_collects_secondary_metric_definitions(monkeypatch): metadata = metrics[MetricsSecondaryTierManager.MY_TIER_METRIC] assert metadata.documentation == "Number of bytes served by the test tier." assert metadata.labelnames == ("tier",) + assert metrics[TieringOffloadingMetrics.BLOCK_QUERIES].labelnames == ("tier",) + assert metrics[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES].labelnames == ( + "tier", + ) def test_tiering_manager_aggregates_secondary_stats(): @@ -181,6 +193,116 @@ def test_tiering_manager_aggregates_secondary_stats(): assert MetricsSecondaryTierManager.MY_TIER_METRIC not in second_stats.data["data"] +def test_tiering_manager_records_secondary_lookup_metrics(): + mock_region = _mock_mmap_region(5) + primary_tier = CPUPrimaryTierOffloadingManager( + num_blocks=5, mmap_region=mock_region + ) + tier0 = MetricsSecondaryTierManager( + offloading_spec=_MOCK_OFFLOADING_SPEC, + primary_kv_view=mock_region.create_kv_memoryview(), + tier_type="fs", + ) + tier1 = MetricsSecondaryTierManager( + offloading_spec=_MOCK_OFFLOADING_SPEC, + primary_kv_view=mock_region.create_kv_memoryview(), + tier_type="p2p", + ) + tier1.lookup_result = LookupResult.HIT + manager = TieringOffloadingManager( + primary_tier=primary_tier, + secondary_tiers=[tier0, tier1], + ) + + assert manager.lookup(to_keys([1])[0], _CTX) is LookupResult.RETRY + + stats = manager.get_stats() + assert stats is not None + values = stats.data["data"] + assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("0:fs",)] == 1 + assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("1:p2p",)] == 1 + assert values[TieringOffloadingMetrics.BLOCK_HITS][("1:p2p",)] == 1 + assert ("0:fs",) not in values[TieringOffloadingMetrics.BLOCK_HITS] + + +def test_tiering_manager_records_finished_job_metrics(): + mock_region = _mock_mmap_region(5) + primary_tier = CPUPrimaryTierOffloadingManager( + num_blocks=5, mmap_region=mock_region + ) + tier = MetricsSecondaryTierManager( + offloading_spec=_MOCK_OFFLOADING_SPEC, + primary_kv_view=mock_region.create_kv_memoryview(), + tier_type="fs", + ) + manager = TieringOffloadingManager( + primary_tier=primary_tier, + secondary_tiers=[tier], + ) + manager.on_new_request(_CTX) + cascade_key, promotion_key, failed_cascade_key, failed_promotion_key = to_keys( + range(4) + ) + + prepared = primary_tier.prepare_store([cascade_key, failed_cascade_key], _CTX) + assert prepared is not None + primary_tier.complete_store(prepared.keys_to_store, _CTX, success=True) + cascade_job = manager.create_store_job([cascade_key], _CTX) + failed_cascade_job = manager.create_store_job([failed_cascade_key], _CTX) + + promotion_prepared = primary_tier.prepare_write([promotion_key], _CTX) + assert promotion_prepared is not None + promotion_job = JobMetadata( + job_id=manager._next_job_id(), + keys=[promotion_key], + block_ids=promotion_prepared.store_spec.block_ids, + is_promotion=True, + req_context=_CTX, + ) + manager._transfer_jobs[promotion_job.job_id] = promotion_job + + failed_promotion_prepared = primary_tier.prepare_write([failed_promotion_key], _CTX) + assert failed_promotion_prepared is not None + failed_promotion_job = JobMetadata( + job_id=manager._next_job_id(), + keys=[failed_promotion_key], + block_ids=failed_promotion_prepared.store_spec.block_ids, + is_promotion=True, + req_context=_CTX, + ) + manager._transfer_jobs[failed_promotion_job.job_id] = failed_promotion_job + + tier.finished_jobs = [ + JobResult( + job_id=cascade_job.job_id, + success=True, + transfer_size=100, + transfer_time=0.5, + ), + JobResult( + job_id=promotion_job.job_id, + success=True, + transfer_size=80, + transfer_time=0.25, + ), + JobResult(job_id=failed_cascade_job.job_id, success=False), + JobResult(job_id=failed_promotion_job.job_id, success=False), + ] + + manager._process_finished_jobs() + + stats = manager.get_stats() + assert stats is not None + values = stats.data["data"] + label = ("0:fs",) + assert values[TieringOffloadingMetrics.WRITE_BYTES][label] == 100 + assert values[TieringOffloadingMetrics.WRITE_TIME][label] == 0.5 + assert values[TieringOffloadingMetrics.READ_BYTES][label] == 80 + assert values[TieringOffloadingMetrics.READ_TIME][label] == 0.25 + assert values[TieringOffloadingMetrics.CASCADE_JOB_FAILURES][label] == 1 + assert values[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES][label] == 1 + + class TestExampleSecondaryTierManager: """Tests for ExampleSecondaryTierManager implementation.""" diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 2075b556fbaf..eeee061aa194 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -6,7 +6,7 @@ from abc import ABC, abstractmethod from collections.abc import Collection, Iterable -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar import numpy as np @@ -39,6 +39,14 @@ class TieringOffloadingMetrics: LOOKUP_SYNC_DELAY = "vllm:kv_offload_tiering_lookup_sync_delay_seconds" LOOKUP_ASYNC_DELAY = "vllm:kv_offload_tiering_lookup_async_delay_seconds" + READ_BYTES = "vllm:kv_offload_tiering_read_bytes" + READ_TIME = "vllm:kv_offload_tiering_read_time" + WRITE_BYTES = "vllm:kv_offload_tiering_write_bytes" + WRITE_TIME = "vllm:kv_offload_tiering_write_time" + PROMOTION_JOB_FAILURES = "vllm:kv_offload_tiering_promotion_job_failures" + CASCADE_JOB_FAILURES = "vllm:kv_offload_tiering_cascade_job_failures" + BLOCK_QUERIES = "vllm:kv_offload_tiering_block_queries" + BLOCK_HITS = "vllm:kv_offload_tiering_block_hits" @dataclass @@ -63,6 +71,8 @@ class JobResult: # keys that were successfully loaded. None means all keys share the fate # indicated by `success`. Must be a subset of the job's original keys. successful_keys: Collection[OffloadKey] | None = None + transfer_size: int | None = field(default=None, compare=False) + transfer_time: float | None = field(default=None, compare=False) class ParentManager(ABC): diff --git a/vllm/v1/kv_offload/tiering/example/manager.py b/vllm/v1/kv_offload/tiering/example/manager.py index 1a78a3feb71d..dd45abcc1996 100644 --- a/vllm/v1/kv_offload/tiering/example/manager.py +++ b/vllm/v1/kv_offload/tiering/example/manager.py @@ -73,6 +73,8 @@ def __init__( # Completed jobs waiting to be retrieved by get_finished_jobs() self.completed_jobs: list[JobResult] = [] + assert primary_kv_view.strides is not None + self._block_size = primary_kv_view.strides[0] @override def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: @@ -106,7 +108,14 @@ def submit_store(self, job_metadata: JobMetadata) -> None: for key in keys: self.blocks[key] = True - self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True)) + self.completed_jobs.append( + JobResult( + job_id=job_metadata.job_id, + success=True, + transfer_size=len(keys) * self._block_size, + transfer_time=0.0, + ) + ) @override def submit_load(self, job_metadata: JobMetadata) -> None: @@ -131,7 +140,14 @@ def submit_load(self, job_metadata: JobMetadata) -> None: ) return - self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True)) + self.completed_jobs.append( + JobResult( + job_id=job_metadata.job_id, + success=True, + transfer_size=len(keys) * self._block_size, + transfer_time=0.0, + ) + ) @override def get_finished_jobs(self) -> Iterable[JobResult]: diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index 560e41093a81..a8fb83f717a8 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -18,6 +18,7 @@ import functools import json import os +import time from collections.abc import Iterable from typing import TYPE_CHECKING, ClassVar @@ -201,6 +202,8 @@ def __init__( ) self._lookup_manager = FsAsyncLookupManager(tier=self, tier_type=self.tier_type) + self._job_transfer_sizes: dict[int, int] = {} + self._job_submitted_at: dict[int, float] = {} @override def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: @@ -215,16 +218,19 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: @override def submit_store(self, job_metadata: JobMetadata) -> None: + keys = list(job_metadata.keys) if self.events is not None: - self._store_job_keys[job_metadata.job_id] = list(job_metadata.keys) + self._store_job_keys[job_metadata.job_id] = keys task = functools.partial( batch_store_block, - [self.file_mapper.get_file_name(key) for key in job_metadata.keys], + [self.file_mapper.get_file_name(key) for key in keys], self._primary_kv_view, [int(bid) * self._block_size for bid in job_metadata.block_ids], self._block_size, self._use_o_direct, ) + self._job_transfer_sizes[job_metadata.job_id] = len(keys) * self._block_size + self._job_submitted_at[job_metadata.job_id] = time.monotonic() self._pool.enqueue_store(job_metadata.job_id, 1, [task]) @override @@ -264,6 +270,8 @@ def load_task() -> None: ) raise + self._job_transfer_sizes[job_id] = len(keys) * self._block_size + self._job_submitted_at[job_id] = time.monotonic() self._pool.enqueue_load(job_id, 1, [load_task]) @override @@ -271,6 +279,7 @@ def get_finished_jobs(self) -> Iterable[JobResult]: """Collect finished jobs; a failed promotion marks only its failed keys as a miss here (scheduler thread).""" results = [] + now = time.monotonic() for job_id, success in self._pool.get_finished(): if self.events is not None: keys = self._store_job_keys.pop(job_id, None) @@ -283,6 +292,9 @@ def get_finished_jobs(self) -> Iterable[JobResult]: locality=self.locality, ) ) + submitted_at = self._job_submitted_at.pop(job_id, None) + transfer_time = None if submitted_at is None else now - submitted_at + transfer_size = self._job_transfer_sizes.pop(job_id, None) load_keys = self._load_job_keys.pop(job_id, None) num_succeeded = self._load_progress.pop(job_id, 0) if load_keys is not None and not success: @@ -293,15 +305,29 @@ def get_finished_jobs(self) -> Iterable[JobResult]: successful = load_keys[:num_succeeded] failed = load_keys[num_succeeded:] self._lookup_manager.mark_miss(failed) + successful_transfer_size = ( + num_succeeded * self._block_size + if transfer_size is not None + else None + ) results.append( JobResult( job_id=job_id, success=False, successful_keys=tuple(successful) if successful else None, + transfer_size=successful_transfer_size, + transfer_time=transfer_time, ) ) else: - results.append(JobResult(job_id=job_id, success=success)) + results.append( + JobResult( + job_id=job_id, + success=success, + transfer_size=transfer_size, + transfer_time=transfer_time, + ) + ) return results @override diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 5b8f92a45845..e0b52a2bf71f 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -220,7 +220,6 @@ def __init__( tier: _SecondaryTierFacingParent(self, tier) for tier in self.secondary_tiers } - # Buffers manager-level observations (e.g. lookup delay) between # get_stats() calls; merged in and reset each time get_stats() runs. self._stats = OffloadingConnectorStats() @@ -231,6 +230,46 @@ def _next_job_id(self) -> JobId: self._job_id_counter += 1 return job_id + @staticmethod + def _tier_label(tier_idx: int, tier: SecondaryTierManager) -> tuple[str]: + return (f"{tier_idx}:{tier.tier_type}",) + + def _record_finished_job_stats( + self, + tier_idx: int, + tier: SecondaryTierManager, + job_metadata: JobMetadata, + completed_job: JobResult, + ) -> None: + labelvalues = self._tier_label(tier_idx, tier) + if not completed_job.success: + failure_metric = ( + TieringOffloadingMetrics.PROMOTION_JOB_FAILURES + if job_metadata.is_promotion + else TieringOffloadingMetrics.CASCADE_JOB_FAILURES + ) + self._stats.increase_counter(failure_metric, labelvalues=labelvalues) + return + + bytes_metric = ( + TieringOffloadingMetrics.READ_BYTES + if job_metadata.is_promotion + else TieringOffloadingMetrics.WRITE_BYTES + ) + time_metric = ( + TieringOffloadingMetrics.READ_TIME + if job_metadata.is_promotion + else TieringOffloadingMetrics.WRITE_TIME + ) + if completed_job.transfer_size is not None: + self._stats.increase_counter( + bytes_metric, completed_job.transfer_size, labelvalues + ) + if completed_job.transfer_time is not None: + self._stats.increase_counter( + time_metric, completed_job.transfer_time, labelvalues + ) + def _maybe_process_finished_jobs(self): """ Poll secondary tiers for completed jobs (at most once per step). @@ -296,6 +335,7 @@ def _process_finished_jobs(self): f"Finished job_id {job_id} from tier #{i}" f" ({tier.tier_type}) not in _transfer_jobs" ) + self._record_finished_job_stats(i, tier, job_metadata, completed_job) if job_metadata.is_promotion: # secondary→primary transfer (promotion) completed. @@ -353,13 +393,22 @@ def lookup( lookup_start = time.monotonic() any_retry = False - for tier in self.secondary_tiers: + for i, tier in enumerate(self.secondary_tiers): if tier is exclude_tier: continue if not req_context.load_tier_filter.allows(tier.medium, tier.locality): continue + labelvalues = self._tier_label(i, tier) + self._stats.increase_counter( + TieringOffloadingMetrics.BLOCK_QUERIES, + labelvalues=labelvalues, + ) result = tier.lookup(key, req_context) if result is LookupResult.HIT: + self._stats.increase_counter( + TieringOffloadingMetrics.BLOCK_HITS, + labelvalues=labelvalues, + ) promoted = self._initiate_promotion(tier, key, req_context) self._accumulate_lookup_sync_delay(req_state, lookup_start) if ( @@ -855,6 +904,13 @@ def get_stats(self) -> OffloadingConnectorStats | None: if stats is not None and stats.is_empty(): stats = None + if not self._stats.is_empty(): + if stats is None: + stats = self._stats + else: + stats.aggregate(self._stats) + self._stats = OffloadingConnectorStats() + for tier in self.secondary_tiers: tier_stats = tier.get_stats() if tier_stats is None or tier_stats.is_empty(): diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index fcfea6909fd3..56de3e53f389 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -58,6 +58,8 @@ class TransferEntry(NamedTuple): xfer_handle: "nixl_xfer_handle" files_desc: object obj_handle: "nixl_prepped_dlist_handle" + transfer_size: int + submitted_at: float class ObjAsyncLookupManager(AsyncLookupManager): @@ -267,7 +269,13 @@ def _submit_transfer( self._pending_results.append(JobResult(job_id=job_id, success=False)) return - self._transfers[job_id] = TransferEntry(xfer_handle, files_desc, obj_handle) + self._transfers[job_id] = TransferEntry( + xfer_handle=xfer_handle, + files_desc=files_desc, + obj_handle=obj_handle, + transfer_size=len(block_ids_list) * self._block_size_bytes, + submitted_at=time.monotonic(), + ) def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: result = self._lookup_manager.lookup(key, req_context) @@ -341,7 +349,14 @@ def _poll_active_transfers(self) -> None: logger.warning("deregister_memory failed for job %d: %s", job_id, exc) del self._transfers[job_id] - self._pending_results.append(JobResult(job_id=job_id, success=success)) + self._pending_results.append( + JobResult( + job_id=job_id, + success=success, + transfer_size=entry.transfer_size, + transfer_time=time.monotonic() - entry.submitted_at, + ) + ) def get_finished_jobs(self) -> Iterable[JobResult]: """Poll transfers; a failed promotion marks its cached verdicts False diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index f1f557ce38e7..655a9e83b55e 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -444,6 +444,8 @@ def submit_store(self, job_metadata: JobMetadata) -> None: self._finished_jobs.append(JobResult(job_id=job_id, success=False)) return + + # Fast path: a session has already received FetchMsg for this id, # so we can route the batch straight into its ServerRole. session = self._kv_to_session.get(kv_request_id) @@ -475,6 +477,7 @@ def submit_load(self, job_metadata: JobMetadata) -> None: keys = list(job_metadata.keys) block_ids = job_metadata.block_ids + source = job_metadata.req_context.get_state(P2PSourceInfo) logger.debug( "P2P %s: submit_load ENTRY job_id=%d blocks=%d kv_request_id=%s peer=%s", @@ -718,9 +721,7 @@ def _reap_unbound_stores(self) -> None: batches = self._unbound_stores.pop(kid) self._failed_req_ids.add(kid) for batch in batches: - self._finished_jobs.append( - JobResult(job_id=batch.job_id, success=False) - ) + self._finished_jobs.append(JobResult(job_id=batch.job_id, success=False)) logger.warning( "P2P %s: unbound store kv_request_id=%s timed out after %.0fs " "without a fetch — failing %d job(s)", @@ -755,15 +756,11 @@ def _poll_once(self) -> None: for session in self._sessions.values(): result = session.poll() for lr in result.loads: - self._finished_jobs.append( - JobResult(job_id=lr.job_id, success=lr.success) - ) + self._finished_jobs.append(JobResult(job_id=lr.job_id, success=lr.success)) if not lr.success: self._failed_req_ids.add(lr.kv_request_id) for sr in result.stores: - self._finished_jobs.append( - JobResult(job_id=sr.job_id, success=sr.success) - ) + self._finished_jobs.append(JobResult(job_id=sr.job_id, success=sr.success)) # Bind kv_request_id → session for any FetchMsg this tick and # replay any submit_store batches parked while no peer was # asking. ServerRole.on_fetch already recorded the demand @@ -797,9 +794,7 @@ def shutdown(self) -> None: # leak them; the manager is going away after this call. for batches in self._unbound_stores.values(): for batch in batches: - self._finished_jobs.append( - JobResult(job_id=batch.job_id, success=False) - ) + self._finished_jobs.append(JobResult(job_id=batch.job_id, success=False)) self._unbound_stores.clear() self._control.close() self._data.close() diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index ccce426bfdf6..90f954584269 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -61,6 +61,7 @@ from vllm.logger import init_logger from vllm.v1.kv_offload.base import ( CanonicalKVCaches, + OffloadingCounterMetadata, OffloadingHistogramMetadata, OffloadingManager, OffloadingMetricMetadata, @@ -144,6 +145,46 @@ def build_metric_definitions( ), ) ) + for metric_name, documentation in ( + ( + TieringOffloadingMetrics.READ_BYTES, + "Total bytes read from secondary tiers into the primary tier.", + ), + ( + TieringOffloadingMetrics.READ_TIME, + "Total time spent reading from secondary tiers into the primary " + "tier, in seconds.", + ), + ( + TieringOffloadingMetrics.WRITE_BYTES, + "Total bytes written from the primary tier to secondary tiers.", + ), + ( + TieringOffloadingMetrics.WRITE_TIME, + "Total time spent writing from the primary tier to secondary " + "tiers, in seconds.", + ), + ( + TieringOffloadingMetrics.PROMOTION_JOB_FAILURES, + "Number of failed secondary-tier promotion jobs.", + ), + ( + TieringOffloadingMetrics.CASCADE_JOB_FAILURES, + "Number of failed secondary-tier cascade jobs.", + ), + ( + TieringOffloadingMetrics.BLOCK_QUERIES, + "Number of block lookup queries sent to secondary tiers.", + ), + ( + TieringOffloadingMetrics.BLOCK_HITS, + "Number of block lookup hits in secondary tiers.", + ), + ): + metrics[metric_name] = OffloadingCounterMetadata( + documentation=documentation, + labelnames=("tier",), + ) secondary_tier_configs = extra_config.get("secondary_tiers", []) if not isinstance(secondary_tier_configs, list): raise ValueError("secondary_tiers must be a list of tier configurations") From 5105c09d00c7635d86f3fd864715a8651572e793 Mon Sep 17 00:00:00 2001 From: srinivas_oo7 Date: Tue, 21 Jul 2026 21:34:44 -0500 Subject: [PATCH 02/13] Add remaining tiering offloading metrics Signed-off-by: srinivas_oo7 --- .../tiering/test_tiering_offloading.py | 111 ++++++++++++++++++ vllm/v1/kv_offload/tiering/base.py | 7 ++ vllm/v1/kv_offload/tiering/manager.py | 87 ++++++++++++-- vllm/v1/kv_offload/tiering/spec.py | 33 ++++++ 4 files changed, 228 insertions(+), 10 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 4d354d398cca..d78bebbf6861 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -29,6 +29,7 @@ Medium, OffloadingCounterMetadata, OffloadingEvent, + OffloadingGaugeMetadata, OffloadKey, OffloadPolicy, ReqContext, @@ -153,6 +154,18 @@ def test_tiering_spec_collects_secondary_metric_definitions(monkeypatch): assert metrics[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES].labelnames == ( "tier", ) + assert isinstance( + metrics[TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES], + OffloadingCounterMetadata, + ) + for metric_name in ( + TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC, + TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC, + TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS, + TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS, + ): + assert isinstance(metrics[metric_name], OffloadingGaugeMetadata) + assert metrics[metric_name].labelnames == ("tier",) def test_tiering_manager_aggregates_secondary_stats(): @@ -303,6 +316,104 @@ def test_tiering_manager_records_finished_job_metrics(): assert values[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES][label] == 1 +def test_tiering_manager_reports_active_job_and_primary_usage_gauges(): + mock_region = _mock_mmap_region(6) + primary_tier = CPUPrimaryTierOffloadingManager( + num_blocks=6, mmap_region=mock_region + ) + tier0 = MetricsSecondaryTierManager( + offloading_spec=_MOCK_OFFLOADING_SPEC, + primary_kv_view=mock_region.create_kv_memoryview(), + tier_type="fs", + ) + tier1 = MetricsSecondaryTierManager( + offloading_spec=_MOCK_OFFLOADING_SPEC, + primary_kv_view=mock_region.create_kv_memoryview(), + tier_type="p2p", + ) + manager = TieringOffloadingManager( + primary_tier=primary_tier, + secondary_tiers=[tier0, tier1], + ) + manager.on_new_request(_CTX) + cascade_keys = to_keys([0, 1]) + promotion_key = to_keys([2])[0] + + prepared = manager.prepare_store(cascade_keys, _CTX) + assert prepared is not None + manager.complete_store(cascade_keys, _CTX, success=True) + tier0.lookup_result = LookupResult.HIT + assert manager.lookup(promotion_key, _CTX) is LookupResult.RETRY + manager.on_schedule_end(ScheduleEndContext(new_req_ids=(), preempted_req_ids=())) + + stats = manager.get_stats() + assert stats is not None + values = stats.data["data"] + fs_label = ("0:fs",) + p2p_label = ("1:p2p",) + assert values[TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC][ + fs_label + ] == pytest.approx(2 / 6) + assert values[TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC][ + p2p_label + ] == pytest.approx(2 / 6) + assert values[TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC][ + fs_label + ] == pytest.approx(1 / 6) + assert values[TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC][p2p_label] == 0 + assert values[TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS][fs_label] == 1 + assert values[TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS][p2p_label] == 1 + assert values[TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS][fs_label] == 1 + assert values[TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS][p2p_label] == 0 + + tier0.finished_jobs = [ + JobResult(job_id=tier0.submitted_stores[0].job_id, success=True), + JobResult(job_id=tier0.submitted_loads[0].job_id, success=True), + ] + tier1.finished_jobs = [ + JobResult(job_id=tier1.submitted_stores[0].job_id, success=True), + ] + manager._process_finished_jobs() + + stats = manager.get_stats() + assert stats is not None + values = stats.data["data"] + assert values[TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC][fs_label] == 0 + assert values[TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC][p2p_label] == 0 + assert values[TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC][fs_label] == 0 + assert values[TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC][p2p_label] == 0 + assert values[TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS][fs_label] == 0 + assert values[TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS][p2p_label] == 0 + assert values[TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS][fs_label] == 0 + assert values[TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS][p2p_label] == 0 + + +def test_tiering_manager_records_promotion_allocation_failures(): + mock_region = _mock_mmap_region(1) + primary_tier = CPUPrimaryTierOffloadingManager( + num_blocks=1, mmap_region=mock_region + ) + tier = MetricsSecondaryTierManager( + offloading_spec=_MOCK_OFFLOADING_SPEC, + primary_kv_view=mock_region.create_kv_memoryview(), + tier_type="fs", + ) + tier.lookup_result = LookupResult.HIT + manager = TieringOffloadingManager( + primary_tier=primary_tier, + secondary_tiers=[tier], + ) + manager.on_new_request(_CTX) + assert primary_tier.prepare_store(to_keys([0]), _CTX) is not None + + assert manager.lookup(to_keys([1])[0], _CTX) is LookupResult.MISS + + stats = manager.get_stats() + assert stats is not None + values = stats.data["data"] + assert values[TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES][()] == 1 + + class TestExampleSecondaryTierManager: """Tests for ExampleSecondaryTierManager implementation.""" diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index eeee061aa194..a760a4ab3628 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -47,6 +47,13 @@ class TieringOffloadingMetrics: CASCADE_JOB_FAILURES = "vllm:kv_offload_tiering_cascade_job_failures" BLOCK_QUERIES = "vllm:kv_offload_tiering_block_queries" BLOCK_HITS = "vllm:kv_offload_tiering_block_hits" + PRIMARY_WRITE_USAGE_PERC = "vllm:kv_offload_tiering_primary_write_usage_perc" + PRIMARY_READ_USAGE_PERC = "vllm:kv_offload_tiering_primary_read_usage_perc" + PROMOTION_ALLOCATION_FAILURES = ( + "vllm:kv_offload_tiering_promotion_allocation_failures" + ) + ACTIVE_PROMOTION_JOBS = "vllm:kv_offload_tiering_active_promotion_jobs" + ACTIVE_CASCADE_JOBS = "vllm:kv_offload_tiering_active_cascade_jobs" @dataclass diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index e0b52a2bf71f..9b24522537cf 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -153,7 +153,7 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: def create_store_job( self, keys: Collection[OffloadKey], req_context: ReqContext ) -> JobMetadata: - return self._m.create_store_job(keys, req_context) + return self._m.create_store_job(keys, req_context, self._origin) def on_request_finished(self, req_context: ReqContext) -> None: return self._m.on_request_finished(req_context, exclude_tier=self._origin) @@ -197,6 +197,7 @@ def __init__( # True: secondary → primary (promotion) # False: primary → secondary (cascade) self._transfer_jobs: dict[JobId, JobMetadata] = {} + self._transfer_job_tiers: dict[JobId, SecondaryTierManager] = {} # Pending promotion requests accumulated during lookup() calls; flushed # as one batched submit_load() per (tier, request) in on_schedule_end(). @@ -270,6 +271,62 @@ def _record_finished_job_stats( time_metric, completed_job.transfer_time, labelvalues ) + def _record_active_transfer_stats(self, stats: OffloadingConnectorStats) -> None: + active_promotion_jobs: dict[SecondaryTierManager, int] = {} + active_cascade_jobs: dict[SecondaryTierManager, int] = {} + primary_write_blocks: dict[SecondaryTierManager, set[int]] = {} + primary_read_blocks: dict[SecondaryTierManager, set[int]] = {} + + for job_id, job_metadata in self._transfer_jobs.items(): + tier = self._transfer_job_tiers.get(job_id) + if tier is None: + continue + block_ids = set(job_metadata.block_ids.tolist()) + if job_metadata.is_promotion: + active_promotion_jobs[tier] = active_promotion_jobs.get(tier, 0) + 1 + primary_write_blocks.setdefault(tier, set()).update(block_ids) + else: + active_cascade_jobs[tier] = active_cascade_jobs.get(tier, 0) + 1 + primary_read_blocks.setdefault(tier, set()).update(block_ids) + + for pending_tier, pending_by_ctx in self._pending_load_submissions.items(): + for pending in pending_by_ctx.values(): + primary_write_blocks.setdefault(pending_tier, set()).update( + pending.block_ids + ) + + num_primary_blocks = self.primary_tier._num_blocks + for tier_idx, tier in enumerate(self.secondary_tiers): + labelvalues = self._tier_label(tier_idx, tier) + write_blocks = len(primary_write_blocks.get(tier, set())) + read_blocks = len(primary_read_blocks.get(tier, set())) + write_usage = ( + write_blocks / num_primary_blocks if num_primary_blocks > 0 else 0.0 + ) + read_usage = ( + read_blocks / num_primary_blocks if num_primary_blocks > 0 else 0.0 + ) + stats.set_gauge( + TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC, + write_usage, + labelvalues, + ) + stats.set_gauge( + TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC, + read_usage, + labelvalues, + ) + stats.set_gauge( + TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS, + active_promotion_jobs.get(tier, 0), + labelvalues, + ) + stats.set_gauge( + TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS, + active_cascade_jobs.get(tier, 0), + labelvalues, + ) + def _maybe_process_finished_jobs(self): """ Poll secondary tiers for completed jobs (at most once per step). @@ -331,6 +388,7 @@ def _process_finished_jobs(self): for completed_job in tier.get_finished_jobs(): job_id = completed_job.job_id job_metadata = self._transfer_jobs.pop(job_id, None) + self._transfer_job_tiers.pop(job_id, None) assert job_metadata is not None, ( f"Finished job_id {job_id} from tier #{i}" f" ({tier.tier_type}) not in _transfer_jobs" @@ -488,6 +546,9 @@ def _initiate_promotion( if primary_write_result is None: # Primary tier is full; caller should treat the block as unavailable # rather than retrying indefinitely. + self._stats.increase_counter( + TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES + ) return False store_spec = primary_write_result.store_spec @@ -525,6 +586,7 @@ def _flush_pending_promotions(self) -> None: req_context=entry.req_context, ) self._transfer_jobs[job_id] = job_metadata + self._transfer_job_tiers[job_id] = tier tier.submit_load(job_metadata) self._pending_load_submissions.clear() @@ -659,7 +721,7 @@ def _cascade_existing_blocks_to_request_level_tiers( return for tier in request_level_tiers: - job_metadata = self.create_store_job(ready_keys, req_context) + job_metadata = self.create_store_job(ready_keys, req_context, tier) tier.submit_store(job_metadata) @override @@ -697,7 +759,7 @@ def complete_store( # eviction during the async transfer). One prepare_read() call per # secondary tier. for tier in self.secondary_tiers: - job_metadata = self.create_store_job(keys, req_context) + job_metadata = self.create_store_job(keys, req_context, tier) tier.submit_store(job_metadata) # Note: The async transfers are now in flight. Their completion is @@ -712,6 +774,7 @@ def create_store_job( self, keys: Collection[OffloadKey], req_context: ReqContext, + tier: SecondaryTierManager | None = None, ) -> JobMetadata: """Pin blocks in the primary tier and create a tracked store job. @@ -733,6 +796,8 @@ def create_store_job( req_context=req_context, ) self._transfer_jobs[job_id] = job_metadata + if tier is not None: + self._transfer_job_tiers[job_id] = tier return job_metadata @override @@ -879,6 +944,7 @@ def reset_cache(self) -> None: # reset below invalidates; their submit_load() has not yet been # called so no tier I/O is touching that memory. self._pending_load_submissions.clear() + self._transfer_job_tiers.clear() finished_req_ids = [] for req_id, state in self._req_state.items(): @@ -904,6 +970,14 @@ def get_stats(self) -> OffloadingConnectorStats | None: if stats is not None and stats.is_empty(): stats = None + active_transfer_stats = OffloadingConnectorStats() + self._record_active_transfer_stats(active_transfer_stats) + if not active_transfer_stats.is_empty(): + if stats is None: + stats = active_transfer_stats + else: + stats.aggregate(active_transfer_stats) + if not self._stats.is_empty(): if stats is None: stats = self._stats @@ -920,13 +994,6 @@ def get_stats(self) -> OffloadingConnectorStats | None: else: stats.aggregate(tier_stats) - if not self._stats.is_empty(): - if stats is None: - stats = self._stats - else: - stats.aggregate(self._stats) - self._stats = OffloadingConnectorStats() - return stats @override diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 90f954584269..16ec9a3f85bb 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -62,6 +62,7 @@ from vllm.v1.kv_offload.base import ( CanonicalKVCaches, OffloadingCounterMetadata, + OffloadingGaugeMetadata, OffloadingHistogramMetadata, OffloadingManager, OffloadingMetricMetadata, @@ -185,6 +186,38 @@ def build_metric_definitions( documentation=documentation, labelnames=("tier",), ) + metrics[TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES] = ( + OffloadingCounterMetadata( + documentation=( + "Number of promotion attempts that failed because the " + "primary tier could not allocate space." + ), + ) + ) + for metric_name, documentation in ( + ( + TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC, + "Current fraction of primary-tier space used by writes from " + "secondary tiers.", + ), + ( + TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC, + "Current fraction of primary-tier space used by reads to " + "secondary tiers.", + ), + ( + TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS, + "Number of active secondary-tier promotion jobs.", + ), + ( + TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS, + "Number of active secondary-tier cascade jobs.", + ), + ): + metrics[metric_name] = OffloadingGaugeMetadata( + documentation=documentation, + labelnames=("tier",), + ) secondary_tier_configs = extra_config.get("secondary_tiers", []) if not isinstance(secondary_tier_configs, list): raise ValueError("secondary_tiers must be a list of tier configurations") From 795cd324f0b72137ce220fb707b090b727e8ec57 Mon Sep 17 00:00:00 2001 From: srinivas_oo7 Date: Wed, 22 Jul 2026 09:17:17 -0500 Subject: [PATCH 03/13] Address tiering metrics review feedback Signed-off-by: srinivas_oo7 --- .../v1/kv_offload/tiering/p2p/test_manager.py | 6 +- .../kv_offload/tiering/p2p/test_sessions.py | 6 +- tests/v1/kv_offload/tiering/test_fs_tier.py | 6 +- tests/v1/kv_offload/tiering/test_obj_tier.py | 9 +- .../tiering/test_tiering_offloading.py | 60 ++--- vllm/v1/kv_offload/tiering/base.py | 13 +- vllm/v1/kv_offload/tiering/example/manager.py | 8 +- vllm/v1/kv_offload/tiering/fs/manager.py | 33 ++- vllm/v1/kv_offload/tiering/fs/thread_pool.py | 33 ++- vllm/v1/kv_offload/tiering/manager.py | 220 +++++++++++------- vllm/v1/kv_offload/tiering/obj/manager.py | 18 +- vllm/v1/kv_offload/tiering/p2p/manager.py | 8 +- vllm/v1/kv_offload/tiering/spec.py | 130 ++++++----- 13 files changed, 320 insertions(+), 230 deletions(-) diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index c8a9d7f3ea35..df7a35219206 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -16,7 +16,7 @@ import pytest from vllm.v1.kv_offload.base import LookupResult, ReqContext, ScheduleEndContext -from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult +from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob from vllm.v1.kv_offload.tiering.p2p import manager as manager_module from vllm.v1.kv_offload.tiering.p2p.manager import ( _UNBOUND_STORE_TIMEOUT_S, @@ -86,12 +86,12 @@ def _job_metadata( keys: list[bytes] | None = None, block_ids: list[int] | None = None, kv_params: dict | None = None, -) -> JobMetadata: +) -> TransferJob: if keys is None: keys = [b"key1"] if block_ids is None: block_ids = list(range(len(keys))) - return JobMetadata( + return TransferJob( job_id=job_id, keys=keys, block_ids=np.array(block_ids), diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index 20e4bad21863..d5692b6690c9 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -24,7 +24,7 @@ ReqContext, RequestOffloadingContext, ) -from vllm.v1.kv_offload.tiering.base import JobMetadata +from vllm.v1.kv_offload.tiering.base import TransferJob from vllm.v1.kv_offload.tiering.p2p.session import ( LoadResult, P2PSession, @@ -256,13 +256,13 @@ def create_store_job( self, keys: Sequence[OffloadKey], ctx: ReqContext, - ) -> JobMetadata: + ) -> TransferJob: keys_list = list(keys) self.calls.append(("create_store_job", tuple(keys_list), ctx.req_id)) block_ids = np.array([self.stored[k] for k in keys_list], dtype=np.int32) job_id = self._next_job_id self._next_job_id += 1 - return JobMetadata( + return TransferJob( job_id=job_id, keys=keys_list, block_ids=block_ids, diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 7d850ca62ae6..f007eb49e190 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -35,7 +35,7 @@ OffloadingModelConfig, OffloadingParallelConfig, ) -from vllm.v1.kv_offload.tiering.base import JobMetadata +from vllm.v1.kv_offload.tiering.base import TransferJob from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory from vllm.v1.kv_offload.tiering.fs.manager import ( FileSystemTierManager, @@ -105,10 +105,10 @@ def make_job( keys: list[OffloadKey], block_ids: list[int] | None = None, is_promotion: bool = False, -) -> JobMetadata: +) -> TransferJob: if block_ids is None: block_ids = list(range(len(keys))) - return JobMetadata( + return TransferJob( job_id=job_id, keys=keys, block_ids=np.array(block_ids, dtype=np.int64), diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index 0e86038dd6ac..54ec4c1f9f91 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -34,7 +34,7 @@ OffloadingModelConfig, OffloadingParallelConfig, ) -from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult +from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob from vllm.v1.kv_offload.tiering.manager import ( CPUPrimaryTierOffloadingManager, TieringOffloadingManager, @@ -105,10 +105,10 @@ def make_job( job_id: int, keys: list[OffloadKey], block_ids: list[int] | None = None, -) -> JobMetadata: +) -> TransferJob: if block_ids is None: block_ids = list(range(len(keys))) - return JobMetadata( + return TransferJob( job_id=job_id, keys=keys, block_ids=np.array(block_ids, dtype=np.int64), @@ -201,6 +201,9 @@ def release_xfer_handle(self, handle): def release_dlist_handle(self, handle): pass + def get_xfer_telemetry(self, handle): + return SimpleNamespace(xferDuration=1000) + def _query_memory(self, queries, mem_type, agent_name): return [object() if q[3] in self._stored_obj_keys else None for q in queries] diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index d78bebbf6861..c70c12d0c26f 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -40,10 +40,10 @@ make_offload_key, ) from vllm.v1.kv_offload.tiering.base import ( - JobMetadata, JobResult, SecondaryTierManager, TieringOffloadingMetrics, + TransferJob, ) from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory @@ -105,17 +105,17 @@ def __init__(self, *args, **kwargs): self.stats: OffloadingConnectorStats | None = None self.lookup_result = LookupResult.MISS self.finished_jobs: list[JobResult] = [] - self.submitted_loads: list[JobMetadata] = [] - self.submitted_stores: list[JobMetadata] = [] + self.submitted_loads: list[TransferJob] = [] + self.submitted_stores: list[TransferJob] = [] def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: return self.lookup_result - def submit_store(self, job_metadata: JobMetadata) -> None: + def submit_store(self, job_metadata: TransferJob) -> None: self.submitted_stores.append(job_metadata) return - def submit_load(self, job_metadata: JobMetadata) -> None: + def submit_load(self, job_metadata: TransferJob) -> None: self.submitted_loads.append(job_metadata) return @@ -226,16 +226,20 @@ def test_tiering_manager_records_secondary_lookup_metrics(): primary_tier=primary_tier, secondary_tiers=[tier0, tier1], ) + manager.on_new_request(_CTX) - assert manager.lookup(to_keys([1])[0], _CTX) is LookupResult.RETRY + key = to_keys([1])[0] + assert manager.lookup(key, _CTX) is LookupResult.RETRY + assert manager.lookup(key, _CTX) is LookupResult.HIT_PENDING stats = manager.get_stats() assert stats is not None values = stats.data["data"] - assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("0:fs",)] == 1 - assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("1:p2p",)] == 1 - assert values[TieringOffloadingMetrics.BLOCK_HITS][("1:p2p",)] == 1 - assert ("0:fs",) not in values[TieringOffloadingMetrics.BLOCK_HITS] + assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("0:primary",)] == 1 + assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("1:fs",)] == 1 + assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("2:p2p",)] == 1 + assert values[TieringOffloadingMetrics.BLOCK_HITS][("2:p2p",)] == 1 + assert ("1:fs",) not in values[TieringOffloadingMetrics.BLOCK_HITS] def test_tiering_manager_records_finished_job_metrics(): @@ -260,42 +264,40 @@ def test_tiering_manager_records_finished_job_metrics(): prepared = primary_tier.prepare_store([cascade_key, failed_cascade_key], _CTX) assert prepared is not None primary_tier.complete_store(prepared.keys_to_store, _CTX, success=True) - cascade_job = manager.create_store_job([cascade_key], _CTX) - failed_cascade_job = manager.create_store_job([failed_cascade_key], _CTX) + cascade_job = manager.create_store_job([cascade_key], _CTX, 0) + failed_cascade_job = manager.create_store_job([failed_cascade_key], _CTX, 0) promotion_prepared = primary_tier.prepare_write([promotion_key], _CTX) assert promotion_prepared is not None - promotion_job = JobMetadata( + promotion_job = TransferJob( job_id=manager._next_job_id(), keys=[promotion_key], block_ids=promotion_prepared.store_spec.block_ids, is_promotion=True, req_context=_CTX, ) - manager._transfer_jobs[promotion_job.job_id] = promotion_job + manager._register_job(promotion_job, 0) failed_promotion_prepared = primary_tier.prepare_write([failed_promotion_key], _CTX) assert failed_promotion_prepared is not None - failed_promotion_job = JobMetadata( + failed_promotion_job = TransferJob( job_id=manager._next_job_id(), keys=[failed_promotion_key], block_ids=failed_promotion_prepared.store_spec.block_ids, is_promotion=True, req_context=_CTX, ) - manager._transfer_jobs[failed_promotion_job.job_id] = failed_promotion_job + manager._register_job(failed_promotion_job, 0) tier.finished_jobs = [ JobResult( job_id=cascade_job.job_id, success=True, - transfer_size=100, transfer_time=0.5, ), JobResult( job_id=promotion_job.job_id, success=True, - transfer_size=80, transfer_time=0.25, ), JobResult(job_id=failed_cascade_job.job_id, success=False), @@ -307,10 +309,10 @@ def test_tiering_manager_records_finished_job_metrics(): stats = manager.get_stats() assert stats is not None values = stats.data["data"] - label = ("0:fs",) - assert values[TieringOffloadingMetrics.WRITE_BYTES][label] == 100 + label = ("1:fs",) + assert values[TieringOffloadingMetrics.WRITE_BYTES][label] == 16 assert values[TieringOffloadingMetrics.WRITE_TIME][label] == 0.5 - assert values[TieringOffloadingMetrics.READ_BYTES][label] == 80 + assert values[TieringOffloadingMetrics.READ_BYTES][label] == 16 assert values[TieringOffloadingMetrics.READ_TIME][label] == 0.25 assert values[TieringOffloadingMetrics.CASCADE_JOB_FAILURES][label] == 1 assert values[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES][label] == 1 @@ -349,8 +351,8 @@ def test_tiering_manager_reports_active_job_and_primary_usage_gauges(): stats = manager.get_stats() assert stats is not None values = stats.data["data"] - fs_label = ("0:fs",) - p2p_label = ("1:p2p",) + fs_label = ("1:fs",) + p2p_label = ("2:p2p",) assert values[TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC][ fs_label ] == pytest.approx(2 / 6) @@ -1229,13 +1231,13 @@ def test_reset_cache_clears_orchestrator_state(self, manager_setup): primary tier; pending submissions are dropped without being sent to the secondary tier. Active request state is retained.""" # Cascade — populates primary blocks and leaves cascade jobs - # in _transfer_jobs (the synchronous example tier has already + # in _jobs (the synchronous example tier has already # queued completions); reset_cache's drain loop will pick them up. blocks = to_keys(range(3)) self._start_request() self.manager.prepare_store(blocks, _CTX) self.manager.complete_store(blocks, _CTX, success=True) - assert self.manager._transfer_jobs + assert self.manager._jobs # Pending promotion submission (deferred — no on_schedule_end after # the lookup that staged it). @@ -1268,7 +1270,7 @@ def test_reset_cache_clears_orchestrator_state(self, manager_setup): self.manager.reset_cache() # Orchestrator state cleared. - assert self.manager._transfer_jobs == {} + assert self.manager._jobs == {} assert self.manager._pending_load_submissions == {} assert set(self.manager._req_state) == {_CTX.req_id, rl_ctx.req_id} assert self.manager._processed_jobs_this_step is False @@ -1296,18 +1298,18 @@ def test_reset_cache_drains_all_tiers(self, manager_setup): wraps=self.secondary_tier2.drain_jobs ) - # Drive a cascade so a job lands in _transfer_jobs. + # Drive a cascade so a job lands in _jobs. blocks = to_keys(range(3)) self._start_request() self.manager.prepare_store(blocks, _CTX) self.manager.complete_store(blocks, _CTX, success=True) - assert self.manager._transfer_jobs + assert self.manager._jobs self.manager.reset_cache() self.secondary_tier1.drain_jobs.assert_called_once() self.secondary_tier2.drain_jobs.assert_called_once() - assert self.manager._transfer_jobs == {} + assert self.manager._jobs == {} @pytest.mark.parametrize( "load_tier_filter", diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index a760a4ab3628..768d7e328ad7 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -6,7 +6,7 @@ from abc import ABC, abstractmethod from collections.abc import Collection, Iterable -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar import numpy as np @@ -57,7 +57,7 @@ class TieringOffloadingMetrics: @dataclass -class JobMetadata: +class TransferJob: """Metadata for an in-flight async transfer job.""" job_id: JobId @@ -78,8 +78,7 @@ class JobResult: # keys that were successfully loaded. None means all keys share the fate # indicated by `success`. Must be a subset of the job's original keys. successful_keys: Collection[OffloadKey] | None = None - transfer_size: int | None = field(default=None, compare=False) - transfer_time: float | None = field(default=None, compare=False) + transfer_time: float | None = None class ParentManager(ABC): @@ -113,7 +112,7 @@ def create_store_job( self, keys: Collection[OffloadKey], req_context: ReqContext, - ) -> JobMetadata: ... + ) -> TransferJob: ... @abstractmethod def on_request_finished(self, req_context: ReqContext) -> None: ... @@ -170,7 +169,7 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: pass @abstractmethod - def submit_store(self, job_metadata: JobMetadata) -> None: + def submit_store(self, job_metadata: TransferJob) -> None: """ Submit an async job to store blocks from the primary tier to this secondary tier. @@ -198,7 +197,7 @@ def submit_store(self, job_metadata: JobMetadata) -> None: pass @abstractmethod - def submit_load(self, job_metadata: JobMetadata) -> None: + def submit_load(self, job_metadata: TransferJob) -> None: """ Submit an async job to load blocks from this secondary tier to the primary tier. diff --git a/vllm/v1/kv_offload/tiering/example/manager.py b/vllm/v1/kv_offload/tiering/example/manager.py index dd45abcc1996..b2652a6d555e 100644 --- a/vllm/v1/kv_offload/tiering/example/manager.py +++ b/vllm/v1/kv_offload/tiering/example/manager.py @@ -23,9 +23,9 @@ RequestOffloadingContext, ) from vllm.v1.kv_offload.tiering.base import ( - JobMetadata, JobResult, SecondaryTierManager, + TransferJob, ) logger = logging.getLogger(__name__) @@ -91,7 +91,7 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: return LookupResult.HIT if key in self.blocks else LookupResult.MISS @override - def submit_store(self, job_metadata: JobMetadata) -> None: + def submit_store(self, job_metadata: TransferJob) -> None: """ Submit a job to store blocks from primary tier to this tier. @@ -112,13 +112,12 @@ def submit_store(self, job_metadata: JobMetadata) -> None: JobResult( job_id=job_metadata.job_id, success=True, - transfer_size=len(keys) * self._block_size, transfer_time=0.0, ) ) @override - def submit_load(self, job_metadata: JobMetadata) -> None: + def submit_load(self, job_metadata: TransferJob) -> None: """ Submit a job to load blocks from this tier to primary tier. @@ -144,7 +143,6 @@ def submit_load(self, job_metadata: JobMetadata) -> None: JobResult( job_id=job_metadata.job_id, success=True, - transfer_size=len(keys) * self._block_size, transfer_time=0.0, ) ) diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index a8fb83f717a8..0ae07fcd66a2 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -18,7 +18,6 @@ import functools import json import os -import time from collections.abc import Iterable from typing import TYPE_CHECKING, ClassVar @@ -44,11 +43,11 @@ from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager from vllm.v1.kv_offload.tiering.base import ( JobId, - JobMetadata, JobResult, RequestOffloadingContext, ScheduleEndContext, SecondaryTierManager, + TransferJob, ) from vllm.v1.kv_offload.tiering.fs.io import ( batch_load_block, @@ -202,8 +201,7 @@ def __init__( ) self._lookup_manager = FsAsyncLookupManager(tier=self, tier_type=self.tier_type) - self._job_transfer_sizes: dict[int, int] = {} - self._job_submitted_at: dict[int, float] = {} + self._job_transfer_sizes: dict[JobId, int] = {} @override def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: @@ -217,7 +215,7 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: return LookupResult.HIT if result else LookupResult.MISS @override - def submit_store(self, job_metadata: JobMetadata) -> None: + def submit_store(self, job_metadata: TransferJob) -> None: keys = list(job_metadata.keys) if self.events is not None: self._store_job_keys[job_metadata.job_id] = keys @@ -230,11 +228,10 @@ def submit_store(self, job_metadata: JobMetadata) -> None: self._use_o_direct, ) self._job_transfer_sizes[job_metadata.job_id] = len(keys) * self._block_size - self._job_submitted_at[job_metadata.job_id] = time.monotonic() self._pool.enqueue_store(job_metadata.job_id, 1, [task]) @override - def submit_load(self, job_metadata: JobMetadata) -> None: + def submit_load(self, job_metadata: TransferJob) -> None: job_id = job_metadata.job_id # Track this load's keys so a failed promotion can mark only its failed # keys as a miss (see get_finished_jobs). @@ -271,7 +268,6 @@ def load_task() -> None: raise self._job_transfer_sizes[job_id] = len(keys) * self._block_size - self._job_submitted_at[job_id] = time.monotonic() self._pool.enqueue_load(job_id, 1, [load_task]) @override @@ -279,8 +275,7 @@ def get_finished_jobs(self) -> Iterable[JobResult]: """Collect finished jobs; a failed promotion marks only its failed keys as a miss here (scheduler thread).""" results = [] - now = time.monotonic() - for job_id, success in self._pool.get_finished(): + for job_id, success, transfer_time in self._pool.get_finished(): if self.events is not None: keys = self._store_job_keys.pop(job_id, None) if success and keys: @@ -292,8 +287,6 @@ def get_finished_jobs(self) -> Iterable[JobResult]: locality=self.locality, ) ) - submitted_at = self._job_submitted_at.pop(job_id, None) - transfer_time = None if submitted_at is None else now - submitted_at transfer_size = self._job_transfer_sizes.pop(job_id, None) load_keys = self._load_job_keys.pop(job_id, None) num_succeeded = self._load_progress.pop(job_id, 0) @@ -319,15 +312,15 @@ def get_finished_jobs(self) -> Iterable[JobResult]: transfer_time=transfer_time, ) ) - else: - results.append( - JobResult( - job_id=job_id, - success=success, - transfer_size=transfer_size, - transfer_time=transfer_time, - ) + continue + results.append( + JobResult( + job_id=job_id, + success=success, + transfer_size=transfer_size, + transfer_time=transfer_time, ) + ) return results @override diff --git a/vllm/v1/kv_offload/tiering/fs/thread_pool.py b/vllm/v1/kv_offload/tiering/fs/thread_pool.py index 9bf8fe508f0d..9db0629192bd 100644 --- a/vllm/v1/kv_offload/tiering/fs/thread_pool.py +++ b/vllm/v1/kv_offload/tiering/fs/thread_pool.py @@ -9,6 +9,7 @@ """ import threading +import time from collections import deque from collections.abc import Callable, Iterable @@ -25,26 +26,37 @@ class JobState: Each task calls task_done(success) when it finishes. """ - __slots__ = ("_job_id", "_n_tasks", "_completed", "_success", "_lock") + __slots__ = ( + "_job_id", + "_n_tasks", + "_completed", + "_success", + "_transfer_time", + "_lock", + ) def __init__(self, job_id: JobId, n_tasks: int) -> None: self._job_id: JobId = job_id self._n_tasks = n_tasks self._completed = 0 self._success = True + self._transfer_time = 0.0 self._lock = threading.Lock() @property def job_id(self) -> JobId: return self._job_id - def task_done(self, success: bool) -> tuple[bool, bool]: + def task_done( + self, success: bool, transfer_time: float + ) -> tuple[bool, bool, float]: """Returns if job completed and success flag""" with self._lock: self._completed += 1 + self._transfer_time += transfer_time if not success: self._success = False - return self._completed == self._n_tasks, self._success + return self._completed == self._n_tasks, self._success, self._transfer_time class DualQueueThreadPool: @@ -67,7 +79,7 @@ def __init__( self._condition = threading.Condition(threading.Lock()) self._stop = False self._threads: list[threading.Thread] = [] - self._finished_q: deque[tuple[JobId, bool]] = deque() + self._finished_q: deque[tuple[JobId, bool, float]] = deque() self._inflight_jobs = 0 # guarded by _condition for i in range(n_read_threads): @@ -118,7 +130,7 @@ def enqueue_store( self._store_q.append((fn, state)) self._condition.notify(n_tasks) - def get_finished(self) -> list[tuple[JobId, bool]]: + def get_finished(self) -> list[tuple[JobId, bool, float]]: # No lock needed: deque is thread-safe for concurrent append/popleft, # and the manager is the sole popper. jobs = [] @@ -163,18 +175,23 @@ def _worker(self, load_priority: bool) -> None: secondary = self._store_q if load_priority else self._load_q task, state = primary.popleft() if primary else secondary.popleft() try: + start_time = time.monotonic() task() - job_finished, success = state.task_done(True) + transfer_time = time.monotonic() - start_time + job_finished, success, total_time = state.task_done(True, transfer_time) except Exception as exc: + transfer_time = time.monotonic() - start_time logger.error( "Job %s block I/O failed: %s", state.job_id, exc, ) - job_finished, success = state.task_done(False) + job_finished, success, total_time = state.task_done( + False, transfer_time + ) if job_finished: with self._condition: - self._finished_q.append((state.job_id, success)) + self._finished_q.append((state.job_id, success, total_time)) self._inflight_jobs -= 1 self._condition.notify_all() diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 9b24522537cf..eb9b9c8a5d80 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -20,9 +20,11 @@ protecting blocks from eviction until complete_read() is called """ +import functools import time from collections.abc import Collection, Iterable, Sequence from dataclasses import dataclass, field +from typing import NamedTuple import numpy as np from typing_extensions import override @@ -48,15 +50,18 @@ from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.kv_offload.tiering.base import ( JobId, - JobMetadata, JobResult, ParentManager, SecondaryTierManager, TieringOffloadingMetrics, + TransferJob, ) logger = init_logger(__name__) +TierLabel = tuple[str] +_PRIMARY_TIER_LABEL: TierLabel = ("0:primary",) + @dataclass class PendingPromotion: @@ -73,12 +78,28 @@ class RequestState: pending_primary_stores: int = 0 is_finished: bool = False request_level_tiers: set[SecondaryTierManager] | None = None + observed_lookups: dict[TierLabel, set[OffloadKey]] | None = field( + default_factory=dict + ) sync_lookup_delay: float = 0.0 # time.monotonic() of this request's first deferred secondary-tier lookup; # None once consumed (observed) or while no secondary lookup is pending. secondary_lookup_start_time: float | None = None +@dataclass +class _TierState: + active_promotion_count: int = 0 + active_cascade_count: int = 0 + primary_write_block_ids: set[int] = field(default_factory=set) + primary_read_block_ids: set[int] = field(default_factory=set) + + +class _JobMetadata(NamedTuple): + transfer_job: TransferJob + tier_idx: int + + class CPUPrimaryTierOffloadingManager(CPUOffloadingManager): """CPUOffloadingManager with a primary/secondary transfer interface. @@ -134,15 +155,17 @@ class _SecondaryTierFacingParent(ParentManager): """Wrapper that implements ParentManager by delegating to the TieringOffloadingManager with exclude_tier set to the origin tier.""" - __slots__ = ("_m", "_origin") + __slots__ = ("_m", "_origin", "_origin_idx") def __init__( self, manager: "TieringOffloadingManager", tier: SecondaryTierManager, + tier_idx: int, ): self._m = manager self._origin = tier + self._origin_idx = tier_idx def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: return self._m.on_new_request(req_context, exclude_tier=self._origin) @@ -152,8 +175,8 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: def create_store_job( self, keys: Collection[OffloadKey], req_context: ReqContext - ) -> JobMetadata: - return self._m.create_store_job(keys, req_context, self._origin) + ) -> TransferJob: + return self._m.create_store_job(keys, req_context, self._origin_idx) def on_request_finished(self, req_context: ReqContext) -> None: return self._m.on_request_finished(req_context, exclude_tier=self._origin) @@ -193,11 +216,19 @@ def __init__( self._job_id_counter: int = 0 # Job tracking: maps job_id to metadata for all in-flight transfers. - # JobMetadata.is_promotion distinguishes direction: + # TransferJob.is_promotion distinguishes direction: # True: secondary → primary (promotion) # False: primary → secondary (cascade) - self._transfer_jobs: dict[JobId, JobMetadata] = {} - self._transfer_job_tiers: dict[JobId, SecondaryTierManager] = {} + self._jobs: dict[JobId, _JobMetadata] = {} + self._tier_states: list[_TierState] = [ + _TierState() for _ in self.secondary_tiers + ] + self._tier_indices = { + tier: tier_idx for tier_idx, tier in enumerate(self.secondary_tiers) + } + primary_view = self.primary_tier.get_kv_memoryview() + assert primary_view.strides is not None + self._primary_block_size = primary_view.strides[0] # Pending promotion requests accumulated during lookup() calls; flushed # as one batched submit_load() per (tier, request) in on_schedule_end(). @@ -218,8 +249,8 @@ def __init__( # Cached ParentManager wrappers for each secondary tier. self._tier_parents: dict[SecondaryTierManager, _SecondaryTierFacingParent] = { - tier: _SecondaryTierFacingParent(self, tier) - for tier in self.secondary_tiers + tier: _SecondaryTierFacingParent(self, tier, tier_idx) + for tier_idx, tier in enumerate(self.secondary_tiers) } # Buffers manager-level observations (e.g. lookup delay) between # get_stats() calls; merged in and reset each time get_stats() runs. @@ -231,22 +262,48 @@ def _next_job_id(self) -> JobId: self._job_id_counter += 1 return job_id - @staticmethod - def _tier_label(tier_idx: int, tier: SecondaryTierManager) -> tuple[str]: - return (f"{tier_idx}:{tier.tier_type}",) + @functools.cache # noqa: B019 + def _tier_label(self, tier_idx: int) -> TierLabel: + return (f"{tier_idx + 1}:{self.secondary_tiers[tier_idx].tier_type}",) + + def _register_job(self, transfer_job: TransferJob, tier_idx: int) -> None: + self._jobs[transfer_job.job_id] = _JobMetadata(transfer_job, tier_idx) + state = self._tier_states[tier_idx] + block_ids = {int(block_id) for block_id in transfer_job.block_ids} + if transfer_job.is_promotion: + state.active_promotion_count += 1 + state.primary_write_block_ids.update(block_ids) + else: + state.active_cascade_count += 1 + state.primary_read_block_ids.update(block_ids) + + def _pop_job(self, job_id: JobId) -> _JobMetadata | None: + job_metadata = self._jobs.pop(job_id, None) + if job_metadata is None: + return None + + transfer_job = job_metadata.transfer_job + state = self._tier_states[job_metadata.tier_idx] + block_ids = {int(block_id) for block_id in transfer_job.block_ids} + if transfer_job.is_promotion: + state.active_promotion_count -= 1 + state.primary_write_block_ids.difference_update(block_ids) + else: + state.active_cascade_count -= 1 + state.primary_read_block_ids.difference_update(block_ids) + return job_metadata def _record_finished_job_stats( self, - tier_idx: int, - tier: SecondaryTierManager, - job_metadata: JobMetadata, + job_metadata: _JobMetadata, completed_job: JobResult, ) -> None: - labelvalues = self._tier_label(tier_idx, tier) + transfer_job = job_metadata.transfer_job + labelvalues = self._tier_label(job_metadata.tier_idx) if not completed_job.success: failure_metric = ( TieringOffloadingMetrics.PROMOTION_JOB_FAILURES - if job_metadata.is_promotion + if transfer_job.is_promotion else TieringOffloadingMetrics.CASCADE_JOB_FAILURES ) self._stats.increase_counter(failure_metric, labelvalues=labelvalues) @@ -254,52 +311,27 @@ def _record_finished_job_stats( bytes_metric = ( TieringOffloadingMetrics.READ_BYTES - if job_metadata.is_promotion + if transfer_job.is_promotion else TieringOffloadingMetrics.WRITE_BYTES ) time_metric = ( TieringOffloadingMetrics.READ_TIME - if job_metadata.is_promotion + if transfer_job.is_promotion else TieringOffloadingMetrics.WRITE_TIME ) - if completed_job.transfer_size is not None: - self._stats.increase_counter( - bytes_metric, completed_job.transfer_size, labelvalues - ) + transfer_size = len(transfer_job.keys) * self._primary_block_size + self._stats.increase_counter(bytes_metric, transfer_size, labelvalues) if completed_job.transfer_time is not None: self._stats.increase_counter( time_metric, completed_job.transfer_time, labelvalues ) def _record_active_transfer_stats(self, stats: OffloadingConnectorStats) -> None: - active_promotion_jobs: dict[SecondaryTierManager, int] = {} - active_cascade_jobs: dict[SecondaryTierManager, int] = {} - primary_write_blocks: dict[SecondaryTierManager, set[int]] = {} - primary_read_blocks: dict[SecondaryTierManager, set[int]] = {} - - for job_id, job_metadata in self._transfer_jobs.items(): - tier = self._transfer_job_tiers.get(job_id) - if tier is None: - continue - block_ids = set(job_metadata.block_ids.tolist()) - if job_metadata.is_promotion: - active_promotion_jobs[tier] = active_promotion_jobs.get(tier, 0) + 1 - primary_write_blocks.setdefault(tier, set()).update(block_ids) - else: - active_cascade_jobs[tier] = active_cascade_jobs.get(tier, 0) + 1 - primary_read_blocks.setdefault(tier, set()).update(block_ids) - - for pending_tier, pending_by_ctx in self._pending_load_submissions.items(): - for pending in pending_by_ctx.values(): - primary_write_blocks.setdefault(pending_tier, set()).update( - pending.block_ids - ) - num_primary_blocks = self.primary_tier._num_blocks - for tier_idx, tier in enumerate(self.secondary_tiers): - labelvalues = self._tier_label(tier_idx, tier) - write_blocks = len(primary_write_blocks.get(tier, set())) - read_blocks = len(primary_read_blocks.get(tier, set())) + for tier_idx, state in enumerate(self._tier_states): + labelvalues = self._tier_label(tier_idx) + write_blocks = len(state.primary_write_block_ids) + read_blocks = len(state.primary_read_block_ids) write_usage = ( write_blocks / num_primary_blocks if num_primary_blocks > 0 else 0.0 ) @@ -318,15 +350,38 @@ def _record_active_transfer_stats(self, stats: OffloadingConnectorStats) -> None ) stats.set_gauge( TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS, - active_promotion_jobs.get(tier, 0), + state.active_promotion_count, labelvalues, ) stats.set_gauge( TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS, - active_cascade_jobs.get(tier, 0), + state.active_cascade_count, labelvalues, ) + def _observe_lookup( + self, + req_state: RequestState | None, + labelvalues: TierLabel, + key: OffloadKey, + result: LookupResult, + ) -> None: + if req_state is None or req_state.observed_lookups is None: + return + observed = req_state.observed_lookups.setdefault(labelvalues, set()) + if key in observed: + return + observed.add(key) + self._stats.increase_counter( + TieringOffloadingMetrics.BLOCK_QUERIES, + labelvalues=labelvalues, + ) + if result is LookupResult.HIT: + self._stats.increase_counter( + TieringOffloadingMetrics.BLOCK_HITS, + labelvalues=labelvalues, + ) + def _maybe_process_finished_jobs(self): """ Poll secondary tiers for completed jobs (at most once per step). @@ -387,15 +442,19 @@ def _process_finished_jobs(self): for i, tier in enumerate(self.secondary_tiers): for completed_job in tier.get_finished_jobs(): job_id = completed_job.job_id - job_metadata = self._transfer_jobs.pop(job_id, None) - self._transfer_job_tiers.pop(job_id, None) + job_metadata = self._pop_job(job_id) assert job_metadata is not None, ( f"Finished job_id {job_id} from tier #{i}" - f" ({tier.tier_type}) not in _transfer_jobs" + f" ({tier.tier_type}) not in _jobs" ) - self._record_finished_job_stats(i, tier, job_metadata, completed_job) + assert job_metadata.tier_idx == i, ( + f"Finished job_id {job_id} reported by tier #{i}" + f" but belongs to tier #{job_metadata.tier_idx}" + ) + transfer_job = job_metadata.transfer_job + self._record_finished_job_stats(job_metadata, completed_job) - if job_metadata.is_promotion: + if transfer_job.is_promotion: # secondary→primary transfer (promotion) completed. # Make blocks available in primary tier. self._complete_promotion(job_metadata, completed_job) @@ -403,7 +462,7 @@ def _process_finished_jobs(self): # primary→secondary transfer completed. # Decrement ref_cnt on primary blocks. self.primary_tier.complete_read( - job_metadata.keys, job_metadata.req_context + transfer_job.keys, transfer_job.req_context ) @override @@ -444,6 +503,8 @@ def lookup( req_state = self._req_state.get(req_context.req_id) primary_hit = self.primary_tier.lookup(key, req_context) + if primary_hit in (LookupResult.HIT, LookupResult.MISS): + self._observe_lookup(req_state, _PRIMARY_TIER_LABEL, key, primary_hit) if primary_hit is LookupResult.HIT: return LookupResult.HIT if primary_hit is LookupResult.HIT_PENDING: @@ -456,17 +517,11 @@ def lookup( continue if not req_context.load_tier_filter.allows(tier.medium, tier.locality): continue - labelvalues = self._tier_label(i, tier) - self._stats.increase_counter( - TieringOffloadingMetrics.BLOCK_QUERIES, - labelvalues=labelvalues, - ) + labelvalues = self._tier_label(i) result = tier.lookup(key, req_context) + if result in (LookupResult.HIT, LookupResult.MISS): + self._observe_lookup(req_state, labelvalues, key, result) if result is LookupResult.HIT: - self._stats.increase_counter( - TieringOffloadingMetrics.BLOCK_HITS, - labelvalues=labelvalues, - ) promoted = self._initiate_promotion(tier, key, req_context) self._accumulate_lookup_sync_delay(req_state, lookup_start) if ( @@ -576,17 +631,17 @@ def _flush_pending_promotions(self) -> None: return for tier, pending_by_ctx in self._pending_load_submissions.items(): + tier_idx = self._tier_indices[tier] for entry in pending_by_ctx.values(): job_id = self._next_job_id() - job_metadata = JobMetadata( + job_metadata = TransferJob( job_id=job_id, keys=entry.keys, block_ids=np.array(entry.block_ids, dtype=np.int64), is_promotion=True, req_context=entry.req_context, ) - self._transfer_jobs[job_id] = job_metadata - self._transfer_job_tiers[job_id] = tier + self._register_job(job_metadata, tier_idx) tier.submit_load(job_metadata) self._pending_load_submissions.clear() @@ -721,7 +776,9 @@ def _cascade_existing_blocks_to_request_level_tiers( return for tier in request_level_tiers: - job_metadata = self.create_store_job(ready_keys, req_context, tier) + job_metadata = self.create_store_job( + ready_keys, req_context, self._tier_indices[tier] + ) tier.submit_store(job_metadata) @override @@ -759,7 +816,9 @@ def complete_store( # eviction during the async transfer). One prepare_read() call per # secondary tier. for tier in self.secondary_tiers: - job_metadata = self.create_store_job(keys, req_context, tier) + job_metadata = self.create_store_job( + keys, req_context, self._tier_indices[tier] + ) tier.submit_store(job_metadata) # Note: The async transfers are now in flight. Their completion is @@ -774,13 +833,13 @@ def create_store_job( self, keys: Collection[OffloadKey], req_context: ReqContext, - tier: SecondaryTierManager | None = None, - ) -> JobMetadata: + tier_idx: int | None = None, + ) -> TransferJob: """Pin blocks in the primary tier and create a tracked store job. Calls prepare_read() to increment ref_cnt (protecting blocks from eviction during the async transfer), allocates a job ID, - and registers the job in _transfer_jobs. + and registers the job in _jobs. The caller is responsible for the actual data transfer and reporting completion via get_finished_jobs(). @@ -788,16 +847,15 @@ def create_store_job( primary_blocks_spec = self.primary_tier.prepare_read(keys, req_context) assert isinstance(primary_blocks_spec, CPULoadStoreSpec) job_id = self._next_job_id() - job_metadata = JobMetadata( + job_metadata = TransferJob( job_id=job_id, keys=keys, block_ids=primary_blocks_spec.block_ids, is_promotion=False, req_context=req_context, ) - self._transfer_jobs[job_id] = job_metadata - if tier is not None: - self._transfer_job_tiers[job_id] = tier + if tier_idx is not None: + self._register_job(job_metadata, tier_idx) return job_metadata @override @@ -896,6 +954,7 @@ def on_schedule_end(self, context: ScheduleEndContext) -> None: state = self._req_state.get(req_id) if state is None: continue + state.observed_lookups = None self._maybe_observe_lookup_sync_delay(state) self._maybe_observe_lookup_async_delay(state) @@ -904,7 +963,7 @@ def has_pending_work(self) -> bool: # In-flight primary<->secondary transfers (pending promotions are # translated to transfer jobs in on_schedule_end), plus any work the # secondary tiers themselves still have outstanding. - return bool(self._transfer_jobs) or any( + return bool(self._jobs) or any( tier.has_pending_work() for tier in self.secondary_tiers ) @@ -944,7 +1003,6 @@ def reset_cache(self) -> None: # reset below invalidates; their submit_load() has not yet been # called so no tier I/O is touching that memory. self._pending_load_submissions.clear() - self._transfer_job_tiers.clear() finished_req_ids = [] for req_id, state in self._req_state.items(): diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index 56de3e53f389..fa055c2f0ad7 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -22,11 +22,11 @@ from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager from vllm.v1.kv_offload.tiering.base import ( JobId, - JobMetadata, JobResult, RequestOffloadingContext, ScheduleEndContext, SecondaryTierManager, + TransferJob, ) from vllm.v1.kv_offload.tiering.obj.config import ObjStoreConfig @@ -58,8 +58,6 @@ class TransferEntry(NamedTuple): xfer_handle: "nixl_xfer_handle" files_desc: object obj_handle: "nixl_prepped_dlist_handle" - transfer_size: int - submitted_at: float class ObjAsyncLookupManager(AsyncLookupManager): @@ -273,8 +271,6 @@ def _submit_transfer( xfer_handle=xfer_handle, files_desc=files_desc, obj_handle=obj_handle, - transfer_size=len(block_ids_list) * self._block_size_bytes, - submitted_at=time.monotonic(), ) def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: @@ -283,7 +279,7 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: return LookupResult.RETRY return LookupResult.HIT if result else LookupResult.MISS - def submit_store(self, job_metadata: JobMetadata) -> None: + def submit_store(self, job_metadata: TransferJob) -> None: if self.events is not None: self._store_job_keys[job_metadata.job_id] = list(job_metadata.keys) obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys) @@ -291,7 +287,7 @@ def submit_store(self, job_metadata: JobMetadata) -> None: job_metadata.job_id, job_metadata.block_ids, obj_keys, NIXL_WRITE ) - def submit_load(self, job_metadata: JobMetadata) -> None: + def submit_load(self, job_metadata: TransferJob) -> None: self._load_job_keys[job_metadata.job_id] = list(job_metadata.keys) obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys) self._submit_transfer( @@ -325,6 +321,11 @@ def _poll_active_transfers(self) -> None: success = False logger.warning("transfer failed job=%d state=%s", job_id, state) + transfer_time = None + if success: + telemetry = self._agent.get_xfer_telemetry(entry.xfer_handle) + transfer_time = telemetry.xferDuration / 1e6 + try: self._agent.release_xfer_handle(entry.xfer_handle) except Exception as exc: @@ -353,8 +354,7 @@ def _poll_active_transfers(self) -> None: JobResult( job_id=job_id, success=success, - transfer_size=entry.transfer_size, - transfer_time=time.monotonic() - entry.submitted_at, + transfer_time=transfer_time, ) ) diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 655a9e83b55e..54b4553b4049 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -28,9 +28,9 @@ ) from vllm.v1.kv_offload.file_mapper import FileMapper from vllm.v1.kv_offload.tiering.base import ( - JobMetadata, JobResult, SecondaryTierManager, + TransferJob, ) from vllm.v1.kv_offload.tiering.p2p.control import ControlTransport, ZmqTransport from vllm.v1.kv_offload.tiering.p2p.data import DataTransport, NixlTransport @@ -411,7 +411,7 @@ def on_request_finished(self, req_context: ReqContext) -> None: return @override - def submit_store(self, job_metadata: JobMetadata) -> None: + def submit_store(self, job_metadata: TransferJob) -> None: job_id = job_metadata.job_id keys = list(job_metadata.keys) block_ids = job_metadata.block_ids @@ -444,8 +444,6 @@ def submit_store(self, job_metadata: JobMetadata) -> None: self._finished_jobs.append(JobResult(job_id=job_id, success=False)) return - - # Fast path: a session has already received FetchMsg for this id, # so we can route the batch straight into its ServerRole. session = self._kv_to_session.get(kv_request_id) @@ -472,7 +470,7 @@ def submit_store(self, job_metadata: JobMetadata) -> None: ) @override - def submit_load(self, job_metadata: JobMetadata) -> None: + def submit_load(self, job_metadata: TransferJob) -> None: job_id = job_metadata.job_id keys = list(job_metadata.keys) block_ids = job_metadata.block_ids diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 16ec9a3f85bb..55b81112ce56 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -146,46 +146,60 @@ def build_metric_definitions( ), ) ) - for metric_name, documentation in ( - ( - TieringOffloadingMetrics.READ_BYTES, - "Total bytes read from secondary tiers into the primary tier.", + metrics[TieringOffloadingMetrics.READ_BYTES] = OffloadingCounterMetadata( + documentation=( + "Total bytes read from secondary tiers into the primary tier, " + "labeled by tier." ), - ( - TieringOffloadingMetrics.READ_TIME, + labelnames=("tier",), + ) + metrics[TieringOffloadingMetrics.READ_TIME] = OffloadingCounterMetadata( + documentation=( "Total time spent reading from secondary tiers into the primary " - "tier, in seconds.", + "tier, in seconds, labeled by tier." ), - ( - TieringOffloadingMetrics.WRITE_BYTES, - "Total bytes written from the primary tier to secondary tiers.", + labelnames=("tier",), + ) + metrics[TieringOffloadingMetrics.WRITE_BYTES] = OffloadingCounterMetadata( + documentation=( + "Total bytes written from the primary tier to secondary tiers, " + "labeled by tier." ), - ( - TieringOffloadingMetrics.WRITE_TIME, + labelnames=("tier",), + ) + metrics[TieringOffloadingMetrics.WRITE_TIME] = OffloadingCounterMetadata( + documentation=( "Total time spent writing from the primary tier to secondary " - "tiers, in seconds.", - ), - ( - TieringOffloadingMetrics.PROMOTION_JOB_FAILURES, - "Number of failed secondary-tier promotion jobs.", - ), - ( - TieringOffloadingMetrics.CASCADE_JOB_FAILURES, - "Number of failed secondary-tier cascade jobs.", + "tiers, in seconds, labeled by tier." ), - ( - TieringOffloadingMetrics.BLOCK_QUERIES, - "Number of block lookup queries sent to secondary tiers.", - ), - ( - TieringOffloadingMetrics.BLOCK_HITS, - "Number of block lookup hits in secondary tiers.", - ), - ): - metrics[metric_name] = OffloadingCounterMetadata( - documentation=documentation, + labelnames=("tier",), + ) + metrics[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES] = ( + OffloadingCounterMetadata( + documentation=( + "Number of failed secondary-tier promotion jobs, labeled by tier." + ), + labelnames=("tier",), + ) + ) + metrics[TieringOffloadingMetrics.CASCADE_JOB_FAILURES] = ( + OffloadingCounterMetadata( + documentation=( + "Number of failed secondary-tier cascade jobs, labeled by tier." + ), labelnames=("tier",), ) + ) + metrics[TieringOffloadingMetrics.BLOCK_QUERIES] = OffloadingCounterMetadata( + documentation=( + "Number of block lookup queries sent to a tier, labeled by tier." + ), + labelnames=("tier",), + ) + metrics[TieringOffloadingMetrics.BLOCK_HITS] = OffloadingCounterMetadata( + documentation="Number of block lookup hits in a tier, labeled by tier.", + labelnames=("tier",), + ) metrics[TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES] = ( OffloadingCounterMetadata( documentation=( @@ -194,30 +208,38 @@ def build_metric_definitions( ), ) ) - for metric_name, documentation in ( - ( - TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC, - "Current fraction of primary-tier space used by writes from " - "secondary tiers.", - ), - ( - TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC, - "Current fraction of primary-tier space used by reads to " - "secondary tiers.", - ), - ( - TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS, - "Number of active secondary-tier promotion jobs.", - ), - ( - TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS, - "Number of active secondary-tier cascade jobs.", - ), - ): - metrics[metric_name] = OffloadingGaugeMetadata( - documentation=documentation, + metrics[TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC] = ( + OffloadingGaugeMetadata( + documentation=( + "Current fraction of primary-tier space used by writes from " + "secondary tiers, labeled by tier." + ), + labelnames=("tier",), + ) + ) + metrics[TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC] = ( + OffloadingGaugeMetadata( + documentation=( + "Current fraction of primary-tier space used by reads to " + "secondary tiers, labeled by tier." + ), + labelnames=("tier",), + ) + ) + metrics[TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS] = ( + OffloadingGaugeMetadata( + documentation=( + "Number of active secondary-tier promotion jobs, labeled by tier." + ), labelnames=("tier",), ) + ) + metrics[TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS] = OffloadingGaugeMetadata( + documentation=( + "Number of active secondary-tier cascade jobs, labeled by tier." + ), + labelnames=("tier",), + ) secondary_tier_configs = extra_config.get("secondary_tiers", []) if not isinstance(secondary_tier_configs, list): raise ValueError("secondary_tiers must be a list of tier configurations") From 19b2e7c88251646a837af7f9aa509419aa5bb073 Mon Sep 17 00:00:00 2001 From: srinivas_oo7 Date: Tue, 28 Jul 2026 22:02:46 -0500 Subject: [PATCH 04/13] updated code Signed-off-by: srinivas_oo7 --- .../tiering/test_tiering_offloading.py | 39 ++++++++++++++++ vllm/v1/kv_offload/tiering/manager.py | 44 ++++++++++++------- vllm/v1/kv_offload/tiering/p2p/manager.py | 17 ++++--- 3 files changed, 80 insertions(+), 20 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index c70c12d0c26f..12e3caf9852d 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -242,6 +242,45 @@ def test_tiering_manager_records_secondary_lookup_metrics(): assert ("1:fs",) not in values[TieringOffloadingMetrics.BLOCK_HITS] +def test_tiering_manager_stops_lookup_metrics_after_allocation(): + mock_region = _mock_mmap_region(5) + primary_tier = CPUPrimaryTierOffloadingManager( + num_blocks=5, mmap_region=mock_region + ) + tier = MetricsSecondaryTierManager( + offloading_spec=_MOCK_OFFLOADING_SPEC, + primary_kv_view=mock_region.create_kv_memoryview(), + tier_type="fs", + ) + tier.lookup_result = LookupResult.HIT + manager = TieringOffloadingManager( + primary_tier=primary_tier, + secondary_tiers=[tier], + ) + manager.on_new_request(_CTX) + + key = to_keys([1])[0] + assert manager.lookup(key, _CTX) is LookupResult.RETRY + + stats = manager.get_stats() + assert stats is not None + values = stats.data["data"] + assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("0:primary",)] == 1 + assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("1:fs",)] == 1 + assert values[TieringOffloadingMetrics.BLOCK_HITS][("1:fs",)] == 1 + + manager.on_schedule_end( + ScheduleEndContext(new_req_ids=(_CTX.req_id,), preempted_req_ids=()) + ) + assert manager.lookup(key, _CTX) is LookupResult.HIT_PENDING + + stats = manager.get_stats() + assert stats is not None + values = stats.data["data"] + assert TieringOffloadingMetrics.BLOCK_QUERIES not in values + assert TieringOffloadingMetrics.BLOCK_HITS not in values + + def test_tiering_manager_records_finished_job_metrics(): mock_region = _mock_mmap_region(5) primary_tier = CPUPrimaryTierOffloadingManager( diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index eb9b9c8a5d80..0806cb65c7a0 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -91,8 +91,8 @@ class RequestState: class _TierState: active_promotion_count: int = 0 active_cascade_count: int = 0 - primary_write_block_ids: set[int] = field(default_factory=set) - primary_read_block_ids: set[int] = field(default_factory=set) + primary_write_block_count: int = 0 + primary_read_block_count: int = 0 class _JobMetadata(NamedTuple): @@ -252,6 +252,7 @@ def __init__( tier: _SecondaryTierFacingParent(self, tier, tier_idx) for tier_idx, tier in enumerate(self.secondary_tiers) } + # Buffers manager-level observations (e.g. lookup delay) between # get_stats() calls; merged in and reset each time get_stats() runs. self._stats = OffloadingConnectorStats() @@ -269,13 +270,13 @@ def _tier_label(self, tier_idx: int) -> TierLabel: def _register_job(self, transfer_job: TransferJob, tier_idx: int) -> None: self._jobs[transfer_job.job_id] = _JobMetadata(transfer_job, tier_idx) state = self._tier_states[tier_idx] - block_ids = {int(block_id) for block_id in transfer_job.block_ids} + block_count = len(transfer_job.block_ids) if transfer_job.is_promotion: state.active_promotion_count += 1 - state.primary_write_block_ids.update(block_ids) + state.primary_write_block_count += block_count else: state.active_cascade_count += 1 - state.primary_read_block_ids.update(block_ids) + state.primary_read_block_count += block_count def _pop_job(self, job_id: JobId) -> _JobMetadata | None: job_metadata = self._jobs.pop(job_id, None) @@ -284,13 +285,17 @@ def _pop_job(self, job_id: JobId) -> _JobMetadata | None: transfer_job = job_metadata.transfer_job state = self._tier_states[job_metadata.tier_idx] - block_ids = {int(block_id) for block_id in transfer_job.block_ids} + block_count = len(transfer_job.block_ids) if transfer_job.is_promotion: + assert state.active_promotion_count > 0 state.active_promotion_count -= 1 - state.primary_write_block_ids.difference_update(block_ids) + state.primary_write_block_count -= block_count + assert state.primary_write_block_count >= 0 else: + assert state.active_cascade_count > 0 state.active_cascade_count -= 1 - state.primary_read_block_ids.difference_update(block_ids) + state.primary_read_block_count -= block_count + assert state.primary_read_block_count >= 0 return job_metadata def _record_finished_job_stats( @@ -330,8 +335,8 @@ def _record_active_transfer_stats(self, stats: OffloadingConnectorStats) -> None num_primary_blocks = self.primary_tier._num_blocks for tier_idx, state in enumerate(self._tier_states): labelvalues = self._tier_label(tier_idx) - write_blocks = len(state.primary_write_block_ids) - read_blocks = len(state.primary_read_block_ids) + write_blocks = state.primary_write_block_count + read_blocks = state.primary_read_block_count write_usage = ( write_blocks / num_primary_blocks if num_primary_blocks > 0 else 0.0 ) @@ -361,12 +366,12 @@ def _record_active_transfer_stats(self, stats: OffloadingConnectorStats) -> None def _observe_lookup( self, - req_state: RequestState | None, + req_state: RequestState, labelvalues: TierLabel, key: OffloadKey, result: LookupResult, ) -> None: - if req_state is None or req_state.observed_lookups is None: + if req_state.observed_lookups is None: return observed = req_state.observed_lookups.setdefault(labelvalues, set()) if key in observed: @@ -504,6 +509,7 @@ def lookup( primary_hit = self.primary_tier.lookup(key, req_context) if primary_hit in (LookupResult.HIT, LookupResult.MISS): + assert req_state is not None self._observe_lookup(req_state, _PRIMARY_TIER_LABEL, key, primary_hit) if primary_hit is LookupResult.HIT: return LookupResult.HIT @@ -520,6 +526,7 @@ def lookup( labelvalues = self._tier_label(i) result = tier.lookup(key, req_context) if result in (LookupResult.HIT, LookupResult.MISS): + assert req_state is not None self._observe_lookup(req_state, labelvalues, key, result) if result is LookupResult.HIT: promoted = self._initiate_promotion(tier, key, req_context) @@ -833,7 +840,7 @@ def create_store_job( self, keys: Collection[OffloadKey], req_context: ReqContext, - tier_idx: int | None = None, + tier_idx: int, ) -> TransferJob: """Pin blocks in the primary tier and create a tracked store job. @@ -854,8 +861,7 @@ def create_store_job( is_promotion=False, req_context=req_context, ) - if tier_idx is not None: - self._register_job(job_metadata, tier_idx) + self._register_job(job_metadata, tier_idx) return job_metadata @override @@ -998,11 +1004,19 @@ def reset_cache(self) -> None: # All tier I/O has stopped; consume their completion notifications # so manager bookkeeping is consistent before the primary reset. self._process_finished_jobs() + assert not self._jobs # Deferred promotion submissions reserve primary slots that the # reset below invalidates; their submit_load() has not yet been # called so no tier I/O is touching that memory. self._pending_load_submissions.clear() + assert all( + state.active_promotion_count == 0 + and state.active_cascade_count == 0 + and state.primary_write_block_count == 0 + and state.primary_read_block_count == 0 + for state in self._tier_states + ) finished_req_ids = [] for req_id, state in self._req_state.items(): diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 54b4553b4049..f76ed59b5c78 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -475,7 +475,6 @@ def submit_load(self, job_metadata: TransferJob) -> None: keys = list(job_metadata.keys) block_ids = job_metadata.block_ids - source = job_metadata.req_context.get_state(P2PSourceInfo) logger.debug( "P2P %s: submit_load ENTRY job_id=%d blocks=%d kv_request_id=%s peer=%s", @@ -719,7 +718,9 @@ def _reap_unbound_stores(self) -> None: batches = self._unbound_stores.pop(kid) self._failed_req_ids.add(kid) for batch in batches: - self._finished_jobs.append(JobResult(job_id=batch.job_id, success=False)) + self._finished_jobs.append( + JobResult(job_id=batch.job_id, success=False) + ) logger.warning( "P2P %s: unbound store kv_request_id=%s timed out after %.0fs " "without a fetch — failing %d job(s)", @@ -754,11 +755,15 @@ def _poll_once(self) -> None: for session in self._sessions.values(): result = session.poll() for lr in result.loads: - self._finished_jobs.append(JobResult(job_id=lr.job_id, success=lr.success)) + self._finished_jobs.append( + JobResult(job_id=lr.job_id, success=lr.success) + ) if not lr.success: self._failed_req_ids.add(lr.kv_request_id) for sr in result.stores: - self._finished_jobs.append(JobResult(job_id=sr.job_id, success=sr.success)) + self._finished_jobs.append( + JobResult(job_id=sr.job_id, success=sr.success) + ) # Bind kv_request_id → session for any FetchMsg this tick and # replay any submit_store batches parked while no peer was # asking. ServerRole.on_fetch already recorded the demand @@ -792,7 +797,9 @@ def shutdown(self) -> None: # leak them; the manager is going away after this call. for batches in self._unbound_stores.values(): for batch in batches: - self._finished_jobs.append(JobResult(job_id=batch.job_id, success=False)) + self._finished_jobs.append( + JobResult(job_id=batch.job_id, success=False) + ) self._unbound_stores.clear() self._control.close() self._data.close() From 90a701057186f34abf1ef00150e30db188e2d4bf Mon Sep 17 00:00:00 2001 From: srinivas_oo7 Date: Wed, 29 Jul 2026 07:14:16 -0500 Subject: [PATCH 05/13] Extract tiering metrics tracker Signed-off-by: srinivas_oo7 --- tests/v1/kv_offload/tiering/test_metrics.py | 193 +++++++++++ .../tiering/test_tiering_offloading.py | 249 --------------- vllm/v1/kv_offload/tiering/manager.py | 300 ++++-------------- vllm/v1/kv_offload/tiering/metrics.py | 263 +++++++++++++++ 4 files changed, 519 insertions(+), 486 deletions(-) create mode 100644 tests/v1/kv_offload/tiering/test_metrics.py create mode 100644 vllm/v1/kv_offload/tiering/metrics.py diff --git a/tests/v1/kv_offload/tiering/test_metrics.py b/tests/v1/kv_offload/tiering/test_metrics.py new file mode 100644 index 000000000000..264f4873eda3 --- /dev/null +++ b/tests/v1/kv_offload/tiering/test_metrics.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterable + +import numpy as np +import pytest + +from vllm.v1.kv_offload.base import LookupResult, ReqContext, make_offload_key +from vllm.v1.kv_offload.tiering.base import ( + JobResult, + TieringOffloadingMetrics, + TransferJob, +) +from vllm.v1.kv_offload.tiering.manager import JobMetadata +from vllm.v1.kv_offload.tiering.metrics import TieringMetricsTracker + +_CTX = ReqContext(req_id="test") + + +def to_keys(int_ids: Iterable[int]): + return [make_offload_key(str(i).encode(), 0) for i in int_ids] + + +def test_tiering_metrics_tracker_records_lookup_metrics(): + tracker = TieringMetricsTracker( + tier_types=["fs", "p2p"], + num_primary_blocks=5, + primary_block_size=16, + ) + tracker.on_new_request(_CTX) + + key = to_keys([1])[0] + tracker.on_lookup( + _CTX, + key, + tracker.primary_tier_label, + LookupResult.MISS, + elapsed=0.0, + ) + tracker.on_lookup(_CTX, key, tracker.tier_label(0), LookupResult.MISS, 0.0) + tracker.on_lookup(_CTX, key, tracker.tier_label(1), LookupResult.HIT, 0.0) + tracker.on_lookup(_CTX, key, tracker.tier_label(1), LookupResult.HIT, 0.0) + + stats = tracker.take_stats() + assert stats is not None + values = stats.data["data"] + assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("0:primary",)] == 1 + assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("1:fs",)] == 1 + assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("2:p2p",)] == 1 + assert values[TieringOffloadingMetrics.BLOCK_HITS][("2:p2p",)] == 1 + assert ("1:fs",) not in values[TieringOffloadingMetrics.BLOCK_HITS] + + +def test_tiering_metrics_tracker_stops_lookup_metrics_after_allocation(): + tracker = TieringMetricsTracker( + tier_types=["fs"], + num_primary_blocks=5, + primary_block_size=16, + ) + tracker.on_new_request(_CTX) + key = to_keys([1])[0] + + tracker.on_lookup( + _CTX, + key, + tracker.primary_tier_label, + LookupResult.MISS, + elapsed=0.0, + ) + tracker.on_lookup(_CTX, key, tracker.tier_label(0), LookupResult.HIT, 0.0) + + stats = tracker.take_stats() + assert stats is not None + values = stats.data["data"] + assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("0:primary",)] == 1 + assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("1:fs",)] == 1 + assert values[TieringOffloadingMetrics.BLOCK_HITS][("1:fs",)] == 1 + + tracker.on_request_allocated(_CTX) + tracker.on_lookup( + _CTX, + key, + tracker.primary_tier_label, + LookupResult.MISS, + elapsed=0.0, + ) + tracker.on_lookup(_CTX, key, tracker.tier_label(0), LookupResult.HIT, 0.0) + + stats = tracker.take_stats() + assert stats is not None + values = stats.data["data"] + assert TieringOffloadingMetrics.BLOCK_QUERIES not in values + assert TieringOffloadingMetrics.BLOCK_HITS not in values + + +def test_tiering_metrics_tracker_records_finished_job_metrics(): + tracker = TieringMetricsTracker( + tier_types=["fs"], + num_primary_blocks=5, + primary_block_size=16, + ) + cascade_key, promotion_key, failed_cascade_key, failed_promotion_key = to_keys( + range(4) + ) + jobs = [ + JobMetadata( + TransferJob(0, [cascade_key], np.array([0]), False, _CTX), + 0, + ), + JobMetadata( + TransferJob(1, [promotion_key], np.array([1]), True, _CTX), + 0, + ), + JobMetadata( + TransferJob(2, [failed_cascade_key], np.array([2]), False, _CTX), + 0, + ), + JobMetadata( + TransferJob(3, [failed_promotion_key], np.array([3]), True, _CTX), + 0, + ), + ] + for job in jobs: + tracker.on_job_registered(job) + + tracker.on_job_finished( + jobs[0], JobResult(job_id=0, success=True, transfer_time=0.5) + ) + tracker.on_job_finished( + jobs[1], JobResult(job_id=1, success=True, transfer_time=0.25) + ) + tracker.on_job_finished(jobs[2], JobResult(job_id=2, success=False)) + tracker.on_job_finished(jobs[3], JobResult(job_id=3, success=False)) + + stats = tracker.take_stats() + assert stats is not None + values = stats.data["data"] + label = ("1:fs",) + assert values[TieringOffloadingMetrics.WRITE_BYTES][label] == 16 + assert values[TieringOffloadingMetrics.WRITE_TIME][label] == 0.5 + assert values[TieringOffloadingMetrics.READ_BYTES][label] == 16 + assert values[TieringOffloadingMetrics.READ_TIME][label] == 0.25 + assert values[TieringOffloadingMetrics.CASCADE_JOB_FAILURES][label] == 1 + assert values[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES][label] == 1 + tracker.assert_idle() + + +def test_tiering_metrics_tracker_reports_active_job_and_primary_usage_gauges(): + tracker = TieringMetricsTracker( + tier_types=["fs", "p2p"], + num_primary_blocks=6, + primary_block_size=16, + ) + fs_job = JobMetadata( + TransferJob(0, to_keys([0, 1]), np.array([0, 1]), False, _CTX), + 0, + ) + p2p_job = JobMetadata( + TransferJob(1, to_keys([2, 3, 4]), np.array([2, 3, 4]), True, _CTX), + 1, + ) + tracker.on_job_registered(fs_job) + tracker.on_job_registered(p2p_job) + + stats = tracker.take_stats() + assert stats is not None + values = stats.data["data"] + fs_label = ("1:fs",) + p2p_label = ("2:p2p",) + assert values[TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC][ + fs_label + ] == pytest.approx(2 / 6) + assert values[TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS][fs_label] == 1 + assert values[TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC][ + p2p_label + ] == pytest.approx(3 / 6) + assert values[TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS][p2p_label] == 1 + + +def test_tiering_metrics_tracker_records_promotion_allocation_failures(): + tracker = TieringMetricsTracker( + tier_types=["fs"], + num_primary_blocks=1, + primary_block_size=16, + ) + + tracker.on_promotion_allocation_failure() + + stats = tracker.take_stats() + assert stats is not None + values = stats.data["data"] + assert values[TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES][()] == 1 diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 12e3caf9852d..045b2f5b6df1 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -206,255 +206,6 @@ def test_tiering_manager_aggregates_secondary_stats(): assert MetricsSecondaryTierManager.MY_TIER_METRIC not in second_stats.data["data"] -def test_tiering_manager_records_secondary_lookup_metrics(): - mock_region = _mock_mmap_region(5) - primary_tier = CPUPrimaryTierOffloadingManager( - num_blocks=5, mmap_region=mock_region - ) - tier0 = MetricsSecondaryTierManager( - offloading_spec=_MOCK_OFFLOADING_SPEC, - primary_kv_view=mock_region.create_kv_memoryview(), - tier_type="fs", - ) - tier1 = MetricsSecondaryTierManager( - offloading_spec=_MOCK_OFFLOADING_SPEC, - primary_kv_view=mock_region.create_kv_memoryview(), - tier_type="p2p", - ) - tier1.lookup_result = LookupResult.HIT - manager = TieringOffloadingManager( - primary_tier=primary_tier, - secondary_tiers=[tier0, tier1], - ) - manager.on_new_request(_CTX) - - key = to_keys([1])[0] - assert manager.lookup(key, _CTX) is LookupResult.RETRY - assert manager.lookup(key, _CTX) is LookupResult.HIT_PENDING - - stats = manager.get_stats() - assert stats is not None - values = stats.data["data"] - assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("0:primary",)] == 1 - assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("1:fs",)] == 1 - assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("2:p2p",)] == 1 - assert values[TieringOffloadingMetrics.BLOCK_HITS][("2:p2p",)] == 1 - assert ("1:fs",) not in values[TieringOffloadingMetrics.BLOCK_HITS] - - -def test_tiering_manager_stops_lookup_metrics_after_allocation(): - mock_region = _mock_mmap_region(5) - primary_tier = CPUPrimaryTierOffloadingManager( - num_blocks=5, mmap_region=mock_region - ) - tier = MetricsSecondaryTierManager( - offloading_spec=_MOCK_OFFLOADING_SPEC, - primary_kv_view=mock_region.create_kv_memoryview(), - tier_type="fs", - ) - tier.lookup_result = LookupResult.HIT - manager = TieringOffloadingManager( - primary_tier=primary_tier, - secondary_tiers=[tier], - ) - manager.on_new_request(_CTX) - - key = to_keys([1])[0] - assert manager.lookup(key, _CTX) is LookupResult.RETRY - - stats = manager.get_stats() - assert stats is not None - values = stats.data["data"] - assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("0:primary",)] == 1 - assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("1:fs",)] == 1 - assert values[TieringOffloadingMetrics.BLOCK_HITS][("1:fs",)] == 1 - - manager.on_schedule_end( - ScheduleEndContext(new_req_ids=(_CTX.req_id,), preempted_req_ids=()) - ) - assert manager.lookup(key, _CTX) is LookupResult.HIT_PENDING - - stats = manager.get_stats() - assert stats is not None - values = stats.data["data"] - assert TieringOffloadingMetrics.BLOCK_QUERIES not in values - assert TieringOffloadingMetrics.BLOCK_HITS not in values - - -def test_tiering_manager_records_finished_job_metrics(): - mock_region = _mock_mmap_region(5) - primary_tier = CPUPrimaryTierOffloadingManager( - num_blocks=5, mmap_region=mock_region - ) - tier = MetricsSecondaryTierManager( - offloading_spec=_MOCK_OFFLOADING_SPEC, - primary_kv_view=mock_region.create_kv_memoryview(), - tier_type="fs", - ) - manager = TieringOffloadingManager( - primary_tier=primary_tier, - secondary_tiers=[tier], - ) - manager.on_new_request(_CTX) - cascade_key, promotion_key, failed_cascade_key, failed_promotion_key = to_keys( - range(4) - ) - - prepared = primary_tier.prepare_store([cascade_key, failed_cascade_key], _CTX) - assert prepared is not None - primary_tier.complete_store(prepared.keys_to_store, _CTX, success=True) - cascade_job = manager.create_store_job([cascade_key], _CTX, 0) - failed_cascade_job = manager.create_store_job([failed_cascade_key], _CTX, 0) - - promotion_prepared = primary_tier.prepare_write([promotion_key], _CTX) - assert promotion_prepared is not None - promotion_job = TransferJob( - job_id=manager._next_job_id(), - keys=[promotion_key], - block_ids=promotion_prepared.store_spec.block_ids, - is_promotion=True, - req_context=_CTX, - ) - manager._register_job(promotion_job, 0) - - failed_promotion_prepared = primary_tier.prepare_write([failed_promotion_key], _CTX) - assert failed_promotion_prepared is not None - failed_promotion_job = TransferJob( - job_id=manager._next_job_id(), - keys=[failed_promotion_key], - block_ids=failed_promotion_prepared.store_spec.block_ids, - is_promotion=True, - req_context=_CTX, - ) - manager._register_job(failed_promotion_job, 0) - - tier.finished_jobs = [ - JobResult( - job_id=cascade_job.job_id, - success=True, - transfer_time=0.5, - ), - JobResult( - job_id=promotion_job.job_id, - success=True, - transfer_time=0.25, - ), - JobResult(job_id=failed_cascade_job.job_id, success=False), - JobResult(job_id=failed_promotion_job.job_id, success=False), - ] - - manager._process_finished_jobs() - - stats = manager.get_stats() - assert stats is not None - values = stats.data["data"] - label = ("1:fs",) - assert values[TieringOffloadingMetrics.WRITE_BYTES][label] == 16 - assert values[TieringOffloadingMetrics.WRITE_TIME][label] == 0.5 - assert values[TieringOffloadingMetrics.READ_BYTES][label] == 16 - assert values[TieringOffloadingMetrics.READ_TIME][label] == 0.25 - assert values[TieringOffloadingMetrics.CASCADE_JOB_FAILURES][label] == 1 - assert values[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES][label] == 1 - - -def test_tiering_manager_reports_active_job_and_primary_usage_gauges(): - mock_region = _mock_mmap_region(6) - primary_tier = CPUPrimaryTierOffloadingManager( - num_blocks=6, mmap_region=mock_region - ) - tier0 = MetricsSecondaryTierManager( - offloading_spec=_MOCK_OFFLOADING_SPEC, - primary_kv_view=mock_region.create_kv_memoryview(), - tier_type="fs", - ) - tier1 = MetricsSecondaryTierManager( - offloading_spec=_MOCK_OFFLOADING_SPEC, - primary_kv_view=mock_region.create_kv_memoryview(), - tier_type="p2p", - ) - manager = TieringOffloadingManager( - primary_tier=primary_tier, - secondary_tiers=[tier0, tier1], - ) - manager.on_new_request(_CTX) - cascade_keys = to_keys([0, 1]) - promotion_key = to_keys([2])[0] - - prepared = manager.prepare_store(cascade_keys, _CTX) - assert prepared is not None - manager.complete_store(cascade_keys, _CTX, success=True) - tier0.lookup_result = LookupResult.HIT - assert manager.lookup(promotion_key, _CTX) is LookupResult.RETRY - manager.on_schedule_end(ScheduleEndContext(new_req_ids=(), preempted_req_ids=())) - - stats = manager.get_stats() - assert stats is not None - values = stats.data["data"] - fs_label = ("1:fs",) - p2p_label = ("2:p2p",) - assert values[TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC][ - fs_label - ] == pytest.approx(2 / 6) - assert values[TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC][ - p2p_label - ] == pytest.approx(2 / 6) - assert values[TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC][ - fs_label - ] == pytest.approx(1 / 6) - assert values[TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC][p2p_label] == 0 - assert values[TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS][fs_label] == 1 - assert values[TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS][p2p_label] == 1 - assert values[TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS][fs_label] == 1 - assert values[TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS][p2p_label] == 0 - - tier0.finished_jobs = [ - JobResult(job_id=tier0.submitted_stores[0].job_id, success=True), - JobResult(job_id=tier0.submitted_loads[0].job_id, success=True), - ] - tier1.finished_jobs = [ - JobResult(job_id=tier1.submitted_stores[0].job_id, success=True), - ] - manager._process_finished_jobs() - - stats = manager.get_stats() - assert stats is not None - values = stats.data["data"] - assert values[TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC][fs_label] == 0 - assert values[TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC][p2p_label] == 0 - assert values[TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC][fs_label] == 0 - assert values[TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC][p2p_label] == 0 - assert values[TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS][fs_label] == 0 - assert values[TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS][p2p_label] == 0 - assert values[TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS][fs_label] == 0 - assert values[TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS][p2p_label] == 0 - - -def test_tiering_manager_records_promotion_allocation_failures(): - mock_region = _mock_mmap_region(1) - primary_tier = CPUPrimaryTierOffloadingManager( - num_blocks=1, mmap_region=mock_region - ) - tier = MetricsSecondaryTierManager( - offloading_spec=_MOCK_OFFLOADING_SPEC, - primary_kv_view=mock_region.create_kv_memoryview(), - tier_type="fs", - ) - tier.lookup_result = LookupResult.HIT - manager = TieringOffloadingManager( - primary_tier=primary_tier, - secondary_tiers=[tier], - ) - manager.on_new_request(_CTX) - assert primary_tier.prepare_store(to_keys([0]), _CTX) is not None - - assert manager.lookup(to_keys([1])[0], _CTX) is LookupResult.MISS - - stats = manager.get_stats() - assert stats is not None - values = stats.data["data"] - assert values[TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES][()] == 1 - - class TestExampleSecondaryTierManager: """Tests for ExampleSecondaryTierManager implementation.""" diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 0806cb65c7a0..71ccb24293b9 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -20,7 +20,6 @@ protecting blocks from eviction until complete_read() is called """ -import functools import time from collections.abc import Collection, Iterable, Sequence from dataclasses import dataclass, field @@ -50,18 +49,14 @@ from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.kv_offload.tiering.base import ( JobId, - JobResult, ParentManager, SecondaryTierManager, - TieringOffloadingMetrics, TransferJob, ) +from vllm.v1.kv_offload.tiering.metrics import TieringMetricsTracker logger = init_logger(__name__) -TierLabel = tuple[str] -_PRIMARY_TIER_LABEL: TierLabel = ("0:primary",) - @dataclass class PendingPromotion: @@ -78,24 +73,9 @@ class RequestState: pending_primary_stores: int = 0 is_finished: bool = False request_level_tiers: set[SecondaryTierManager] | None = None - observed_lookups: dict[TierLabel, set[OffloadKey]] | None = field( - default_factory=dict - ) - sync_lookup_delay: float = 0.0 - # time.monotonic() of this request's first deferred secondary-tier lookup; - # None once consumed (observed) or while no secondary lookup is pending. - secondary_lookup_start_time: float | None = None - - -@dataclass -class _TierState: - active_promotion_count: int = 0 - active_cascade_count: int = 0 - primary_write_block_count: int = 0 - primary_read_block_count: int = 0 -class _JobMetadata(NamedTuple): +class JobMetadata(NamedTuple): transfer_job: TransferJob tier_idx: int @@ -155,23 +135,25 @@ class _SecondaryTierFacingParent(ParentManager): """Wrapper that implements ParentManager by delegating to the TieringOffloadingManager with exclude_tier set to the origin tier.""" - __slots__ = ("_m", "_origin", "_origin_idx") + __slots__ = ("_m", "_origin_idx") def __init__( self, manager: "TieringOffloadingManager", - tier: SecondaryTierManager, tier_idx: int, ): self._m = manager - self._origin = tier self._origin_idx = tier_idx def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: - return self._m.on_new_request(req_context, exclude_tier=self._origin) + return self._m.on_new_request( + req_context, exclude_tier=self._m.secondary_tiers[self._origin_idx] + ) def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: - return self._m.lookup(key, req_context, exclude_tier=self._origin) + return self._m.lookup( + key, req_context, exclude_tier=self._m.secondary_tiers[self._origin_idx] + ) def create_store_job( self, keys: Collection[OffloadKey], req_context: ReqContext @@ -179,7 +161,9 @@ def create_store_job( return self._m.create_store_job(keys, req_context, self._origin_idx) def on_request_finished(self, req_context: ReqContext) -> None: - return self._m.on_request_finished(req_context, exclude_tier=self._origin) + return self._m.on_request_finished( + req_context, exclude_tier=self._m.secondary_tiers[self._origin_idx] + ) class TieringOffloadingManager(OffloadingManager): @@ -219,16 +203,17 @@ def __init__( # TransferJob.is_promotion distinguishes direction: # True: secondary → primary (promotion) # False: primary → secondary (cascade) - self._jobs: dict[JobId, _JobMetadata] = {} - self._tier_states: list[_TierState] = [ - _TierState() for _ in self.secondary_tiers - ] + self._jobs: dict[JobId, JobMetadata] = {} self._tier_indices = { tier: tier_idx for tier_idx, tier in enumerate(self.secondary_tiers) } primary_view = self.primary_tier.get_kv_memoryview() assert primary_view.strides is not None - self._primary_block_size = primary_view.strides[0] + self._metrics = TieringMetricsTracker( + tier_types=[tier.tier_type for tier in self.secondary_tiers], + num_primary_blocks=self.primary_tier._num_blocks, + primary_block_size=primary_view.strides[0], + ) # Pending promotion requests accumulated during lookup() calls; flushed # as one batched submit_load() per (tier, request) in on_schedule_end(). @@ -249,143 +234,23 @@ def __init__( # Cached ParentManager wrappers for each secondary tier. self._tier_parents: dict[SecondaryTierManager, _SecondaryTierFacingParent] = { - tier: _SecondaryTierFacingParent(self, tier, tier_idx) + tier: _SecondaryTierFacingParent(self, tier_idx) for tier_idx, tier in enumerate(self.secondary_tiers) } - # Buffers manager-level observations (e.g. lookup delay) between - # get_stats() calls; merged in and reset each time get_stats() runs. - self._stats = OffloadingConnectorStats() - def _next_job_id(self) -> JobId: """Generate a unique job ID for async transfer tracking.""" job_id = self._job_id_counter self._job_id_counter += 1 return job_id - @functools.cache # noqa: B019 - def _tier_label(self, tier_idx: int) -> TierLabel: - return (f"{tier_idx + 1}:{self.secondary_tiers[tier_idx].tier_type}",) - def _register_job(self, transfer_job: TransferJob, tier_idx: int) -> None: - self._jobs[transfer_job.job_id] = _JobMetadata(transfer_job, tier_idx) - state = self._tier_states[tier_idx] - block_count = len(transfer_job.block_ids) - if transfer_job.is_promotion: - state.active_promotion_count += 1 - state.primary_write_block_count += block_count - else: - state.active_cascade_count += 1 - state.primary_read_block_count += block_count - - def _pop_job(self, job_id: JobId) -> _JobMetadata | None: - job_metadata = self._jobs.pop(job_id, None) - if job_metadata is None: - return None - - transfer_job = job_metadata.transfer_job - state = self._tier_states[job_metadata.tier_idx] - block_count = len(transfer_job.block_ids) - if transfer_job.is_promotion: - assert state.active_promotion_count > 0 - state.active_promotion_count -= 1 - state.primary_write_block_count -= block_count - assert state.primary_write_block_count >= 0 - else: - assert state.active_cascade_count > 0 - state.active_cascade_count -= 1 - state.primary_read_block_count -= block_count - assert state.primary_read_block_count >= 0 - return job_metadata - - def _record_finished_job_stats( - self, - job_metadata: _JobMetadata, - completed_job: JobResult, - ) -> None: - transfer_job = job_metadata.transfer_job - labelvalues = self._tier_label(job_metadata.tier_idx) - if not completed_job.success: - failure_metric = ( - TieringOffloadingMetrics.PROMOTION_JOB_FAILURES - if transfer_job.is_promotion - else TieringOffloadingMetrics.CASCADE_JOB_FAILURES - ) - self._stats.increase_counter(failure_metric, labelvalues=labelvalues) - return - - bytes_metric = ( - TieringOffloadingMetrics.READ_BYTES - if transfer_job.is_promotion - else TieringOffloadingMetrics.WRITE_BYTES - ) - time_metric = ( - TieringOffloadingMetrics.READ_TIME - if transfer_job.is_promotion - else TieringOffloadingMetrics.WRITE_TIME - ) - transfer_size = len(transfer_job.keys) * self._primary_block_size - self._stats.increase_counter(bytes_metric, transfer_size, labelvalues) - if completed_job.transfer_time is not None: - self._stats.increase_counter( - time_metric, completed_job.transfer_time, labelvalues - ) - - def _record_active_transfer_stats(self, stats: OffloadingConnectorStats) -> None: - num_primary_blocks = self.primary_tier._num_blocks - for tier_idx, state in enumerate(self._tier_states): - labelvalues = self._tier_label(tier_idx) - write_blocks = state.primary_write_block_count - read_blocks = state.primary_read_block_count - write_usage = ( - write_blocks / num_primary_blocks if num_primary_blocks > 0 else 0.0 - ) - read_usage = ( - read_blocks / num_primary_blocks if num_primary_blocks > 0 else 0.0 - ) - stats.set_gauge( - TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC, - write_usage, - labelvalues, - ) - stats.set_gauge( - TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC, - read_usage, - labelvalues, - ) - stats.set_gauge( - TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS, - state.active_promotion_count, - labelvalues, - ) - stats.set_gauge( - TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS, - state.active_cascade_count, - labelvalues, - ) + job_metadata = JobMetadata(transfer_job, tier_idx) + self._jobs[transfer_job.job_id] = job_metadata + self._metrics.on_job_registered(job_metadata) - def _observe_lookup( - self, - req_state: RequestState, - labelvalues: TierLabel, - key: OffloadKey, - result: LookupResult, - ) -> None: - if req_state.observed_lookups is None: - return - observed = req_state.observed_lookups.setdefault(labelvalues, set()) - if key in observed: - return - observed.add(key) - self._stats.increase_counter( - TieringOffloadingMetrics.BLOCK_QUERIES, - labelvalues=labelvalues, - ) - if result is LookupResult.HIT: - self._stats.increase_counter( - TieringOffloadingMetrics.BLOCK_HITS, - labelvalues=labelvalues, - ) + def _pop_job(self, job_id: JobId) -> JobMetadata | None: + return self._jobs.pop(job_id, None) def _maybe_process_finished_jobs(self): """ @@ -457,7 +322,7 @@ def _process_finished_jobs(self): f" but belongs to tier #{job_metadata.tier_idx}" ) transfer_job = job_metadata.transfer_job - self._record_finished_job_stats(job_metadata, completed_job) + self._metrics.on_job_finished(job_metadata, completed_job) if transfer_job.is_promotion: # secondary→primary transfer (promotion) completed. @@ -505,77 +370,57 @@ def lookup( # in time for a promotion this lookup may initiate. self._maybe_process_finished_jobs() - req_state = self._req_state.get(req_context.req_id) - + assert req_context.req_id in self._req_state + start_time = time.monotonic() primary_hit = self.primary_tier.lookup(key, req_context) - if primary_hit in (LookupResult.HIT, LookupResult.MISS): - assert req_state is not None - self._observe_lookup(req_state, _PRIMARY_TIER_LABEL, key, primary_hit) + elapsed = time.monotonic() - start_time + self._metrics.on_lookup( + req_context, + key, + self._metrics.primary_tier_label, + primary_hit, + elapsed, + ) if primary_hit is LookupResult.HIT: return LookupResult.HIT if primary_hit is LookupResult.HIT_PENDING: return LookupResult.HIT_PENDING - lookup_start = time.monotonic() any_retry = False for i, tier in enumerate(self.secondary_tiers): if tier is exclude_tier: continue if not req_context.load_tier_filter.allows(tier.medium, tier.locality): continue - labelvalues = self._tier_label(i) + labelvalues = self._metrics.tier_label(i) + start_time = time.monotonic() result = tier.lookup(key, req_context) - if result in (LookupResult.HIT, LookupResult.MISS): - assert req_state is not None - self._observe_lookup(req_state, labelvalues, key, result) if result is LookupResult.HIT: promoted = self._initiate_promotion(tier, key, req_context) - self._accumulate_lookup_sync_delay(req_state, lookup_start) - if ( - req_state is not None - and promoted - and req_state.secondary_lookup_start_time is None - ): - req_state.secondary_lookup_start_time = lookup_start + self._metrics.on_lookup( + req_context, + key, + labelvalues, + result, + time.monotonic() - start_time, + starts_async_delay=promoted, + ) return LookupResult.MISS if not promoted else LookupResult.RETRY if result is LookupResult.RETRY: any_retry = True + self._metrics.on_lookup( + req_context, + key, + labelvalues, + result, + time.monotonic() - start_time, + starts_async_delay=result is LookupResult.RETRY, + ) - self._accumulate_lookup_sync_delay(req_state, lookup_start) if any_retry: - if req_state is not None and req_state.secondary_lookup_start_time is None: - req_state.secondary_lookup_start_time = lookup_start return LookupResult.RETRY return LookupResult.MISS - def _accumulate_lookup_sync_delay( - self, req_state: RequestState | None, start_time: float - ) -> None: - """Accumulate secondary-tier lookup time until allocation or finish.""" - if req_state is not None: - req_state.sync_lookup_delay += time.monotonic() - start_time - - def _maybe_observe_lookup_sync_delay(self, req_state: RequestState) -> None: - delay = req_state.sync_lookup_delay - if delay == 0: - return - req_state.sync_lookup_delay = 0.0 - self._stats.observe_histogram( - TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, - delay, - ) - - def _maybe_observe_lookup_async_delay(self, req_state: RequestState) -> None: - """Flush a pending deferred secondary-tier lookup timer, if any.""" - start_time = req_state.secondary_lookup_start_time - if start_time is None: - return - req_state.secondary_lookup_start_time = None - self._stats.observe_histogram( - TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY, - time.monotonic() - start_time, - ) - def _initiate_promotion( self, tier: SecondaryTierManager, @@ -608,9 +453,7 @@ def _initiate_promotion( if primary_write_result is None: # Primary tier is full; caller should treat the block as unavailable # rather than retrying indefinitely. - self._stats.increase_counter( - TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES - ) + self._metrics.on_promotion_allocation_failure() return False store_spec = primary_write_result.store_spec @@ -878,6 +721,7 @@ def on_new_request( Only stores REQUEST_LEVEL tier decisions for use in prepare_store. """ state = RequestState(req_context=req_context) + self._metrics.on_new_request(req_context) for tier in self.secondary_tiers: if tier is exclude_tier: continue @@ -928,8 +772,7 @@ def _maybe_finalize_request( if tier is exclude_tier: continue tier.on_request_finished(state.req_context) - self._maybe_observe_lookup_sync_delay(state) - self._maybe_observe_lookup_async_delay(state) + self._metrics.on_request_finished(state.req_context) del self._req_state[req_id] @override @@ -960,9 +803,7 @@ def on_schedule_end(self, context: ScheduleEndContext) -> None: state = self._req_state.get(req_id) if state is None: continue - state.observed_lookups = None - self._maybe_observe_lookup_sync_delay(state) - self._maybe_observe_lookup_async_delay(state) + self._metrics.on_request_allocated(state.req_context) @override def has_pending_work(self) -> bool: @@ -1010,13 +851,7 @@ def reset_cache(self) -> None: # reset below invalidates; their submit_load() has not yet been # called so no tier I/O is touching that memory. self._pending_load_submissions.clear() - assert all( - state.active_promotion_count == 0 - and state.active_cascade_count == 0 - and state.primary_write_block_count == 0 - and state.primary_read_block_count == 0 - for state in self._tier_states - ) + self._metrics.assert_idle() finished_req_ids = [] for req_id, state in self._req_state.items(): @@ -1025,8 +860,7 @@ def reset_cache(self) -> None: continue for tier in self.secondary_tiers: tier.on_request_finished(state.req_context) - self._maybe_observe_lookup_sync_delay(state) - self._maybe_observe_lookup_async_delay(state) + self._metrics.on_request_finished(state.req_context) finished_req_ids.append(req_id) self.primary_tier.reset_cache() @@ -1042,20 +876,12 @@ def get_stats(self) -> OffloadingConnectorStats | None: if stats is not None and stats.is_empty(): stats = None - active_transfer_stats = OffloadingConnectorStats() - self._record_active_transfer_stats(active_transfer_stats) - if not active_transfer_stats.is_empty(): - if stats is None: - stats = active_transfer_stats - else: - stats.aggregate(active_transfer_stats) - - if not self._stats.is_empty(): + metrics_stats = self._metrics.take_stats() + if metrics_stats is not None: if stats is None: - stats = self._stats + stats = metrics_stats else: - stats.aggregate(self._stats) - self._stats = OffloadingConnectorStats() + stats.aggregate(metrics_stats) for tier in self.secondary_tiers: tier_stats = tier.get_stats() diff --git a/vllm/v1/kv_offload/tiering/metrics.py b/vllm/v1/kv_offload/tiering/metrics.py new file mode 100644 index 000000000000..810fa33ee72a --- /dev/null +++ b/vllm/v1/kv_offload/tiering/metrics.py @@ -0,0 +1,263 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import functools +import time +from dataclasses import dataclass, field +from typing import Protocol + +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( + OffloadingConnectorStats, +) +from vllm.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext +from vllm.v1.kv_offload.tiering.base import ( + JobResult, + TieringOffloadingMetrics, + TransferJob, +) + +TierLabel = tuple[str] +PRIMARY_TIER_LABEL: TierLabel = ("0:primary",) + + +class _JobMetadataLike(Protocol): + @property + def transfer_job(self) -> TransferJob: ... + + @property + def tier_idx(self) -> int: ... + + +@dataclass(slots=True) +class _RequestMetricsState: + observed_lookups: dict[TierLabel, set[OffloadKey]] | None = field( + default_factory=dict + ) + sync_lookup_delay: float = 0.0 + secondary_lookup_start_time: float | None = None + + +@dataclass +class _TierState: + active_promotion_count: int = 0 + active_cascade_count: int = 0 + primary_write_block_count: int = 0 + primary_read_block_count: int = 0 + + +class TieringMetricsTracker: + def __init__( + self, + tier_types: list[str], + num_primary_blocks: int, + primary_block_size: int, + ) -> None: + self._tier_types = tier_types + self._num_primary_blocks = num_primary_blocks + self._primary_block_size = primary_block_size + self._request_states: dict[str, _RequestMetricsState] = {} + self._tier_states = [_TierState() for _ in tier_types] + self._stats = OffloadingConnectorStats() + + @functools.cache # noqa: B019 + def tier_label(self, tier_idx: int) -> TierLabel: + return (f"{tier_idx + 1}:{self._tier_types[tier_idx]}",) + + @property + def primary_tier_label(self) -> TierLabel: + return PRIMARY_TIER_LABEL + + def on_new_request(self, req_context: ReqContext) -> None: + self._request_states[req_context.req_id] = _RequestMetricsState() + + def on_request_allocated(self, req_context: ReqContext) -> None: + state = self._request_states.get(req_context.req_id) + if state is None: + return + state.observed_lookups = None + self._flush_lookup_delays(state) + + def on_request_finished(self, req_context: ReqContext) -> None: + state = self._request_states.pop(req_context.req_id, None) + if state is None: + return + self._flush_lookup_delays(state) + + def on_lookup( + self, + req_context: ReqContext, + key: OffloadKey, + tier_label: TierLabel, + result: LookupResult, + elapsed: float, + *, + starts_async_delay: bool = False, + ) -> None: + state = self._request_states.get(req_context.req_id) + assert state is not None + state.sync_lookup_delay += elapsed + + if state.observed_lookups is not None and result in ( + LookupResult.HIT, + LookupResult.MISS, + ): + observed = state.observed_lookups.setdefault(tier_label, set()) + if key not in observed: + observed.add(key) + self._stats.increase_counter( + TieringOffloadingMetrics.BLOCK_QUERIES, + labelvalues=tier_label, + ) + if result is LookupResult.HIT: + self._stats.increase_counter( + TieringOffloadingMetrics.BLOCK_HITS, + labelvalues=tier_label, + ) + + if starts_async_delay and state.secondary_lookup_start_time is None: + state.secondary_lookup_start_time = time.monotonic() + + def on_job_registered(self, job_metadata: _JobMetadataLike) -> None: + transfer_job = job_metadata.transfer_job + state = self._tier_states[job_metadata.tier_idx] + block_count = len(transfer_job.block_ids) + if transfer_job.is_promotion: + state.active_promotion_count += 1 + state.primary_write_block_count += block_count + else: + state.active_cascade_count += 1 + state.primary_read_block_count += block_count + + def on_job_finished( + self, job_metadata: _JobMetadataLike, result: JobResult + ) -> None: + self._observe_finished_job_stats(job_metadata, result) + self._decrement_tier_state(job_metadata) + + def on_promotion_allocation_failure(self) -> None: + self._stats.increase_counter( + TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES + ) + + def take_stats(self) -> OffloadingConnectorStats | None: + active_transfer_stats = OffloadingConnectorStats() + self._observe_active_transfer_stats(active_transfer_stats) + + stats = None + if not active_transfer_stats.is_empty(): + stats = active_transfer_stats + if not self._stats.is_empty(): + if stats is None: + stats = self._stats + else: + stats.aggregate(self._stats) + self._stats = OffloadingConnectorStats() + return stats + + def assert_idle(self) -> None: + assert all( + state.active_promotion_count == 0 + and state.active_cascade_count == 0 + and state.primary_write_block_count == 0 + and state.primary_read_block_count == 0 + for state in self._tier_states + ) + + def _decrement_tier_state(self, job_metadata: _JobMetadataLike) -> None: + transfer_job = job_metadata.transfer_job + state = self._tier_states[job_metadata.tier_idx] + block_count = len(transfer_job.block_ids) + if transfer_job.is_promotion: + assert state.active_promotion_count > 0 + state.active_promotion_count -= 1 + state.primary_write_block_count -= block_count + assert state.primary_write_block_count >= 0 + else: + assert state.active_cascade_count > 0 + state.active_cascade_count -= 1 + state.primary_read_block_count -= block_count + assert state.primary_read_block_count >= 0 + + def _observe_finished_job_stats( + self, + job_metadata: _JobMetadataLike, + completed_job: JobResult, + ) -> None: + transfer_job = job_metadata.transfer_job + labelvalues = self.tier_label(job_metadata.tier_idx) + if not completed_job.success: + failure_metric = ( + TieringOffloadingMetrics.PROMOTION_JOB_FAILURES + if transfer_job.is_promotion + else TieringOffloadingMetrics.CASCADE_JOB_FAILURES + ) + self._stats.increase_counter(failure_metric, labelvalues=labelvalues) + return + + bytes_metric = ( + TieringOffloadingMetrics.READ_BYTES + if transfer_job.is_promotion + else TieringOffloadingMetrics.WRITE_BYTES + ) + time_metric = ( + TieringOffloadingMetrics.READ_TIME + if transfer_job.is_promotion + else TieringOffloadingMetrics.WRITE_TIME + ) + transfer_size = len(transfer_job.keys) * self._primary_block_size + self._stats.increase_counter(bytes_metric, transfer_size, labelvalues) + if completed_job.transfer_time is not None: + self._stats.increase_counter( + time_metric, completed_job.transfer_time, labelvalues + ) + + def _observe_active_transfer_stats(self, stats: OffloadingConnectorStats) -> None: + for tier_idx, state in enumerate(self._tier_states): + labelvalues = self.tier_label(tier_idx) + write_usage = ( + state.primary_write_block_count / self._num_primary_blocks + if self._num_primary_blocks > 0 + else 0.0 + ) + read_usage = ( + state.primary_read_block_count / self._num_primary_blocks + if self._num_primary_blocks > 0 + else 0.0 + ) + stats.set_gauge( + TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC, + write_usage, + labelvalues, + ) + stats.set_gauge( + TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC, + read_usage, + labelvalues, + ) + stats.set_gauge( + TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS, + state.active_promotion_count, + labelvalues, + ) + stats.set_gauge( + TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS, + state.active_cascade_count, + labelvalues, + ) + + def _flush_lookup_delays(self, state: _RequestMetricsState) -> None: + sync_delay = state.sync_lookup_delay + if sync_delay != 0: + state.sync_lookup_delay = 0.0 + self._stats.observe_histogram( + TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, + sync_delay, + ) + + async_start = state.secondary_lookup_start_time + if async_start is not None: + state.secondary_lookup_start_time = None + self._stats.observe_histogram( + TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY, + time.monotonic() - async_start, + ) From 9d77c6647f856cc28d9484e17523eab7912f542a Mon Sep 17 00:00:00 2001 From: srinivas_oo7 Date: Wed, 29 Jul 2026 22:09:46 -0500 Subject: [PATCH 06/13] Address tiering metrics review comments Signed-off-by: srinivas_oo7 --- .../tiering/test_tiering_offloading.py | 131 +++++++++++++----- vllm/v1/kv_offload/tiering/manager.py | 77 +++++----- vllm/v1/kv_offload/tiering/metrics.py | 74 +++++----- vllm/v1/kv_offload/tiering/spec.py | 13 +- 4 files changed, 171 insertions(+), 124 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 045b2f5b6df1..1957fac29eb4 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -86,6 +86,16 @@ def count_hits(manager, keys: list[OffloadKey]) -> int | None: return count +def histogram_count_key(metric_name: str, labelvalues: tuple[str, ...]) -> str: + return f"{metric_name}:{labelvalues}_count" + + +def has_histogram_count(reduced: dict[str, int | float], metric_name: str) -> bool: + return any( + key.startswith(f"{metric_name}:") and key.endswith("_count") for key in reduced + ) + + class MetricsSecondaryTierManager(SecondaryTierManager): """Test-only secondary tier that declares and emits one labeled metric.""" @@ -151,6 +161,8 @@ def test_tiering_spec_collects_secondary_metric_definitions(monkeypatch): assert metadata.documentation == "Number of bytes served by the test tier." assert metadata.labelnames == ("tier",) assert metrics[TieringOffloadingMetrics.BLOCK_QUERIES].labelnames == ("tier",) + assert metrics[TieringOffloadingMetrics.LOOKUP_SYNC_DELAY].labelnames == ("tier",) + assert metrics[TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY].labelnames == ("tier",) assert metrics[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES].labelnames == ( "tier", ) @@ -525,7 +537,7 @@ def submit_partial(job_metadata: JobMetadata) -> None: ] == expected_results def test_lookup_reports_sync_delay_for_resolved_lookups(self, manager_setup): - """Resolved lookups report one sync delay sample on allocation.""" + """Resolved lookups report one sync delay sample per tier and block.""" self._start_request() blocks = to_keys(range(2)) @@ -533,22 +545,39 @@ def test_lookup_reports_sync_delay_for_resolved_lookups(self, manager_setup): for block in blocks: assert self.manager.lookup(block, _CTX) is LookupResult.MISS - stats = self.manager.get_stats() - if stats is not None: - assert f"{TieringOffloadingMetrics.LOOKUP_SYNC_DELAY}_count" not in ( - stats.reduce() - ) - - self._simulate_on_schedule_end(new_req_ids=[_CTX.req_id]) - stats = self.manager.get_stats() assert stats is not None reduced = stats.reduce() - assert reduced[f"{TieringOffloadingMetrics.LOOKUP_SYNC_DELAY}_count"] == 1 - assert f"{TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY}_count" not in reduced + assert ( + reduced[ + histogram_count_key( + TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, ("0:primary",) + ) + ] + == 2 + ) + assert ( + reduced[ + histogram_count_key( + TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, ("1:example",) + ) + ] + == 2 + ) + assert ( + reduced[ + histogram_count_key( + TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, ("2:example",) + ) + ] + == 2 + ) + assert not has_histogram_count( + reduced, TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY + ) - def test_lookup_reports_async_delay_across_promotion(self, manager_setup): - """A new request reports async delay at schedule end.""" + def test_lookup_does_not_report_async_delay_for_promotion(self, manager_setup): + """Promotion time is not included in lookup async delay.""" self._start_request() block = to_keys(range(1))[0] self.secondary_tier1.blocks[block] = True @@ -556,13 +585,30 @@ def test_lookup_reports_async_delay_across_promotion(self, manager_setup): # First lookup finds the block in a secondary tier and defers. assert self.manager.lookup(block, _CTX) is LookupResult.RETRY - # The first scheduler step reports async delay for the new request. + # Promotion is not treated as unresolved secondary lookup time. self._simulate_on_schedule_end(new_req_ids=[_CTX.req_id]) stats = self.manager.get_stats() assert stats is not None reduced = stats.reduce() - assert reduced[f"{TieringOffloadingMetrics.LOOKUP_SYNC_DELAY}_count"] == 1 - assert reduced[f"{TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY}_count"] == 1 + assert ( + reduced[ + histogram_count_key( + TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, ("0:primary",) + ) + ] + == 1 + ) + assert ( + reduced[ + histogram_count_key( + TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, ("1:example",) + ) + ] + == 1 + ) + assert not has_histogram_count( + reduced, TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY + ) # Promotion completes on the next scheduler step. self._simulate_on_schedule_end() @@ -572,38 +618,51 @@ def test_lookup_reports_async_delay_across_promotion(self, manager_setup): stats = self.manager.get_stats() if stats is not None: - assert f"{TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY}_count" not in ( - stats.reduce() + assert not has_histogram_count( + stats.reduce(), TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY ) - def test_lookup_reports_async_delay_on_request_finish(self, manager_setup): - """Never-allocated lookup delays flush at teardown.""" + def test_lookup_reports_async_delay_when_deferred_lookup_resolves( + self, manager_setup + ): + """Async delay is observed when a RETRY lookup later resolves.""" ctx = ReqContext(req_id="req_lookup_finish") self._start_request(ctx) block = to_keys(range(1))[0] - self.secondary_tier1.blocks[block] = True + self.secondary_tier1.lookup = MagicMock( + side_effect=[LookupResult.RETRY, LookupResult.HIT] + ) - # Lookup finds the block in a secondary tier and defers. + # First lookup is deferred by the secondary tier. assert self.manager.lookup(block, ctx) is LookupResult.RETRY - self._simulate_on_schedule_end() stats = self.manager.get_stats() if stats is not None: - assert f"{TieringOffloadingMetrics.LOOKUP_SYNC_DELAY}_count" not in ( - stats.reduce() - ) - assert f"{TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY}_count" not in ( - stats.reduce() + assert not has_histogram_count( + stats.reduce(), TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY ) - # Request finishes before the deferred lookup is ever resolved. - self.manager.on_request_finished(ctx) + assert self.manager.lookup(block, ctx) is LookupResult.RETRY stats = self.manager.get_stats() assert stats is not None reduced = stats.reduce() - assert reduced[f"{TieringOffloadingMetrics.LOOKUP_SYNC_DELAY}_count"] == 1 - assert reduced[f"{TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY}_count"] == 1 + assert ( + reduced[ + histogram_count_key( + TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, ("1:example",) + ) + ] + == 1 + ) + assert ( + reduced[ + histogram_count_key( + TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY, ("1:example",) + ) + ] + == 1 + ) def test_partial_lookup(self, manager_setup): """Test lookup with partial hits.""" @@ -963,9 +1022,7 @@ def test_on_new_request_lifecycle(self, manager_setup): ctx = ReqContext(req_id="req_policy_lifecycle_2") result = self.manager.on_new_request(ctx) assert result.policy == OffloadPolicy.REQUEST_LEVEL - assert self.manager._req_state[ctx.req_id].request_level_tiers == { - self.secondary_tier1 - } + assert self.manager._req_state[ctx.req_id].request_level_tiers == {0} # Cleanup self.manager.on_request_finished(ctx) @@ -1045,9 +1102,7 @@ def test_reset_cache_clears_orchestrator_state(self, manager_setup): ) rl_ctx = ReqContext(req_id="rl") self.manager.on_new_request(rl_ctx) - assert self.manager._req_state[rl_ctx.req_id].request_level_tiers == { - self.secondary_tier1 - } + assert self.manager._req_state[rl_ctx.req_id].request_level_tiers == {0} # Mark this step as already polled (reset_cache must clear it). self.manager._processed_jobs_this_step = True diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 71ccb24293b9..427a57afad65 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -72,7 +72,7 @@ class RequestState: req_context: ReqContext pending_primary_stores: int = 0 is_finished: bool = False - request_level_tiers: set[SecondaryTierManager] | None = None + request_level_tiers: set[int] | None = None class JobMetadata(NamedTuple): @@ -133,7 +133,7 @@ def shutdown(self) -> None: class _SecondaryTierFacingParent(ParentManager): """Wrapper that implements ParentManager by delegating to the - TieringOffloadingManager with exclude_tier set to the origin tier.""" + TieringOffloadingManager with exclude_tier_idx set to the origin tier.""" __slots__ = ("_m", "_origin_idx") @@ -146,14 +146,10 @@ def __init__( self._origin_idx = tier_idx def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: - return self._m.on_new_request( - req_context, exclude_tier=self._m.secondary_tiers[self._origin_idx] - ) + return self._m.on_new_request(req_context, exclude_tier_idx=self._origin_idx) def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: - return self._m.lookup( - key, req_context, exclude_tier=self._m.secondary_tiers[self._origin_idx] - ) + return self._m.lookup(key, req_context, exclude_tier_idx=self._origin_idx) def create_store_job( self, keys: Collection[OffloadKey], req_context: ReqContext @@ -162,7 +158,7 @@ def create_store_job( def on_request_finished(self, req_context: ReqContext) -> None: return self._m.on_request_finished( - req_context, exclude_tier=self._m.secondary_tiers[self._origin_idx] + req_context, exclude_tier_idx=self._origin_idx ) @@ -204,9 +200,6 @@ def __init__( # True: secondary → primary (promotion) # False: primary → secondary (cascade) self._jobs: dict[JobId, JobMetadata] = {} - self._tier_indices = { - tier: tier_idx for tier_idx, tier in enumerate(self.secondary_tiers) - } primary_view = self.primary_tier.get_kv_memoryview() assert primary_view.strides is not None self._metrics = TieringMetricsTracker( @@ -217,11 +210,9 @@ def __init__( # Pending promotion requests accumulated during lookup() calls; flushed # as one batched submit_load() per (tier, request) in on_schedule_end(). - # Outer key: tier. Inner key: req_context.req_id — the same ReqContext + # Outer key: tier index. Inner key: req_context.req_id — the same ReqContext # object is reused for all block lookups of a given request per engine step. - self._pending_load_submissions: dict[ - SecondaryTierManager, dict[str, PendingPromotion] - ] = {} + self._pending_load_submissions: dict[int, dict[str, PendingPromotion]] = {} # Gate for once-per-step execution of _maybe_process_finished_jobs(). # Reset at the end of each step in on_schedule_end(). @@ -341,7 +332,7 @@ def lookup( key: OffloadKey, req_context: ReqContext, *, - exclude_tier: SecondaryTierManager | None = None, + exclude_tier_idx: int | None = None, ) -> LookupResult: """ Check whether a single block is offloaded and ready. @@ -388,7 +379,7 @@ def lookup( any_retry = False for i, tier in enumerate(self.secondary_tiers): - if tier is exclude_tier: + if i == exclude_tier_idx: continue if not req_context.load_tier_filter.allows(tier.medium, tier.locality): continue @@ -396,14 +387,13 @@ def lookup( start_time = time.monotonic() result = tier.lookup(key, req_context) if result is LookupResult.HIT: - promoted = self._initiate_promotion(tier, key, req_context) + promoted = self._initiate_promotion(i, key, req_context) self._metrics.on_lookup( req_context, key, labelvalues, result, time.monotonic() - start_time, - starts_async_delay=promoted, ) return LookupResult.MISS if not promoted else LookupResult.RETRY if result is LookupResult.RETRY: @@ -414,7 +404,9 @@ def lookup( labelvalues, result, time.monotonic() - start_time, - starts_async_delay=result is LookupResult.RETRY, + async_delay_start_time=( + start_time if result is LookupResult.RETRY else None + ), ) if any_retry: @@ -423,7 +415,7 @@ def lookup( def _initiate_promotion( self, - tier: SecondaryTierManager, + tier_idx: int, key: OffloadKey, req_context: ReqContext, ) -> bool: @@ -437,7 +429,7 @@ def _initiate_promotion( batched job. Args: - tier: The secondary tier to promote from + tier_idx: The secondary tier index to promote from key: Block to promote req_context: Per-request context forwarded to primary.prepare_write(). @@ -460,7 +452,7 @@ def _initiate_promotion( assert isinstance(store_spec, CPULoadStoreSpec) # Defer submit_load to on_schedule_end(). Group by (tier, request) so # each request's blocks are submitted as one batched job per tier. - tier_pending = self._pending_load_submissions.setdefault(tier, {}) + tier_pending = self._pending_load_submissions.setdefault(tier_idx, {}) ctx_id = req_context.req_id if ctx_id not in tier_pending: tier_pending[ctx_id] = PendingPromotion( @@ -480,8 +472,8 @@ def _flush_pending_promotions(self) -> None: if not self._pending_load_submissions: return - for tier, pending_by_ctx in self._pending_load_submissions.items(): - tier_idx = self._tier_indices[tier] + for tier_idx, pending_by_ctx in self._pending_load_submissions.items(): + tier = self.secondary_tiers[tier_idx] for entry in pending_by_ctx.values(): job_id = self._next_job_id() job_metadata = TransferJob( @@ -610,7 +602,7 @@ def _cascade_existing_blocks_to_request_level_tiers( self, keys: Sequence[OffloadKey], req_context: ReqContext, - request_level_tiers: set[SecondaryTierManager], + request_level_tiers: set[int], ) -> None: """ For tiers that requested request-level policy, submit_store() for @@ -625,10 +617,9 @@ def _cascade_existing_blocks_to_request_level_tiers( if not ready_keys: return - for tier in request_level_tiers: - job_metadata = self.create_store_job( - ready_keys, req_context, self._tier_indices[tier] - ) + for tier_idx in request_level_tiers: + job_metadata = self.create_store_job(ready_keys, req_context, tier_idx) + tier = self.secondary_tiers[tier_idx] tier.submit_store(job_metadata) @override @@ -665,10 +656,8 @@ def complete_store( # LoadStoreSpec AND to increment ref_cnt (protecting blocks from # eviction during the async transfer). One prepare_read() call per # secondary tier. - for tier in self.secondary_tiers: - job_metadata = self.create_store_job( - keys, req_context, self._tier_indices[tier] - ) + for tier_idx, tier in enumerate(self.secondary_tiers): + job_metadata = self.create_store_job(keys, req_context, tier_idx) tier.submit_store(job_metadata) # Note: The async transfers are now in flight. Their completion is @@ -712,7 +701,7 @@ def on_new_request( self, req_context: ReqContext, *, - exclude_tier: SecondaryTierManager | None = None, + exclude_tier_idx: int | None = None, ) -> RequestOffloadingContext: """ Query each secondary tier for its offload policy preference. @@ -722,14 +711,14 @@ def on_new_request( """ state = RequestState(req_context=req_context) self._metrics.on_new_request(req_context) - for tier in self.secondary_tiers: - if tier is exclude_tier: + for tier_idx, tier in enumerate(self.secondary_tiers): + if tier_idx == exclude_tier_idx: continue tier_ctx = tier.on_new_request(req_context) if tier_ctx.policy == OffloadPolicy.REQUEST_LEVEL: if state.request_level_tiers is None: state.request_level_tiers = set() - state.request_level_tiers.add(tier) + state.request_level_tiers.add(tier_idx) self._req_state[req_context.req_id] = state policy = ( @@ -744,17 +733,17 @@ def on_request_finished( self, req_context: ReqContext, *, - exclude_tier: SecondaryTierManager | None = None, + exclude_tier_idx: int | None = None, ) -> None: self.primary_tier.on_request_finished(req_context) state = self._req_state[req_context.req_id] state.is_finished = True - self._maybe_finalize_request(req_context.req_id, exclude_tier) + self._maybe_finalize_request(req_context.req_id, exclude_tier_idx) def _maybe_finalize_request( self, req_id: str, - exclude_tier: SecondaryTierManager | None = None, + exclude_tier_idx: int | None = None, ) -> None: """Finalize secondary tiers once no more store cascades can be submitted. @@ -768,8 +757,8 @@ def _maybe_finalize_request( if state.pending_primary_stores != 0: return - for tier in self.secondary_tiers: - if tier is exclude_tier: + for tier_idx, tier in enumerate(self.secondary_tiers): + if tier_idx == exclude_tier_idx: continue tier.on_request_finished(state.req_context) self._metrics.on_request_finished(state.req_context) diff --git a/vllm/v1/kv_offload/tiering/metrics.py b/vllm/v1/kv_offload/tiering/metrics.py index 810fa33ee72a..28f86b9a6235 100644 --- a/vllm/v1/kv_offload/tiering/metrics.py +++ b/vllm/v1/kv_offload/tiering/metrics.py @@ -30,11 +30,9 @@ def tier_idx(self) -> int: ... @dataclass(slots=True) class _RequestMetricsState: - observed_lookups: dict[TierLabel, set[OffloadKey]] | None = field( + observed_lookups: dict[TierLabel, dict[OffloadKey, float | None]] | None = field( default_factory=dict ) - sync_lookup_delay: float = 0.0 - secondary_lookup_start_time: float | None = None @dataclass @@ -75,13 +73,9 @@ def on_request_allocated(self, req_context: ReqContext) -> None: if state is None: return state.observed_lookups = None - self._flush_lookup_delays(state) def on_request_finished(self, req_context: ReqContext) -> None: - state = self._request_states.pop(req_context.req_id, None) - if state is None: - return - self._flush_lookup_delays(state) + self._request_states.pop(req_context.req_id, None) def on_lookup( self, @@ -91,31 +85,28 @@ def on_lookup( result: LookupResult, elapsed: float, *, - starts_async_delay: bool = False, + async_delay_start_time: float | None = None, ) -> None: state = self._request_states.get(req_context.req_id) assert state is not None - state.sync_lookup_delay += elapsed if state.observed_lookups is not None and result in ( LookupResult.HIT, LookupResult.MISS, ): - observed = state.observed_lookups.setdefault(tier_label, set()) - if key not in observed: - observed.add(key) - self._stats.increase_counter( - TieringOffloadingMetrics.BLOCK_QUERIES, - labelvalues=tier_label, + observed = state.observed_lookups.setdefault(tier_label, {}) + start_time = observed.get(key) + if key not in observed or start_time is not None: + observed[key] = None + self._observe_resolved_lookup( + tier_label, + result, + elapsed, + start_time, ) - if result is LookupResult.HIT: - self._stats.increase_counter( - TieringOffloadingMetrics.BLOCK_HITS, - labelvalues=tier_label, - ) - - if starts_async_delay and state.secondary_lookup_start_time is None: - state.secondary_lookup_start_time = time.monotonic() + elif async_delay_start_time is not None and state.observed_lookups is not None: + observed = state.observed_lookups.setdefault(tier_label, {}) + observed.setdefault(key, async_delay_start_time) def on_job_registered(self, job_metadata: _JobMetadataLike) -> None: transfer_job = job_metadata.transfer_job @@ -245,19 +236,30 @@ def _observe_active_transfer_stats(self, stats: OffloadingConnectorStats) -> Non labelvalues, ) - def _flush_lookup_delays(self, state: _RequestMetricsState) -> None: - sync_delay = state.sync_lookup_delay - if sync_delay != 0: - state.sync_lookup_delay = 0.0 - self._stats.observe_histogram( - TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, - sync_delay, + def _observe_resolved_lookup( + self, + tier_label: TierLabel, + result: LookupResult, + elapsed: float, + async_start_time: float | None, + ) -> None: + self._stats.increase_counter( + TieringOffloadingMetrics.BLOCK_QUERIES, + labelvalues=tier_label, + ) + if result is LookupResult.HIT: + self._stats.increase_counter( + TieringOffloadingMetrics.BLOCK_HITS, + labelvalues=tier_label, ) - - async_start = state.secondary_lookup_start_time - if async_start is not None: - state.secondary_lookup_start_time = None + self._stats.observe_histogram( + TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, + elapsed, + tier_label, + ) + if async_start_time is not None: self._stats.observe_histogram( TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY, - time.monotonic() - async_start, + time.monotonic() - async_start_time, + tier_label, ) diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 55b81112ce56..d1d84c356e6b 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -105,10 +105,10 @@ def build_metric_definitions( metrics[TieringOffloadingMetrics.LOOKUP_SYNC_DELAY] = ( OffloadingHistogramMetadata( documentation=( - "Histogram of total blocking time spent querying secondary " - "tiers for a request, accumulated from first lookup until " - "the request is allocated or finishes, in seconds." + "Histogram of blocking time spent in a per-block tier lookup " + "that resolved as a hit or miss, labeled by tier, in seconds." ), + labelnames=("tier",), buckets=( 0.00001, 0.00005, @@ -127,10 +127,11 @@ def build_metric_definitions( metrics[TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY] = ( OffloadingHistogramMetadata( documentation=( - "Histogram of wall-clock time from a request's first deferred " - "secondary-tier lookup until the request is allocated or " - "finishes, in seconds." + "Histogram of wall-clock time from a per-block tier lookup " + "first returning retry until that same tier lookup resolves " + "as a hit or miss, labeled by tier, in seconds." ), + labelnames=("tier",), buckets=( 0.0001, 0.0005, From 4b33637fc06407bfda4f4e9f0f2adced05e837f1 Mon Sep 17 00:00:00 2001 From: srinivas_oo7 Date: Tue, 4 Aug 2026 06:09:08 -0500 Subject: [PATCH 07/13] Refine tiering lookup timing Signed-off-by: srinivas_oo7 --- tests/v1/kv_offload/tiering/test_metrics.py | 6 +++--- vllm/v1/kv_offload/tiering/manager.py | 14 ++++++-------- vllm/v1/kv_offload/tiering/metrics.py | 14 ++++++-------- 3 files changed, 15 insertions(+), 19 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_metrics.py b/tests/v1/kv_offload/tiering/test_metrics.py index 264f4873eda3..d801a2f6cc08 100644 --- a/tests/v1/kv_offload/tiering/test_metrics.py +++ b/tests/v1/kv_offload/tiering/test_metrics.py @@ -36,7 +36,7 @@ def test_tiering_metrics_tracker_records_lookup_metrics(): key, tracker.primary_tier_label, LookupResult.MISS, - elapsed=0.0, + lookup_duration=0.0, ) tracker.on_lookup(_CTX, key, tracker.tier_label(0), LookupResult.MISS, 0.0) tracker.on_lookup(_CTX, key, tracker.tier_label(1), LookupResult.HIT, 0.0) @@ -66,7 +66,7 @@ def test_tiering_metrics_tracker_stops_lookup_metrics_after_allocation(): key, tracker.primary_tier_label, LookupResult.MISS, - elapsed=0.0, + lookup_duration=0.0, ) tracker.on_lookup(_CTX, key, tracker.tier_label(0), LookupResult.HIT, 0.0) @@ -83,7 +83,7 @@ def test_tiering_metrics_tracker_stops_lookup_metrics_after_allocation(): key, tracker.primary_tier_label, LookupResult.MISS, - elapsed=0.0, + lookup_duration=0.0, ) tracker.on_lookup(_CTX, key, tracker.tier_label(0), LookupResult.HIT, 0.0) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 427a57afad65..f96b039929e8 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -364,13 +364,13 @@ def lookup( assert req_context.req_id in self._req_state start_time = time.monotonic() primary_hit = self.primary_tier.lookup(key, req_context) - elapsed = time.monotonic() - start_time + lookup_duration = time.monotonic() - start_time self._metrics.on_lookup( req_context, key, self._metrics.primary_tier_label, primary_hit, - elapsed, + lookup_duration, ) if primary_hit is LookupResult.HIT: return LookupResult.HIT @@ -386,15 +386,16 @@ def lookup( labelvalues = self._metrics.tier_label(i) start_time = time.monotonic() result = tier.lookup(key, req_context) + lookup_duration = time.monotonic() - start_time if result is LookupResult.HIT: - promoted = self._initiate_promotion(i, key, req_context) self._metrics.on_lookup( req_context, key, labelvalues, result, - time.monotonic() - start_time, + lookup_duration, ) + promoted = self._initiate_promotion(i, key, req_context) return LookupResult.MISS if not promoted else LookupResult.RETRY if result is LookupResult.RETRY: any_retry = True @@ -403,10 +404,7 @@ def lookup( key, labelvalues, result, - time.monotonic() - start_time, - async_delay_start_time=( - start_time if result is LookupResult.RETRY else None - ), + lookup_duration, ) if any_retry: diff --git a/vllm/v1/kv_offload/tiering/metrics.py b/vllm/v1/kv_offload/tiering/metrics.py index 28f86b9a6235..ee143f67c9f0 100644 --- a/vllm/v1/kv_offload/tiering/metrics.py +++ b/vllm/v1/kv_offload/tiering/metrics.py @@ -83,9 +83,7 @@ def on_lookup( key: OffloadKey, tier_label: TierLabel, result: LookupResult, - elapsed: float, - *, - async_delay_start_time: float | None = None, + lookup_duration: float, ) -> None: state = self._request_states.get(req_context.req_id) assert state is not None @@ -101,12 +99,12 @@ def on_lookup( self._observe_resolved_lookup( tier_label, result, - elapsed, + lookup_duration, start_time, ) - elif async_delay_start_time is not None and state.observed_lookups is not None: + elif result is LookupResult.RETRY and state.observed_lookups is not None: observed = state.observed_lookups.setdefault(tier_label, {}) - observed.setdefault(key, async_delay_start_time) + observed.setdefault(key, time.monotonic()) def on_job_registered(self, job_metadata: _JobMetadataLike) -> None: transfer_job = job_metadata.transfer_job @@ -240,7 +238,7 @@ def _observe_resolved_lookup( self, tier_label: TierLabel, result: LookupResult, - elapsed: float, + lookup_duration: float, async_start_time: float | None, ) -> None: self._stats.increase_counter( @@ -254,7 +252,7 @@ def _observe_resolved_lookup( ) self._stats.observe_histogram( TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, - elapsed, + lookup_duration, tier_label, ) if async_start_time is not None: From 3b7edf972ab0251dc096da2138137ec380da82bd Mon Sep 17 00:00:00 2001 From: srinivas_oo7 Date: Wed, 5 Aug 2026 21:01:47 -0500 Subject: [PATCH 08/13] Adjust tiering async lookup timer Signed-off-by: srinivas_oo7 --- vllm/v1/kv_offload/tiering/metrics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/kv_offload/tiering/metrics.py b/vllm/v1/kv_offload/tiering/metrics.py index ee143f67c9f0..227be6f535d0 100644 --- a/vllm/v1/kv_offload/tiering/metrics.py +++ b/vllm/v1/kv_offload/tiering/metrics.py @@ -104,7 +104,7 @@ def on_lookup( ) elif result is LookupResult.RETRY and state.observed_lookups is not None: observed = state.observed_lookups.setdefault(tier_label, {}) - observed.setdefault(key, time.monotonic()) + observed.setdefault(key, time.monotonic() - lookup_duration) def on_job_registered(self, job_metadata: _JobMetadataLike) -> None: transfer_job = job_metadata.transfer_job From d1506cf834b317842a348ea2677f132a1e56e21a Mon Sep 17 00:00:00 2001 From: srinivas_oo7 Date: Wed, 5 Aug 2026 21:13:41 -0500 Subject: [PATCH 09/13] Fix tiering rebase fallout Signed-off-by: srinivas_oo7 --- .../v1/kv_offload/tiering/test_tiering_offloading.py | 2 +- vllm/v1/kv_offload/tiering/manager.py | 12 +++++++----- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 1957fac29eb4..b77443802a53 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -510,7 +510,7 @@ def test_failed_promotion_keeps_only_successful_blocks( for block in blocks: self.secondary_tier1.blocks[block] = True - def submit_partial(job_metadata: JobMetadata) -> None: + def submit_partial(job_metadata: TransferJob) -> None: successful_keys = ( None if successful_indices is None diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index f96b039929e8..2f44b274dcc9 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -49,6 +49,7 @@ from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.kv_offload.tiering.base import ( JobId, + JobResult, ParentManager, SecondaryTierManager, TransferJob, @@ -259,13 +260,14 @@ def _maybe_process_finished_jobs(self): def _complete_promotion( self, job_metadata: JobMetadata, completed_job: JobResult ) -> None: + transfer_job = job_metadata.transfer_job successful_keys = completed_job.successful_keys failed_keys: Collection[OffloadKey] if completed_job.success: - successful_keys = job_metadata.keys + successful_keys = transfer_job.keys failed_keys = () elif successful_keys: - failed_keys_set = set(job_metadata.keys) + failed_keys_set = set(transfer_job.keys) assert failed_keys_set.issuperset(successful_keys), ( f"Finished promotion job_id {completed_job.job_id} " "reported unknown successful keys" @@ -274,18 +276,18 @@ def _complete_promotion( failed_keys = failed_keys_set else: successful_keys = () - failed_keys = job_metadata.keys + failed_keys = transfer_job.keys if successful_keys: self.primary_tier.complete_write( successful_keys, - job_metadata.req_context, + transfer_job.req_context, True, ) if failed_keys: self.primary_tier.complete_write( failed_keys, - job_metadata.req_context, + transfer_job.req_context, False, ) From 54f6db29a5b8953f0882cef3e4f8fd4307032c4e Mon Sep 17 00:00:00 2001 From: srinivas_oo7 Date: Fri, 7 Aug 2026 15:37:05 -0500 Subject: [PATCH 10/13] Fix tiering metrics CI failure Signed-off-by: srinivas_oo7 --- tests/v1/kv_offload/tiering/test_metrics.py | 33 +++++++++++++++++++++ vllm/v1/kv_offload/tiering/metrics.py | 8 +++-- 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_metrics.py b/tests/v1/kv_offload/tiering/test_metrics.py index d801a2f6cc08..42b973943cae 100644 --- a/tests/v1/kv_offload/tiering/test_metrics.py +++ b/tests/v1/kv_offload/tiering/test_metrics.py @@ -146,6 +146,39 @@ def test_tiering_metrics_tracker_records_finished_job_metrics(): tracker.assert_idle() +def test_tiering_metrics_tracker_records_partial_promotion_success_bytes(): + tracker = TieringMetricsTracker( + tier_types=["fs"], + num_primary_blocks=5, + primary_block_size=16, + ) + keys = to_keys(range(3)) + job = JobMetadata( + TransferJob(0, keys, np.array([0, 1, 2]), True, _CTX), + 0, + ) + tracker.on_job_registered(job) + + tracker.on_job_finished( + job, + JobResult( + job_id=0, + success=False, + successful_keys=(keys[0], keys[2]), + transfer_time=0.5, + ), + ) + + stats = tracker.take_stats() + assert stats is not None + values = stats.data["data"] + label = ("1:fs",) + assert values[TieringOffloadingMetrics.READ_BYTES][label] == 32 + assert values[TieringOffloadingMetrics.READ_TIME][label] == 0.5 + assert values[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES][label] == 1 + tracker.assert_idle() + + def test_tiering_metrics_tracker_reports_active_job_and_primary_usage_gauges(): tracker = TieringMetricsTracker( tier_types=["fs", "p2p"], diff --git a/vllm/v1/kv_offload/tiering/metrics.py b/vllm/v1/kv_offload/tiering/metrics.py index 227be6f535d0..6991967942c1 100644 --- a/vllm/v1/kv_offload/tiering/metrics.py +++ b/vllm/v1/kv_offload/tiering/metrics.py @@ -174,6 +174,7 @@ def _observe_finished_job_stats( ) -> None: transfer_job = job_metadata.transfer_job labelvalues = self.tier_label(job_metadata.tier_idx) + completed_key_count = len(transfer_job.keys) if not completed_job.success: failure_metric = ( TieringOffloadingMetrics.PROMOTION_JOB_FAILURES @@ -181,7 +182,10 @@ def _observe_finished_job_stats( else TieringOffloadingMetrics.CASCADE_JOB_FAILURES ) self._stats.increase_counter(failure_metric, labelvalues=labelvalues) - return + if transfer_job.is_promotion and completed_job.successful_keys: + completed_key_count = len(completed_job.successful_keys) + else: + return bytes_metric = ( TieringOffloadingMetrics.READ_BYTES @@ -193,7 +197,7 @@ def _observe_finished_job_stats( if transfer_job.is_promotion else TieringOffloadingMetrics.WRITE_TIME ) - transfer_size = len(transfer_job.keys) * self._primary_block_size + transfer_size = completed_key_count * self._primary_block_size self._stats.increase_counter(bytes_metric, transfer_size, labelvalues) if completed_job.transfer_time is not None: self._stats.increase_counter( From 2d42c1f1f9d01e6c4b2c947111397cca66f7282b Mon Sep 17 00:00:00 2001 From: srinivas_oo7 Date: Fri, 7 Aug 2026 21:17:33 -0500 Subject: [PATCH 11/13] Fix tiering lookup test regressions Signed-off-by: srinivas_oo7 --- vllm/v1/kv_offload/tiering/manager.py | 7 +++++-- vllm/v1/kv_offload/tiering/metrics.py | 4 +++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 2f44b274dcc9..f8e89b0e6db1 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -230,6 +230,10 @@ def __init__( for tier_idx, tier in enumerate(self.secondary_tiers) } + @property + def _transfer_jobs(self) -> dict[JobId, JobMetadata]: + return self._jobs + def _next_job_id(self) -> JobId: """Generate a unique job ID for async transfer tracking.""" job_id = self._job_id_counter @@ -363,7 +367,6 @@ def lookup( # in time for a promotion this lookup may initiate. self._maybe_process_finished_jobs() - assert req_context.req_id in self._req_state start_time = time.monotonic() primary_hit = self.primary_tier.lookup(key, req_context) lookup_duration = time.monotonic() - start_time @@ -672,7 +675,7 @@ def create_store_job( self, keys: Collection[OffloadKey], req_context: ReqContext, - tier_idx: int, + tier_idx: int = 0, ) -> TransferJob: """Pin blocks in the primary tier and create a tracked store job. diff --git a/vllm/v1/kv_offload/tiering/metrics.py b/vllm/v1/kv_offload/tiering/metrics.py index 6991967942c1..60907142ac38 100644 --- a/vllm/v1/kv_offload/tiering/metrics.py +++ b/vllm/v1/kv_offload/tiering/metrics.py @@ -86,7 +86,9 @@ def on_lookup( lookup_duration: float, ) -> None: state = self._request_states.get(req_context.req_id) - assert state is not None + if state is None: + state = _RequestMetricsState() + self._request_states[req_context.req_id] = state if state.observed_lookups is not None and result in ( LookupResult.HIT, From a6bdf9c1c6063293c18d9345ce1988d2569d9fba Mon Sep 17 00:00:00 2001 From: srinivas_oo7 Date: Sun, 9 Aug 2026 10:56:24 -0500 Subject: [PATCH 12/13] Fix rebased tiering test setup Signed-off-by: srinivas_oo7 --- .../tiering/test_tiering_offloading.py | 30 +++++++++++-------- 1 file changed, 18 insertions(+), 12 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index b77443802a53..dccabbed4185 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -298,12 +298,15 @@ def test_failed_promotion_finalizes_primary_with_failure(self, manager_setup): self._start_request() # Register an in-flight promotion job for tier1 by hand. job_id = self.manager._next_job_id() - self.manager._transfer_jobs[job_id] = JobMetadata( - job_id=job_id, - keys=to_keys([1, 2]), - block_ids=[0, 1], - is_promotion=True, - req_context=_CTX, + self.manager._register_job( + TransferJob( + job_id=job_id, + keys=to_keys([1, 2]), + block_ids=[0, 1], + is_promotion=True, + req_context=_CTX, + ), + 0, ) failed = JobResult(job_id=job_id, success=False) with ( @@ -322,12 +325,15 @@ def test_successful_promotion_finalizes_primary_with_success(self, manager_setup self._start_request() job_id = self.manager._next_job_id() - self.manager._transfer_jobs[job_id] = JobMetadata( - job_id=job_id, - keys=to_keys([1]), - block_ids=[0], - is_promotion=True, - req_context=_CTX, + self.manager._register_job( + TransferJob( + job_id=job_id, + keys=to_keys([1]), + block_ids=[0], + is_promotion=True, + req_context=_CTX, + ), + 0, ) ok = JobResult(job_id=job_id, success=True) with ( From f6b59f8e2545dbb5a2207789e1ca3dbfe3d0f7ac Mon Sep 17 00:00:00 2001 From: srinivas_oo7 Date: Sun, 9 Aug 2026 11:11:05 -0500 Subject: [PATCH 13/13] Fix FS tier JobResult typing Signed-off-by: srinivas_oo7 --- vllm/v1/kv_offload/tiering/fs/manager.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index 0ae07fcd66a2..dae0d8d5008c 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -201,7 +201,6 @@ def __init__( ) self._lookup_manager = FsAsyncLookupManager(tier=self, tier_type=self.tier_type) - self._job_transfer_sizes: dict[JobId, int] = {} @override def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: @@ -227,7 +226,6 @@ def submit_store(self, job_metadata: TransferJob) -> None: self._block_size, self._use_o_direct, ) - self._job_transfer_sizes[job_metadata.job_id] = len(keys) * self._block_size self._pool.enqueue_store(job_metadata.job_id, 1, [task]) @override @@ -267,7 +265,6 @@ def load_task() -> None: ) raise - self._job_transfer_sizes[job_id] = len(keys) * self._block_size self._pool.enqueue_load(job_id, 1, [load_task]) @override @@ -287,7 +284,6 @@ def get_finished_jobs(self) -> Iterable[JobResult]: locality=self.locality, ) ) - transfer_size = self._job_transfer_sizes.pop(job_id, None) load_keys = self._load_job_keys.pop(job_id, None) num_succeeded = self._load_progress.pop(job_id, 0) if load_keys is not None and not success: @@ -298,17 +294,11 @@ def get_finished_jobs(self) -> Iterable[JobResult]: successful = load_keys[:num_succeeded] failed = load_keys[num_succeeded:] self._lookup_manager.mark_miss(failed) - successful_transfer_size = ( - num_succeeded * self._block_size - if transfer_size is not None - else None - ) results.append( JobResult( job_id=job_id, success=False, successful_keys=tuple(successful) if successful else None, - transfer_size=successful_transfer_size, transfer_time=transfer_time, ) ) @@ -317,7 +307,6 @@ def get_finished_jobs(self) -> Iterable[JobResult]: JobResult( job_id=job_id, success=success, - transfer_size=transfer_size, transfer_time=transfer_time, ) )