| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Duration and in-flight metrics for RL weight-transfer operations.""" |
| 4 | |
| 5 | from collections.abc import Iterator |
| 6 | from contextlib import contextmanager |
| 7 | from threading import Lock |
| 8 | from time import perf_counter |
| 9 | from typing import Literal |
| 10 | |
| 11 | from prometheus_client import REGISTRY, CollectorRegistry, Gauge, Histogram |
| 12 | |
| 13 | Operation = Literal["init", "start", "start_draft", "update", "finish"] |
| 14 | |
| 15 | |
| 16 | class WeightOperationMetrics: |
| 17 | def __init__(self, registry: CollectorRegistry = REGISTRY): |
| 18 | self.duration = Histogram( |
| 19 | "vllm:rl_weight_update_operation_duration_seconds", |
| 20 | "Duration of one weight-transfer operation.", |
| 21 | ["operation"], |
| 22 | buckets=(0.01, 0.1, 1, 10, 30, 60, 120, 300, 600), |
| 23 | registry=registry, |
| 24 | ) |
| 25 | self.in_flight = Gauge( |
| 26 | "vllm:rl_weight_update_operations_in_flight", |
| 27 | "Weight-transfer operations currently awaited.", |
| 28 | ["operation"], |
| 29 | multiprocess_mode="livesum", |
| 30 | registry=registry, |
| 31 | ) |
| 32 | |
| 33 | @contextmanager |
| 34 | def record(self, operation: Operation) -> Iterator[None]: |
| 35 | in_flight = self.in_flight.labels(operation) |
| 36 | in_flight.inc() |
| 37 | started = perf_counter() |
| 38 | try: |
| 39 | yield |
| 40 | finally: |
| 41 | in_flight.dec() |
| 42 | self.duration.labels(operation).observe(perf_counter() - started) |
| 43 | |
| 44 | |
| 45 | _metrics: WeightOperationMetrics | None = None |
| 46 | _metrics_lock = Lock() |
| 47 | |
| 48 | |
| 49 | def weight_operation_metrics() -> WeightOperationMetrics: |
| 50 | # Lazy, so the collectors exist only after multiprocess Prometheus setup. |
| 51 | global _metrics |
| 52 | if _metrics is None: |
| 53 | with _metrics_lock: |
| 54 | if _metrics is None: |
| 55 | _metrics = WeightOperationMetrics() |
| 56 | return _metrics |