From a123159f7ab0dbacb4d8f45cdeec3cb366982e50 Mon Sep 17 00:00:00 2001 From: Srinivas Krovvidi <194645829+Srinivasoo7@users.noreply.github.com> Date: Sun, 9 Aug 2026 21:23:14 -0500 Subject: [PATCH] Add tiering offloading metrics (#48798) Signed-off-by: srinivas_oo7 Co-authored-by: srinivas_oo7 --- .../unit/offloading_connector/test_metrics.py | 38 ++- .../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_metrics.py | 226 +++++++++++++++ tests/v1/kv_offload/tiering/test_obj_tier.py | 9 +- .../tiering/test_tiering_offloading.py | 212 ++++++++++---- vllm/v1/kv_offload/tiering/base.py | 24 +- vllm/v1/kv_offload/tiering/example/manager.py | 24 +- vllm/v1/kv_offload/tiering/fs/manager.py | 24 +- vllm/v1/kv_offload/tiering/fs/thread_pool.py | 33 ++- vllm/v1/kv_offload/tiering/manager.py | 251 ++++++++-------- vllm/v1/kv_offload/tiering/metrics.py | 269 ++++++++++++++++++ vllm/v1/kv_offload/tiering/obj/manager.py | 25 +- vllm/v1/kv_offload/tiering/p2p/manager.py | 6 +- vllm/v1/kv_offload/tiering/spec.py | 109 ++++++- 16 files changed, 1031 insertions(+), 237 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_connector/unit/offloading_connector/test_metrics.py b/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py index fb6b785c88..58bbd40915 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 @@ from vllm.v1.kv_offload.base import ( 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 @@ class _FakeMetric: 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/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index c8a9d7f3ea..df7a352192 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 numpy as np 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 20e4bad218..d5692b6690 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 @@ from vllm.v1.kv_offload.base import ( 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 @@ class FakeParent: 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 7d850ca62a..f007eb49e1 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 @@ from vllm.v1.kv_offload.config import ( 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_metrics.py b/tests/v1/kv_offload/tiering/test_metrics.py new file mode 100644 index 0000000000..42b973943c --- /dev/null +++ b/tests/v1/kv_offload/tiering/test_metrics.py @@ -0,0 +1,226 @@ +# 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, + 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) + 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, + lookup_duration=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, + lookup_duration=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_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"], + 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_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index 0e86038dd6..54ec4c1f9f 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 @@ from vllm.v1.kv_offload.config import ( 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 @@ class MockNixlAgent: 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 2e9fa224d2..dccabbed41 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 @@ from vllm.v1.kv_offload.base import ( Medium, OffloadingCounterMetadata, OffloadingEvent, + OffloadingGaugeMetadata, OffloadKey, OffloadPolicy, ReqContext, @@ -39,10 +40,10 @@ from vllm.v1.kv_offload.base import ( 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 @@ -85,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.""" @@ -102,18 +113,26 @@ class MetricsSecondaryTierManager(SecondaryTierManager): 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[TransferJob] = [] + self.submitted_stores: list[TransferJob] = [] - 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: + 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 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 +160,24 @@ 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.LOOKUP_SYNC_DELAY].labelnames == ("tier",) + assert metrics[TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY].labelnames == ("tier",) + 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(): @@ -261,12 +298,15 @@ class TestTieringOffloadingManager: 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 ( @@ -285,12 +325,15 @@ class TestTieringOffloadingManager: 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 ( @@ -473,7 +516,7 @@ class TestTieringOffloadingManager: 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 @@ -500,7 +543,7 @@ class TestTieringOffloadingManager: ] == 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)) @@ -508,22 +551,39 @@ class TestTieringOffloadingManager: 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 @@ -531,13 +591,30 @@ class TestTieringOffloadingManager: # 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() @@ -547,38 +624,51 @@ class TestTieringOffloadingManager: 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.""" @@ -938,9 +1028,7 @@ class TestTieringOffloadingManager: 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) @@ -996,13 +1084,13 @@ class TestTieringOffloadingManager: 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). @@ -1020,9 +1108,7 @@ class TestTieringOffloadingManager: ) 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 @@ -1035,7 +1121,7 @@ class TestTieringOffloadingManager: 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 @@ -1063,18 +1149,18 @@ class TestTieringOffloadingManager: 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 2075b556fb..768d7e328a 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -39,10 +39,25 @@ 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" + 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 -class JobMetadata: +class TransferJob: """Metadata for an in-flight async transfer job.""" job_id: JobId @@ -63,6 +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_time: float | None = None class ParentManager(ABC): @@ -96,7 +112,7 @@ class ParentManager(ABC): self, keys: Collection[OffloadKey], req_context: ReqContext, - ) -> JobMetadata: ... + ) -> TransferJob: ... @abstractmethod def on_request_finished(self, req_context: ReqContext) -> None: ... @@ -153,7 +169,7 @@ class SecondaryTierManager(ABC): 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. @@ -181,7 +197,7 @@ class SecondaryTierManager(ABC): 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 1a78a3feb7..b2652a6d55 100644 --- a/vllm/v1/kv_offload/tiering/example/manager.py +++ b/vllm/v1/kv_offload/tiering/example/manager.py @@ -23,9 +23,9 @@ from vllm.v1.kv_offload.base import ( RequestOffloadingContext, ) from vllm.v1.kv_offload.tiering.base import ( - JobMetadata, JobResult, SecondaryTierManager, + TransferJob, ) logger = logging.getLogger(__name__) @@ -73,6 +73,8 @@ class ExampleSecondaryTierManager(SecondaryTierManager): # 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: @@ -89,7 +91,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager): 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. @@ -106,10 +108,16 @@ class ExampleSecondaryTierManager(SecondaryTierManager): 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_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. @@ -131,7 +139,13 @@ class ExampleSecondaryTierManager(SecondaryTierManager): ) 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_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 560e41093a..dae0d8d500 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -43,11 +43,11 @@ from vllm.v1.kv_offload.file_mapper import FileMapper 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, @@ -214,12 +214,13 @@ class FileSystemTierManager(SecondaryTierManager): 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] = 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, @@ -228,7 +229,7 @@ class FileSystemTierManager(SecondaryTierManager): 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 +272,7 @@ class FileSystemTierManager(SecondaryTierManager): """Collect finished jobs; a failed promotion marks only its failed keys as a miss here (scheduler thread).""" results = [] - 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: @@ -298,10 +299,17 @@ class FileSystemTierManager(SecondaryTierManager): job_id=job_id, success=False, successful_keys=tuple(successful) if successful else None, + transfer_time=transfer_time, ) ) - else: - results.append(JobResult(job_id=job_id, success=success)) + continue + results.append( + JobResult( + job_id=job_id, + success=success, + 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 9bf8fe508f..9db0629192 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 @@ Thread pool: """ 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 @@ class DualQueueThreadPool: 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 @@ class DualQueueThreadPool: 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 @@ class DualQueueThreadPool: 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 5b8f92a458..f8e89b0e6d 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -23,6 +23,7 @@ Key Design Principles: 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,12 +49,12 @@ from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager 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, ) +from vllm.v1.kv_offload.tiering.metrics import TieringMetricsTracker logger = init_logger(__name__) @@ -72,11 +73,12 @@ class RequestState: req_context: ReqContext pending_primary_stores: int = 0 is_finished: bool = False - request_level_tiers: set[SecondaryTierManager] | None = None - 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 + request_level_tiers: set[int] | None = None + + +class JobMetadata(NamedTuple): + transfer_job: TransferJob + tier_idx: int class CPUPrimaryTierOffloadingManager(CPUOffloadingManager): @@ -132,31 +134,33 @@ class CPUPrimaryTierOffloadingManager(CPUOffloadingManager): 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") + __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_idx=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_idx=self._origin_idx) def create_store_job( self, keys: Collection[OffloadKey], req_context: ReqContext - ) -> JobMetadata: - return self._m.create_store_job(keys, req_context) + ) -> 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) + return self._m.on_request_finished( + req_context, exclude_tier_idx=self._origin_idx + ) class TieringOffloadingManager(OffloadingManager): @@ -193,18 +197,23 @@ class TieringOffloadingManager(OffloadingManager): 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._jobs: dict[JobId, JobMetadata] = {} + primary_view = self.primary_tier.get_kv_memoryview() + assert primary_view.strides is not None + 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(). - # 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(). @@ -217,13 +226,13 @@ class TieringOffloadingManager(OffloadingManager): # 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_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() + @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.""" @@ -231,6 +240,14 @@ class TieringOffloadingManager(OffloadingManager): self._job_id_counter += 1 return job_id + def _register_job(self, transfer_job: TransferJob, tier_idx: int) -> None: + job_metadata = JobMetadata(transfer_job, tier_idx) + self._jobs[transfer_job.job_id] = job_metadata + self._metrics.on_job_registered(job_metadata) + + def _pop_job(self, job_id: JobId) -> JobMetadata | None: + return self._jobs.pop(job_id, None) + def _maybe_process_finished_jobs(self): """ Poll secondary tiers for completed jobs (at most once per step). @@ -247,13 +264,14 @@ class TieringOffloadingManager(OffloadingManager): 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" @@ -262,18 +280,18 @@ class TieringOffloadingManager(OffloadingManager): 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, ) @@ -291,13 +309,19 @@ class TieringOffloadingManager(OffloadingManager): 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) + 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" ) + 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._metrics.on_job_finished(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) @@ -305,7 +329,7 @@ class TieringOffloadingManager(OffloadingManager): # 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 @@ -314,7 +338,7 @@ class TieringOffloadingManager(OffloadingManager): 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. @@ -343,73 +367,58 @@ class TieringOffloadingManager(OffloadingManager): # in time for a promotion this lookup may initiate. self._maybe_process_finished_jobs() - req_state = self._req_state.get(req_context.req_id) - + start_time = time.monotonic() primary_hit = self.primary_tier.lookup(key, req_context) + lookup_duration = time.monotonic() - start_time + self._metrics.on_lookup( + req_context, + key, + self._metrics.primary_tier_label, + primary_hit, + lookup_duration, + ) 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 tier in self.secondary_tiers: - if tier is exclude_tier: + for i, tier in enumerate(self.secondary_tiers): + if i == exclude_tier_idx: continue if not req_context.load_tier_filter.allows(tier.medium, tier.locality): continue + 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(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, + 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 + self._metrics.on_lookup( + req_context, + key, + labelvalues, + result, + lookup_duration, + ) - 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, + tier_idx: int, key: OffloadKey, req_context: ReqContext, ) -> bool: @@ -423,7 +432,7 @@ class TieringOffloadingManager(OffloadingManager): 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(). @@ -439,13 +448,14 @@ class TieringOffloadingManager(OffloadingManager): if primary_write_result is None: # Primary tier is full; caller should treat the block as unavailable # rather than retrying indefinitely. + self._metrics.on_promotion_allocation_failure() return False store_spec = primary_write_result.store_spec 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( @@ -465,17 +475,18 @@ class TieringOffloadingManager(OffloadingManager): if not self._pending_load_submissions: return - for tier, pending_by_ctx in self._pending_load_submissions.items(): + 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 = 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._register_job(job_metadata, tier_idx) tier.submit_load(job_metadata) self._pending_load_submissions.clear() @@ -594,7 +605,7 @@ class TieringOffloadingManager(OffloadingManager): 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 @@ -609,8 +620,9 @@ class TieringOffloadingManager(OffloadingManager): if not ready_keys: return - for tier in request_level_tiers: - job_metadata = self.create_store_job(ready_keys, req_context) + 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 @@ -647,8 +659,8 @@ class TieringOffloadingManager(OffloadingManager): # 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) + 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 @@ -663,12 +675,13 @@ class TieringOffloadingManager(OffloadingManager): self, keys: Collection[OffloadKey], req_context: ReqContext, - ) -> JobMetadata: + tier_idx: int = 0, + ) -> 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(). @@ -676,14 +689,14 @@ class TieringOffloadingManager(OffloadingManager): 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 + self._register_job(job_metadata, tier_idx) return job_metadata @override @@ -691,7 +704,7 @@ class TieringOffloadingManager(OffloadingManager): 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. @@ -700,14 +713,15 @@ class TieringOffloadingManager(OffloadingManager): Only stores REQUEST_LEVEL tier decisions for use in prepare_store. """ state = RequestState(req_context=req_context) - for tier in self.secondary_tiers: - if tier is exclude_tier: + self._metrics.on_new_request(req_context) + 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 = ( @@ -722,17 +736,17 @@ class TieringOffloadingManager(OffloadingManager): 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. @@ -746,12 +760,11 @@ class TieringOffloadingManager(OffloadingManager): 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._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 @@ -782,15 +795,14 @@ class TieringOffloadingManager(OffloadingManager): state = self._req_state.get(req_id) if state is None: continue - 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: # 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 ) @@ -825,11 +837,13 @@ class TieringOffloadingManager(OffloadingManager): # 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() + self._metrics.assert_idle() finished_req_ids = [] for req_id, state in self._req_state.items(): @@ -838,8 +852,7 @@ class TieringOffloadingManager(OffloadingManager): 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() @@ -855,6 +868,13 @@ class TieringOffloadingManager(OffloadingManager): if stats is not None and stats.is_empty(): stats = None + metrics_stats = self._metrics.take_stats() + if metrics_stats is not None: + if stats is None: + stats = metrics_stats + else: + stats.aggregate(metrics_stats) + for tier in self.secondary_tiers: tier_stats = tier.get_stats() if tier_stats is None or tier_stats.is_empty(): @@ -864,13 +884,6 @@ class TieringOffloadingManager(OffloadingManager): 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/metrics.py b/vllm/v1/kv_offload/tiering/metrics.py new file mode 100644 index 0000000000..60907142ac --- /dev/null +++ b/vllm/v1/kv_offload/tiering/metrics.py @@ -0,0 +1,269 @@ +# 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, dict[OffloadKey, float | None]] | None = field( + default_factory=dict + ) + + +@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 + + def on_request_finished(self, req_context: ReqContext) -> None: + self._request_states.pop(req_context.req_id, None) + + def on_lookup( + self, + req_context: ReqContext, + key: OffloadKey, + tier_label: TierLabel, + result: LookupResult, + lookup_duration: float, + ) -> None: + state = self._request_states.get(req_context.req_id) + 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, + LookupResult.MISS, + ): + 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, + lookup_duration, + start_time, + ) + elif result is LookupResult.RETRY and state.observed_lookups is not None: + observed = state.observed_lookups.setdefault(tier_label, {}) + observed.setdefault(key, time.monotonic() - lookup_duration) + + 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) + completed_key_count = len(transfer_job.keys) + 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) + 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 + 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 = 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( + 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 _observe_resolved_lookup( + self, + tier_label: TierLabel, + result: LookupResult, + lookup_duration: 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, + ) + self._stats.observe_histogram( + TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, + lookup_duration, + tier_label, + ) + if async_start_time is not None: + self._stats.observe_histogram( + TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY, + time.monotonic() - async_start_time, + tier_label, + ) diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index fcfea6909f..fa055c2f0a 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.file_mapper import FileMapper 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 @@ -267,7 +267,11 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): 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, + ) def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: result = self._lookup_manager.lookup(key, req_context) @@ -275,7 +279,7 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): 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) @@ -283,7 +287,7 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): 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( @@ -317,6 +321,11 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): 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: @@ -341,7 +350,13 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): 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_time=transfer_time, + ) + ) 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 f1f557ce38..f76ed59b5c 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.base import ( ) 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 @@ class P2PSecondaryTierManager(SecondaryTierManager): 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 @@ -470,7 +470,7 @@ class P2PSecondaryTierManager(SecondaryTierManager): ) @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 ccce426bfd..d1d84c356e 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -61,6 +61,8 @@ from typing_extensions import override from vllm.logger import init_logger from vllm.v1.kv_offload.base import ( CanonicalKVCaches, + OffloadingCounterMetadata, + OffloadingGaugeMetadata, OffloadingHistogramMetadata, OffloadingManager, OffloadingMetricMetadata, @@ -103,10 +105,10 @@ class TieringOffloadingSpec(CPUOffloadingSpec): 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, @@ -125,10 +127,11 @@ class TieringOffloadingSpec(CPUOffloadingSpec): 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, @@ -144,6 +147,100 @@ class TieringOffloadingSpec(CPUOffloadingSpec): ), ) ) + metrics[TieringOffloadingMetrics.READ_BYTES] = OffloadingCounterMetadata( + documentation=( + "Total bytes read from secondary tiers into the primary tier, " + "labeled by tier." + ), + labelnames=("tier",), + ) + metrics[TieringOffloadingMetrics.READ_TIME] = OffloadingCounterMetadata( + documentation=( + "Total time spent reading from secondary tiers into the primary " + "tier, in seconds, labeled by tier." + ), + labelnames=("tier",), + ) + metrics[TieringOffloadingMetrics.WRITE_BYTES] = OffloadingCounterMetadata( + documentation=( + "Total bytes written from the primary tier to secondary tiers, " + "labeled by tier." + ), + labelnames=("tier",), + ) + metrics[TieringOffloadingMetrics.WRITE_TIME] = OffloadingCounterMetadata( + documentation=( + "Total time spent writing from the primary tier to secondary " + "tiers, in seconds, labeled by tier." + ), + 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=( + "Number of promotion attempts that failed because the " + "primary tier could not allocate space." + ), + ) + ) + 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")