diff --git a/rust/src/metrics/src/api_server.rs b/rust/src/metrics/src/api_server.rs index 1dc7154998..51599ee51e 100644 --- a/rust/src/metrics/src/api_server.rs +++ b/rust/src/metrics/src/api_server.rs @@ -23,6 +23,10 @@ fn http_request_duration_highr_histogram() -> Histogram { Histogram::new(HTTP_REQUEST_DURATION_HIGHR_BUCKETS.iter().copied()) } +fn weight_operation_duration_histogram() -> Histogram { + Histogram::new([0.01, 0.1, 1.0, 10.0, 30.0, 60.0, 120.0, 300.0, 600.0]) +} + fn sleep_mode_operation_duration_histogram() -> Histogram { Histogram::new([0.01, 0.1, 1.0, 10.0, 30.0, 60.0, 120.0, 300.0, 600.0]) } @@ -40,6 +44,11 @@ pub struct HttpHandlerLabels { pub handler: String, } +#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)] +pub struct WeightOperationLabels { + pub operation: &'static str, +} + #[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)] pub struct SleepModeOperationLabels { pub operation: &'static str, @@ -48,6 +57,9 @@ pub struct SleepModeOperationLabels { pub(crate) type HttpRequestCounterFamily = Family; pub(crate) type HttpHandlerHistogramFamily = Family Histogram>; +pub(crate) type WeightOperationHistogramFamily = + Family Histogram>; +pub(crate) type WeightOperationGaugeFamily = Family; pub(crate) type SleepModeOperationHistogramFamily = Family Histogram>; pub(crate) type SleepModeOperationGaugeFamily = Family; @@ -57,6 +69,8 @@ pub struct ApiServerMetrics { pub http_requests: HttpRequestCounterFamily, pub http_request_duration_seconds: HttpHandlerHistogramFamily, pub http_request_duration_highr_seconds: Histogram, + pub weight_operation_duration_seconds: WeightOperationHistogramFamily, + pub weight_operations_in_flight: WeightOperationGaugeFamily, pub sleep_mode_operation_duration_seconds: SleepModeOperationHistogramFamily, pub sleep_mode_operations_in_flight: SleepModeOperationGaugeFamily, } @@ -87,6 +101,21 @@ impl ApiServerMetrics { http_request_duration_highr_seconds.clone(), ); + let weight_operation_duration_seconds = + Family::new_with_constructor(weight_operation_duration_histogram as fn() -> Histogram); + registry.register( + "vllm:rl_weight_update_operation_duration_seconds", + "Duration of one weight-transfer operation.", + weight_operation_duration_seconds.clone(), + ); + + let weight_operations_in_flight = WeightOperationGaugeFamily::default(); + registry.register( + "vllm:rl_weight_update_operations_in_flight", + "Weight-transfer operations currently awaited.", + weight_operations_in_flight.clone(), + ); + let sleep_mode_operation_duration_seconds = Family::new_with_constructor( sleep_mode_operation_duration_histogram as fn() -> Histogram, ); @@ -107,11 +136,25 @@ impl ApiServerMetrics { http_requests, http_request_duration_seconds, http_request_duration_highr_seconds, + weight_operation_duration_seconds, + weight_operations_in_flight, sleep_mode_operation_duration_seconds, sleep_mode_operations_in_flight, } } + /// Track one weight-transfer operation until the returned guard is dropped. + pub fn record_weight_operation(&self, operation: &'static str) -> WeightOperationRecorder { + let labels = WeightOperationLabels { operation }; + let in_flight = self.weight_operations_in_flight.get_or_create_owned(&labels); + in_flight.inc(); + WeightOperationRecorder { + started_at: Instant::now(), + duration: self.weight_operation_duration_seconds.get_or_create_owned(&labels), + in_flight, + } + } + /// Record one dispatched sleep-mode engine operation. pub fn record_sleep_mode_operation( &self, @@ -128,6 +171,19 @@ impl ApiServerMetrics { } } +pub struct WeightOperationRecorder { + started_at: Instant, + duration: Histogram, + in_flight: U64Gauge, +} + +impl Drop for WeightOperationRecorder { + fn drop(&mut self) { + self.in_flight.dec(); + self.duration.observe(self.started_at.elapsed().as_secs_f64()); + } +} + /// Records duration and clears in-flight tracking on completion or cancellation. pub struct SleepModeOperationRecorder { started_at: Instant, @@ -194,4 +250,25 @@ mod tests { ))); } } + + #[test] + fn weight_operation_recorder_tracks_in_flight_and_duration() { + let mut registry = Registry::default(); + let metrics = ApiServerMetrics::register(&mut registry); + + let recorder = metrics.record_weight_operation("update"); + assert!( + rendered_metrics(®istry) + .contains("vllm:rl_weight_update_operations_in_flight{operation=\"update\"} 1") + ); + drop(recorder); + + let after = rendered_metrics(®istry); + assert!( + after.contains("vllm:rl_weight_update_operations_in_flight{operation=\"update\"} 0") + ); + assert!(after.contains( + "vllm:rl_weight_update_operation_duration_seconds_count{operation=\"update\"} 1" + )); + } } diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index d229be25fa..7225f2bebf 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -6723,6 +6723,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() { }) }) .await; + let metrics_before = METRICS.render().unwrap(); let mut responses = Vec::new(); for (method, path, body) in [ @@ -6785,6 +6786,26 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() { ); } engine_task.await.expect("mock engine task"); + let metrics_after = METRICS.render().unwrap(); + for operation in ["update", "finish"] { + assert_eq!( + metric_delta( + &metrics_before, + &metrics_after, + "vllm:rl_weight_update_operation_duration_seconds_count", + Some(&format!("operation=\"{operation}\"")), + ), + 2.0 + ); + assert_eq!( + metric_value( + &metrics_after, + "vllm:rl_weight_update_operations_in_flight", + Some(&format!("operation=\"{operation}\"")), + ), + Some(0.0) + ); + } expect_test::expect![[r#" [ "{\"message\":\"Weight transfer initialized\"}", @@ -6811,6 +6832,7 @@ async fn weight_transfer_routes_reject_invalid_payloads_before_engine_calls() { }) }) .await; + let metrics_before = METRICS.render().unwrap(); for (path, body) in [ ("/init_weight_transfer_engine", "{"), @@ -6840,6 +6862,18 @@ async fn weight_transfer_routes_reject_invalid_payloads_before_engine_calls() { .expect("call app"); assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{path}: {body}"); } + let metrics_after = METRICS.render().unwrap(); + for operation in ["init", "update", "finish"] { + assert_eq!( + metric_delta( + &metrics_before, + &metrics_after, + "vllm:rl_weight_update_operation_duration_seconds_count", + Some(&format!("operation=\"{operation}\"")), + ), + 0.0 + ); + } engine_task.abort_and_join().await; } diff --git a/rust/src/server/src/routes/weight_transfer.rs b/rust/src/server/src/routes/weight_transfer.rs index df55e25593..5ccd66e39b 100644 --- a/rust/src/server/src/routes/weight_transfer.rs +++ b/rust/src/server/src/routes/weight_transfer.rs @@ -8,6 +8,7 @@ use axum::extract::State; use axum::extract::rejection::JsonRejection; use serde::{Deserialize, Serialize}; use serde_json::Value as JsonValue; +use vllm_metrics::METRICS; use crate::error::{ApiError, invalid_request}; use crate::state::AppState; @@ -62,6 +63,7 @@ pub async fn init_weight_transfer_engine( ) })?; + let _recorder = METRICS.api_server.record_weight_operation("init"); state .engine_core_client() .init_weight_transfer_engine(init_info) @@ -77,6 +79,7 @@ pub async fn init_weight_transfer_engine( pub async fn start_weight_update( State(state): State>, ) -> Result, ApiError> { + let _recorder = METRICS.api_server.record_weight_operation("start"); state .engine_core_client() .start_weight_update() @@ -92,6 +95,7 @@ pub async fn start_weight_update( pub async fn start_draft_weight_update( State(state): State>, ) -> Result, ApiError> { + let _recorder = METRICS.api_server.record_weight_operation("start_draft"); state .engine_core_client() .start_draft_weight_update() @@ -122,6 +126,7 @@ pub async fn update_weights( ) })?; + let _recorder = METRICS.api_server.record_weight_operation("update"); state .engine_core_client() .update_weights(update_info) @@ -142,6 +147,7 @@ pub async fn finish_weight_update( let request = body?.map(|Json(request)| request).unwrap_or_default(); let client = state.engine_core_client(); + let _recorder = METRICS.api_server.record_weight_operation("finish"); client .finish_weight_update() .await diff --git a/tests/entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py b/tests/entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py new file mode 100644 index 0000000000..6c7cb60bbb --- /dev/null +++ b/tests/entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Weight-transfer metrics scraped from a real two-API-server deployment.""" + +from pathlib import Path + +import requests +import torch +from prometheus_client.multiprocess import MultiProcessCollector +from prometheus_client.parser import text_string_to_metric_families +from transformers import AutoModelForCausalLM + +from tests.entrypoints.serve.dev.rlhf.conftest import MODEL_NAME +from tests.utils import RemoteOpenAIServer +from vllm.distributed.weight_transfer import ( + HTTPVLLMWeightSyncClient, + ModuleSource, + WeightTransferTrainerFactory, +) +from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo + +DURATION = "vllm:rl_weight_update_operation_duration_seconds" +IN_FLIGHT = "vllm:rl_weight_update_operations_in_flight" +# Enough requests that the kernel spreads them over both API servers. +SYNCS = 8 + + +def scrape(url: str) -> dict[tuple[str, str], float]: + text = requests.get(f"{url}/metrics", timeout=30).text + return { + (sample.name, sample.labels["operation"]): sample.value + for family in text_string_to_metric_families(text) + for sample in family.samples + if sample.name in (f"{DURATION}_count", IN_FLIGHT) + } + + +def pids_that_recorded(multiproc_dir: Path) -> set[str]: + return { + path.stem.rsplit("_", 1)[1] + for path in multiproc_dir.glob("histogram_*.db") + if any( + metric.name == DURATION + for metric in MultiProcessCollector.merge([str(path)]) + ) + } + + +def test_weight_sync_metrics_aggregate_across_api_servers(tmp_path): + args = [ + "--load-format", + "dummy", + "--enforce-eager", + "--max-model-len", + "512", + "--gpu-memory-utilization", + "0.5", + "--api-server-count", + "2", + "--weight-transfer-config", + '{"backend": "ipc"}', + ] + env = { + "VLLM_SERVER_DEV_MODE": "1", + "VLLM_ALLOW_INSECURE_SERIALIZATION": "1", + "PROMETHEUS_MULTIPROC_DIR": str(tmp_path), + } + with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env) as server: + url = server.url_root + trainer = AutoModelForCausalLM.from_pretrained( + MODEL_NAME, torch_dtype=torch.bfloat16 + ).cuda() + engine = WeightTransferTrainerFactory.trainer_init( + IPCTrainerInitInfo(rank=0), + client=HTTPVLLMWeightSyncClient(url), + source=ModuleSource(trainer), + ) + for _ in range(SYNCS): + engine.send_weights() + metrics = scrape(url) + + calls = {"init": 1, "start": SYNCS, "update": SYNCS, "finish": SYNCS} + for operation, count in calls.items(): + assert metrics[(f"{DURATION}_count", operation)] == count, operation + assert metrics[(IN_FLIGHT, operation)] == 0, operation + # The counts above are sums over processes only if both servers recorded. + assert len(pids_that_recorded(tmp_path)) == 2 diff --git a/tests/entrypoints/serve/lora/test_serving_models.py b/tests/entrypoints/serve/lora/test_serving_models.py index b2aefff219..2fb39b53ab 100644 --- a/tests/entrypoints/serve/lora/test_serving_models.py +++ b/tests/entrypoints/serve/lora/test_serving_models.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from http import HTTPStatus -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import MagicMock import pytest @@ -29,7 +29,9 @@ LORA_UNLOADING_SUCCESS_MESSAGE = ( ) -async def _async_serving_models_init() -> OpenAIServingModels: +async def _async_serving_models_init_with_mock() -> tuple[ + OpenAIServingModels, MagicMock +]: mock_engine_client = MagicMock(spec=EngineClient) # Set the max_model_len attribute to avoid missing attribute mock_model_config = MagicMock(spec=ModelConfig) @@ -45,6 +47,13 @@ async def _async_serving_models_init() -> OpenAIServingModels: ) await serving_models.init_static_loras() + # The mock is returned separately: `engine_client` is typed as the protocol, + # so assertion helpers are not visible through it. + return serving_models, mock_engine_client + + +async def _async_serving_models_init() -> OpenAIServingModels: + serving_models, _ = await _async_serving_models_init_with_mock() return serving_models @@ -103,7 +112,7 @@ async def test_load_lora_adapter_duplicate(): @pytest.mark.parametrize("load_inplace", [False, True]) async def test_load_lora_adapter_base_model_name(load_inplace: bool): """A LoRA named after a served model would shadow it for all requests.""" - serving_models = await _async_serving_models_init() + serving_models, mock_engine_client = await _async_serving_models_init_with_mock() request = LoadLoRAAdapterRequest( lora_name=MODEL_NAME, lora_path="/path/to/adapter", load_inplace=load_inplace ) @@ -112,8 +121,7 @@ async def test_load_lora_adapter_base_model_name(load_inplace: bool): assert response.error.type == "InvalidUserInput" assert response.error.code == HTTPStatus.BAD_REQUEST assert len(serving_models.lora_requests) == 0 - assert isinstance(serving_models.engine_client.add_lora, AsyncMock) - serving_models.engine_client.add_lora.assert_not_called() + mock_engine_client.add_lora.assert_not_called() @pytest.mark.asyncio diff --git a/tests/entrypoints/unit_tests/test_rl_weight_operation_metrics.py b/tests/entrypoints/unit_tests/test_rl_weight_operation_metrics.py new file mode 100644 index 0000000000..bdddf3d247 --- /dev/null +++ b/tests/entrypoints/unit_tests/test_rl_weight_operation_metrics.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import asyncio +import os +import subprocess +import sys +from contextlib import nullcontext +from pathlib import Path +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from prometheus_client import CollectorRegistry + +from vllm.entrypoints.serve.dev.rlhf import api_router +from vllm.entrypoints.serve.dev.rlhf import metrics as rlhf_metrics +from vllm.entrypoints.serve.dev.rlhf.metrics import WeightOperationMetrics +from vllm.entrypoints.serve.exception_handling.register import init_exception_handler +from vllm.v1.metrics.prometheus import get_prometheus_registry + +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] + +DURATION = "vllm:rl_weight_update_operation_duration_seconds_count" +IN_FLIGHT = "vllm:rl_weight_update_operations_in_flight" + + +def sample(registry, name, operation): + return registry.get_sample_value(name, {"operation": operation}) + + +class Request: + def __init__(self, engine): + self.app = SimpleNamespace(state=SimpleNamespace(engine_client=engine)) + + +class Engine: + """Holds every engine call open until the test releases it.""" + + def __init__(self): + self.started = asyncio.Event() + self.release = asyncio.Event() + + def __getattr__(self, name): + async def call(*args, **kwargs): + self.started.set() + await self.release.wait() + + return call + + +@pytest.fixture +def registry(monkeypatch): + registry = CollectorRegistry() + monkeypatch.setattr(rlhf_metrics, "_metrics", WeightOperationMetrics(registry)) + return registry + + +@pytest.mark.parametrize("fail", [False, True]) +def test_record_releases_in_flight_and_observes_duration(registry, fail): + metrics = rlhf_metrics.weight_operation_metrics() + with ( + pytest.raises(RuntimeError) if fail else nullcontext(), + metrics.record("update"), + ): + assert sample(registry, IN_FLIGHT, "update") == 1 + if fail: + raise RuntimeError + assert sample(registry, IN_FLIGHT, "update") == 0 + assert sample(registry, DURATION, "update") == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "operation,route,body", + [ + ("init", api_router.init_weight_transfer_engine, {"init_info": {}}), + ("start", api_router.start_weight_update, None), + ("start_draft", api_router.start_draft_weight_update, None), + ("update", api_router.update_weights, {"update_info": {}}), + ("finish", api_router.finish_weight_update, None), + ], +) +async def test_route_records_its_operation(registry, operation, route, body): + engine = Engine() + task = asyncio.create_task(route(Request(engine), **(body or {}))) + await engine.started.wait() + assert sample(registry, IN_FLIGHT, operation) == 1 + engine.release.set() + await task + assert sample(registry, IN_FLIGHT, operation) == 0 + assert sample(registry, DURATION, operation) == 1 + + +@pytest.mark.asyncio +async def test_cancelled_route_releases_in_flight(registry): + engine = Engine() + task = asyncio.create_task( + api_router.update_weights(Request(engine), update_info={}) + ) + await engine.started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert sample(registry, IN_FLIGHT, "update") == 0 + assert sample(registry, DURATION, "update") == 1 + + +HOLD_OPERATION = """ +import sys +from vllm.entrypoints.serve.dev.rlhf.metrics import weight_operation_metrics +with weight_operation_metrics().record("update"): + print("entered", flush=True) + sys.stdin.readline() +""" + + +def test_in_flight_is_summed_across_api_server_processes(tmp_path, monkeypatch): + """Two API-server processes each hold an operation; /metrics sums them.""" + monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) + servers = [ + subprocess.Popen( + [sys.executable, "-c", HOLD_OPERATION], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + text=True, + env=os.environ.copy(), + ) + for _ in range(2) + ] + try: + for server in servers: + assert server.stdout is not None + assert server.stdout.readline().strip() == "entered" + assert sample(get_prometheus_registry(), IN_FLIGHT, "update") == 2 + finally: + for server in servers: + server.communicate("\n", timeout=60) + + scraped = get_prometheus_registry() + assert sample(scraped, IN_FLIGHT, "update") == 0 + assert sample(scraped, DURATION, "update") == 2 + + +def test_rejected_request_is_not_recorded(registry): + app = FastAPI() + app.include_router(api_router.router) + init_exception_handler(app) + app.state.engine_client = Engine() + app.state.args = SimpleNamespace(log_error_stack=False) + response = TestClient(app).post("/update_weights", json={}) + assert response.status_code == 400 + assert sample(registry, DURATION, "update") is None + + +def test_default_recorder_is_created_lazily_on_the_default_registry(): + """The default recorder registers on the process registry only when first used. + + Runs in a fresh interpreter: the registry is process-global, so a route test + elsewhere in the session may already have registered these collectors, and + "not registered yet" would then be unobservable. + """ + probe = """ +from prometheus_client import REGISTRY + +from vllm.entrypoints.serve.dev.rlhf import metrics as rlhf_metrics + +NAMES = [ + "vllm:rl_weight_update_operation_duration_seconds", + "vllm:rl_weight_update_operations_in_flight", +] +assert rlhf_metrics._metrics is None +assert not set(NAMES) & set(REGISTRY._names_to_collectors) + +metrics = rlhf_metrics.weight_operation_metrics() +assert REGISTRY._names_to_collectors[NAMES[0]] is metrics.duration +assert REGISTRY._names_to_collectors[NAMES[1]] is metrics.in_flight +assert rlhf_metrics.weight_operation_metrics() is metrics +print("lazy registration ok") +""" + result = subprocess.run( + [sys.executable, "-c", probe], + cwd=Path(__file__).parents[3], + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + assert "lazy registration ok" in result.stdout diff --git a/vllm/entrypoints/serve/dev/rlhf/api_router.py b/vllm/entrypoints/serve/dev/rlhf/api_router.py index 9c464e6a11..7e52e51217 100644 --- a/vllm/entrypoints/serve/dev/rlhf/api_router.py +++ b/vllm/entrypoints/serve/dev/rlhf/api_router.py @@ -14,6 +14,7 @@ from vllm.distributed.weight_transfer.base import ( WeightTransferUpdateRequest, ) from vllm.engine.protocol import EngineClient +from vllm.entrypoints.serve.dev.rlhf.metrics import weight_operation_metrics from vllm.logger import init_logger from vllm.v1.engine import PauseMode @@ -163,21 +164,24 @@ async def init_weight_transfer_engine( raw_request: Request, init_info: Annotated[dict[str, Any], Body(embed=True)], ): - await engine_client(raw_request).init_weight_transfer_engine( - WeightTransferInitRequest(init_info=init_info) - ) + with weight_operation_metrics().record("init"): + await engine_client(raw_request).init_weight_transfer_engine( + WeightTransferInitRequest(init_info=init_info) + ) return JSONResponse(content={"message": "Weight transfer initialized"}) @router.post("/start_weight_update") async def start_weight_update(raw_request: Request): - await engine_client(raw_request).start_weight_update() + with weight_operation_metrics().record("start"): + await engine_client(raw_request).start_weight_update() return JSONResponse(content={"message": "Weight update started"}) @router.post("/start_draft_weight_update") async def start_draft_weight_update(raw_request: Request): - await engine_client(raw_request).start_draft_weight_update() + with weight_operation_metrics().record("start_draft"): + await engine_client(raw_request).start_draft_weight_update() return JSONResponse(content={"message": "Draft weight update started"}) @@ -186,9 +190,10 @@ async def update_weights( raw_request: Request, update_info: Annotated[WeightTransferUpdatePayload, Body(embed=True)], ): - await engine_client(raw_request).update_weights( - request=WeightTransferUpdateRequest(update_info=update_info) - ) + with weight_operation_metrics().record("update"): + await engine_client(raw_request).update_weights( + request=WeightTransferUpdateRequest(update_info=update_info) + ) return JSONResponse(content={"message": "Weights updated"}) @@ -197,7 +202,8 @@ async def finish_weight_update( raw_request: Request, weight_version: Annotated[str | None, Body(embed=True)] = None, ): - await engine_client(raw_request).finish_weight_update(weight_version) + with weight_operation_metrics().record("finish"): + await engine_client(raw_request).finish_weight_update(weight_version) return JSONResponse(content={"message": "Weight update finished"}) diff --git a/vllm/entrypoints/serve/dev/rlhf/metrics.py b/vllm/entrypoints/serve/dev/rlhf/metrics.py new file mode 100644 index 0000000000..ed16f9a820 --- /dev/null +++ b/vllm/entrypoints/serve/dev/rlhf/metrics.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Duration and in-flight metrics for RL weight-transfer operations.""" + +from collections.abc import Iterator +from contextlib import contextmanager +from threading import Lock +from time import perf_counter +from typing import Literal + +from prometheus_client import REGISTRY, CollectorRegistry, Gauge, Histogram + +Operation = Literal["init", "start", "start_draft", "update", "finish"] + + +class WeightOperationMetrics: + def __init__(self, registry: CollectorRegistry = REGISTRY): + self.duration = Histogram( + "vllm:rl_weight_update_operation_duration_seconds", + "Duration of one weight-transfer operation.", + ["operation"], + buckets=(0.01, 0.1, 1, 10, 30, 60, 120, 300, 600), + registry=registry, + ) + self.in_flight = Gauge( + "vllm:rl_weight_update_operations_in_flight", + "Weight-transfer operations currently awaited.", + ["operation"], + multiprocess_mode="livesum", + registry=registry, + ) + + @contextmanager + def record(self, operation: Operation) -> Iterator[None]: + in_flight = self.in_flight.labels(operation) + in_flight.inc() + started = perf_counter() + try: + yield + finally: + in_flight.dec() + self.duration.labels(operation).observe(perf_counter() - started) + + +_metrics: WeightOperationMetrics | None = None +_metrics_lock = Lock() + + +def weight_operation_metrics() -> WeightOperationMetrics: + # Lazy, so the collectors exist only after multiprocess Prometheus setup. + global _metrics + if _metrics is None: + with _metrics_lock: + if _metrics is None: + _metrics = WeightOperationMetrics() + return _metrics