@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 713ec07655

[Frontend][RL] Track HTTP weight operation outcomes and concurrency (#55781)

Signed-off-by: Ronald1995 <ronaldautomobile@163.com> Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: aoshen02 <aoshen@inferact.ai>

Ronald committed Sep 30, 2026 at 17:19 UTC 713ec07655a3d39daeeff0fd2fcb9dd251c73979

8 files changed +476 -14

rust/src/metrics/src/api_server.rs

+77

@@ -23,6 +23,10 @@ fn http_request_duration_highr_histogram() -> Histogram {

23 Histogram::new(HTTP_REQUEST_DURATION_HIGHR_BUCKETS.iter().copied())

24 }

26 +fn weight_operation_duration_histogram() -> Histogram {

27 + Histogram::new([0.01, 0.1, 1.0, 10.0, 30.0, 60.0, 120.0, 300.0, 600.0])

28 +}

29 +

30 fn sleep_mode_operation_duration_histogram() -> Histogram {

31 Histogram::new([0.01, 0.1, 1.0, 10.0, 30.0, 60.0, 120.0, 300.0, 600.0])

32 }

@@ -40,6 +44,11 @@ pub struct HttpHandlerLabels {

44 pub handler: String,

45 }

47 +#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]

48 +pub struct WeightOperationLabels {

49 + pub operation: &'static str,

50 +}

51 +

52 #[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]

53 pub struct SleepModeOperationLabels {

54 pub operation: &'static str,

@@ -48,6 +57,9 @@ pub struct SleepModeOperationLabels {

57 pub(crate) type HttpRequestCounterFamily = Family<HttpRequestLabels, U64Counter>;

58 pub(crate) type HttpHandlerHistogramFamily =

59 Family<HttpHandlerLabels, Histogram, fn() -> Histogram>;

60 +pub(crate) type WeightOperationHistogramFamily =

61 + Family<WeightOperationLabels, Histogram, fn() -> Histogram>;

62 +pub(crate) type WeightOperationGaugeFamily = Family<WeightOperationLabels, U64Gauge>;

63 pub(crate) type SleepModeOperationHistogramFamily =

64 Family<SleepModeOperationLabels, Histogram, fn() -> Histogram>;

65 pub(crate) type SleepModeOperationGaugeFamily = Family<SleepModeOperationLabels, U64Gauge>;

@@ -57,6 +69,8 @@ pub struct ApiServerMetrics {

69 pub http_requests: HttpRequestCounterFamily,

70 pub http_request_duration_seconds: HttpHandlerHistogramFamily,

71 pub http_request_duration_highr_seconds: Histogram,

72 + pub weight_operation_duration_seconds: WeightOperationHistogramFamily,

73 + pub weight_operations_in_flight: WeightOperationGaugeFamily,

74 pub sleep_mode_operation_duration_seconds: SleepModeOperationHistogramFamily,

75 pub sleep_mode_operations_in_flight: SleepModeOperationGaugeFamily,

76 }

@@ -87,6 +101,21 @@ impl ApiServerMetrics {

101 http_request_duration_highr_seconds.clone(),

102 );

103

104 + let weight_operation_duration_seconds =

105 + Family::new_with_constructor(weight_operation_duration_histogram as fn() -> Histogram);

106 + registry.register(

107 + "vllm:rl_weight_update_operation_duration_seconds",

108 + "Duration of one weight-transfer operation.",

109 + weight_operation_duration_seconds.clone(),

110 + );

111 +

112 + let weight_operations_in_flight = WeightOperationGaugeFamily::default();

113 + registry.register(

114 + "vllm:rl_weight_update_operations_in_flight",

115 + "Weight-transfer operations currently awaited.",

116 + weight_operations_in_flight.clone(),

117 + );

118 +

119 let sleep_mode_operation_duration_seconds = Family::new_with_constructor(

120 sleep_mode_operation_duration_histogram as fn() -> Histogram,

121 );

@@ -107,11 +136,25 @@ impl ApiServerMetrics {

136 http_requests,

137 http_request_duration_seconds,

138 http_request_duration_highr_seconds,

139 + weight_operation_duration_seconds,

140 + weight_operations_in_flight,

141 sleep_mode_operation_duration_seconds,

142 sleep_mode_operations_in_flight,

143 }

144 }

145

146 + /// Track one weight-transfer operation until the returned guard is dropped.

147 + pub fn record_weight_operation(&self, operation: &'static str) -> WeightOperationRecorder {

148 + let labels = WeightOperationLabels { operation };

149 + let in_flight = self.weight_operations_in_flight.get_or_create_owned(&labels);

150 + in_flight.inc();

151 + WeightOperationRecorder {

152 + started_at: Instant::now(),

153 + duration: self.weight_operation_duration_seconds.get_or_create_owned(&labels),

154 + in_flight,

155 + }

156 + }

157 +

158 /// Record one dispatched sleep-mode engine operation.

159 pub fn record_sleep_mode_operation(

160 &self,

@@ -128,6 +171,19 @@ impl ApiServerMetrics {

171 }

172 }

173

174 +pub struct WeightOperationRecorder {

175 + started_at: Instant,

176 + duration: Histogram,

177 + in_flight: U64Gauge,

178 +}

179 +

180 +impl Drop for WeightOperationRecorder {

181 + fn drop(&mut self) {

182 + self.in_flight.dec();

183 + self.duration.observe(self.started_at.elapsed().as_secs_f64());

184 + }

185 +}

186 +

187 /// Records duration and clears in-flight tracking on completion or cancellation.

188 pub struct SleepModeOperationRecorder {

189 started_at: Instant,

@@ -194,4 +250,25 @@ mod tests {

250 )));

251 }

252 }

253 +

254 + #[test]

255 + fn weight_operation_recorder_tracks_in_flight_and_duration() {

256 + let mut registry = Registry::default();

257 + let metrics = ApiServerMetrics::register(&mut registry);

258 +

259 + let recorder = metrics.record_weight_operation("update");

260 + assert!(

261 + rendered_metrics(&registry)

262 + .contains("vllm:rl_weight_update_operations_in_flight{operation=\"update\"} 1")

263 + );

264 + drop(recorder);

265 +

266 + let after = rendered_metrics(&registry);

267 + assert!(

268 + after.contains("vllm:rl_weight_update_operations_in_flight{operation=\"update\"} 0")

269 + );

270 + assert!(after.contains(

271 + "vllm:rl_weight_update_operation_duration_seconds_count{operation=\"update\"} 1"

272 + ));

273 + }

274 }

rust/src/server/src/routes/tests.rs

+34

@@ -6723,6 +6723,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() {

6723 })

6724 })

6725 .await;

6726 + let metrics_before = METRICS.render().unwrap();

6727

6728 let mut responses = Vec::new();

6729 for (method, path, body) in [

@@ -6785,6 +6786,26 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() {

6786 );

6787 }

6788 engine_task.await.expect("mock engine task");

6789 + let metrics_after = METRICS.render().unwrap();

6790 + for operation in ["update", "finish"] {

6791 + assert_eq!(

6792 + metric_delta(

6793 + &metrics_before,

6794 + &metrics_after,

6795 + "vllm:rl_weight_update_operation_duration_seconds_count",

6796 + Some(&format!("operation=\"{operation}\"")),

6797 + ),

6798 + 2.0

6799 + );

6800 + assert_eq!(

6801 + metric_value(

6802 + &metrics_after,

6803 + "vllm:rl_weight_update_operations_in_flight",

6804 + Some(&format!("operation=\"{operation}\"")),

6805 + ),

6806 + Some(0.0)

6807 + );

6808 + }

6809 expect_test::expect![[r#"

6810 [

6811 "{\"message\":\"Weight transfer initialized\"}",

@@ -6811,6 +6832,7 @@ async fn weight_transfer_routes_reject_invalid_payloads_before_engine_calls() {

6832 })

6833 })

6834 .await;

6835 + let metrics_before = METRICS.render().unwrap();

6836

6837 for (path, body) in [

6838 ("/init_weight_transfer_engine", "{"),

@@ -6840,6 +6862,18 @@ async fn weight_transfer_routes_reject_invalid_payloads_before_engine_calls() {

6862 .expect("call app");

6863 assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{path}: {body}");

6864 }

6865 + let metrics_after = METRICS.render().unwrap();

6866 + for operation in ["init", "update", "finish"] {

6867 + assert_eq!(

6868 + metric_delta(

6869 + &metrics_before,

6870 + &metrics_after,

6871 + "vllm:rl_weight_update_operation_duration_seconds_count",

6872 + Some(&format!("operation=\"{operation}\"")),

6873 + ),

6874 + 0.0

6875 + );

6876 + }

6877 engine_task.abort_and_join().await;

6878 }

6879

rust/src/server/src/routes/weight_transfer.rs

@@ -8,6 +8,7 @@ use axum::extract::State;

8 use axum::extract::rejection::JsonRejection;

9 use serde::{Deserialize, Serialize};

10 use serde_json::Value as JsonValue;

11 +use vllm_metrics::METRICS;

13 use crate::error::{ApiError, invalid_request};

14 use crate::state::AppState;

@@ -62,6 +63,7 @@ pub async fn init_weight_transfer_engine(

63 )

64 })?;

66 + let _recorder = METRICS.api_server.record_weight_operation("init");

67 state

68 .engine_core_client()

69 .init_weight_transfer_engine(init_info)

@@ -77,6 +79,7 @@ pub async fn init_weight_transfer_engine(

79 pub async fn start_weight_update(

80 State(state): State<Arc<AppState>>,

81 ) -> Result<Json<MessageResponse>, ApiError> {

82 + let _recorder = METRICS.api_server.record_weight_operation("start");

83 state

84 .engine_core_client()

85 .start_weight_update()

@@ -92,6 +95,7 @@ pub async fn start_weight_update(

95 pub async fn start_draft_weight_update(

96 State(state): State<Arc<AppState>>,

97 ) -> Result<Json<MessageResponse>, ApiError> {

98 + let _recorder = METRICS.api_server.record_weight_operation("start_draft");

99 state

100 .engine_core_client()

101 .start_draft_weight_update()

@@ -122,6 +126,7 @@ pub async fn update_weights(

126 )

127 })?;

128

129 + let _recorder = METRICS.api_server.record_weight_operation("update");

130 state

131 .engine_core_client()

132 .update_weights(update_info)

@@ -142,6 +147,7 @@ pub async fn finish_weight_update(

147 let request = body?.map(|Json(request)| request).unwrap_or_default();

148

149 let client = state.engine_core_client();

150 + let _recorder = METRICS.api_server.record_weight_operation("finish");

151 client

152 .finish_weight_update()

153 .await

tests/entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py new

+87

@@ -0,0 +1,87 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Weight-transfer metrics scraped from a real two-API-server deployment."""

4 +

5 +from pathlib import Path

6 +

7 +import requests

8 +import torch

9 +from prometheus_client.multiprocess import MultiProcessCollector

10 +from prometheus_client.parser import text_string_to_metric_families

11 +from transformers import AutoModelForCausalLM

12 +

13 +from tests.entrypoints.serve.dev.rlhf.conftest import MODEL_NAME

14 +from tests.utils import RemoteOpenAIServer

15 +from vllm.distributed.weight_transfer import (

16 + HTTPVLLMWeightSyncClient,

17 + ModuleSource,

18 + WeightTransferTrainerFactory,

19 +)

20 +from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo

21 +

22 +DURATION = "vllm:rl_weight_update_operation_duration_seconds"

23 +IN_FLIGHT = "vllm:rl_weight_update_operations_in_flight"

24 +# Enough requests that the kernel spreads them over both API servers.

25 +SYNCS = 8

26 +

27 +

28 +def scrape(url: str) -> dict[tuple[str, str], float]:

29 + text = requests.get(f"{url}/metrics", timeout=30).text

30 + return {

31 + (sample.name, sample.labels["operation"]): sample.value

32 + for family in text_string_to_metric_families(text)

33 + for sample in family.samples

34 + if sample.name in (f"{DURATION}_count", IN_FLIGHT)

35 + }

36 +

37 +

38 +def pids_that_recorded(multiproc_dir: Path) -> set[str]:

39 + return {

40 + path.stem.rsplit("_", 1)[1]

41 + for path in multiproc_dir.glob("histogram_*.db")

42 + if any(

43 + metric.name == DURATION

44 + for metric in MultiProcessCollector.merge([str(path)])

45 + )

46 + }

47 +

48 +

49 +def test_weight_sync_metrics_aggregate_across_api_servers(tmp_path):

50 + args = [

51 + "--load-format",

52 + "dummy",

53 + "--enforce-eager",

54 + "--max-model-len",

55 + "512",

56 + "--gpu-memory-utilization",

57 + "0.5",

58 + "--api-server-count",

59 + "2",

60 + "--weight-transfer-config",

61 + '{"backend": "ipc"}',

62 + ]

63 + env = {

64 + "VLLM_SERVER_DEV_MODE": "1",

65 + "VLLM_ALLOW_INSECURE_SERIALIZATION": "1",

66 + "PROMETHEUS_MULTIPROC_DIR": str(tmp_path),

67 + }

68 + with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env) as server:

69 + url = server.url_root

70 + trainer = AutoModelForCausalLM.from_pretrained(

71 + MODEL_NAME, torch_dtype=torch.bfloat16

72 + ).cuda()

73 + engine = WeightTransferTrainerFactory.trainer_init(

74 + IPCTrainerInitInfo(rank=0),

75 + client=HTTPVLLMWeightSyncClient(url),

76 + source=ModuleSource(trainer),

77 + )

78 + for _ in range(SYNCS):

79 + engine.send_weights()

80 + metrics = scrape(url)

81 +

82 + calls = {"init": 1, "start": SYNCS, "update": SYNCS, "finish": SYNCS}

83 + for operation, count in calls.items():

84 + assert metrics[(f"{DURATION}_count", operation)] == count, operation

85 + assert metrics[(IN_FLIGHT, operation)] == 0, operation

86 + # The counts above are sums over processes only if both servers recorded.

87 + assert len(pids_that_recorded(tmp_path)) == 2

tests/entrypoints/serve/lora/test_serving_models.py

+13 -5

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from http import HTTPStatus

5 -from unittest.mock import AsyncMock, MagicMock

5 +from unittest.mock import MagicMock

7 import pytest

@@ -29,7 +29,9 @@ LORA_UNLOADING_SUCCESS_MESSAGE = (

29 )

32 -async def _async_serving_models_init() -> OpenAIServingModels:

32 +async def _async_serving_models_init_with_mock() -> tuple[

33 + OpenAIServingModels, MagicMock

34 +]:

35 mock_engine_client = MagicMock(spec=EngineClient)

36 # Set the max_model_len attribute to avoid missing attribute

37 mock_model_config = MagicMock(spec=ModelConfig)

@@ -45,6 +47,13 @@ async def _async_serving_models_init() -> OpenAIServingModels:

47 )

48 await serving_models.init_static_loras()

50 + # The mock is returned separately: `engine_client` is typed as the protocol,

51 + # so assertion helpers are not visible through it.

52 + return serving_models, mock_engine_client

53 +

54 +

55 +async def _async_serving_models_init() -> OpenAIServingModels:

56 + serving_models, _ = await _async_serving_models_init_with_mock()

57 return serving_models

@@ -103,7 +112,7 @@ async def test_load_lora_adapter_duplicate():

112 @pytest.mark.parametrize("load_inplace", [False, True])

113 async def test_load_lora_adapter_base_model_name(load_inplace: bool):

114 """A LoRA named after a served model would shadow it for all requests."""

106 - serving_models = await _async_serving_models_init()

115 + serving_models, mock_engine_client = await _async_serving_models_init_with_mock()

116 request = LoadLoRAAdapterRequest(

117 lora_name=MODEL_NAME, lora_path="/path/to/adapter", load_inplace=load_inplace

118 )

@@ -112,8 +121,7 @@ async def test_load_lora_adapter_base_model_name(load_inplace: bool):

121 assert response.error.type == "InvalidUserInput"

122 assert response.error.code == HTTPStatus.BAD_REQUEST

123 assert len(serving_models.lora_requests) == 0

115 - assert isinstance(serving_models.engine_client.add_lora, AsyncMock)

116 - serving_models.engine_client.add_lora.assert_not_called()

124 + mock_engine_client.add_lora.assert_not_called()

125

126

127 @pytest.mark.asyncio

tests/entrypoints/unit_tests/test_rl_weight_operation_metrics.py new

+188

@@ -0,0 +1,188 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import asyncio

4 +import os

5 +import subprocess

6 +import sys

7 +from contextlib import nullcontext

8 +from pathlib import Path

9 +from types import SimpleNamespace

10 +

11 +import pytest

12 +from fastapi import FastAPI

13 +from fastapi.testclient import TestClient

14 +from prometheus_client import CollectorRegistry

15 +

16 +from vllm.entrypoints.serve.dev.rlhf import api_router

17 +from vllm.entrypoints.serve.dev.rlhf import metrics as rlhf_metrics

18 +from vllm.entrypoints.serve.dev.rlhf.metrics import WeightOperationMetrics

19 +from vllm.entrypoints.serve.exception_handling.register import init_exception_handler

20 +from vllm.v1.metrics.prometheus import get_prometheus_registry

21 +

22 +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]

23 +

24 +DURATION = "vllm:rl_weight_update_operation_duration_seconds_count"

25 +IN_FLIGHT = "vllm:rl_weight_update_operations_in_flight"

26 +

27 +

28 +def sample(registry, name, operation):

29 + return registry.get_sample_value(name, {"operation": operation})

30 +

31 +

32 +class Request:

33 + def __init__(self, engine):

34 + self.app = SimpleNamespace(state=SimpleNamespace(engine_client=engine))

35 +

36 +

37 +class Engine:

38 + """Holds every engine call open until the test releases it."""

39 +

40 + def __init__(self):

41 + self.started = asyncio.Event()

42 + self.release = asyncio.Event()

43 +

44 + def __getattr__(self, name):

45 + async def call(*args, **kwargs):

46 + self.started.set()

47 + await self.release.wait()

48 +

49 + return call

50 +

51 +

52 +@pytest.fixture

53 +def registry(monkeypatch):

54 + registry = CollectorRegistry()

55 + monkeypatch.setattr(rlhf_metrics, "_metrics", WeightOperationMetrics(registry))

56 + return registry

57 +

58 +

59 +@pytest.mark.parametrize("fail", [False, True])

60 +def test_record_releases_in_flight_and_observes_duration(registry, fail):

61 + metrics = rlhf_metrics.weight_operation_metrics()

62 + with (

63 + pytest.raises(RuntimeError) if fail else nullcontext(),

64 + metrics.record("update"),

65 + ):

66 + assert sample(registry, IN_FLIGHT, "update") == 1

67 + if fail:

68 + raise RuntimeError

69 + assert sample(registry, IN_FLIGHT, "update") == 0

70 + assert sample(registry, DURATION, "update") == 1

71 +

72 +

73 +@pytest.mark.asyncio

74 +@pytest.mark.parametrize(

75 + "operation,route,body",

76 + [

77 + ("init", api_router.init_weight_transfer_engine, {"init_info": {}}),

78 + ("start", api_router.start_weight_update, None),

79 + ("start_draft", api_router.start_draft_weight_update, None),

80 + ("update", api_router.update_weights, {"update_info": {}}),

81 + ("finish", api_router.finish_weight_update, None),

82 + ],

83 +)

84 +async def test_route_records_its_operation(registry, operation, route, body):

85 + engine = Engine()

86 + task = asyncio.create_task(route(Request(engine), **(body or {})))

87 + await engine.started.wait()

88 + assert sample(registry, IN_FLIGHT, operation) == 1

89 + engine.release.set()

90 + await task

91 + assert sample(registry, IN_FLIGHT, operation) == 0

92 + assert sample(registry, DURATION, operation) == 1

93 +

94 +

95 +@pytest.mark.asyncio

96 +async def test_cancelled_route_releases_in_flight(registry):

97 + engine = Engine()

98 + task = asyncio.create_task(

99 + api_router.update_weights(Request(engine), update_info={})

100 + )

101 + await engine.started.wait()

102 + task.cancel()

103 + with pytest.raises(asyncio.CancelledError):

104 + await task

105 + assert sample(registry, IN_FLIGHT, "update") == 0

106 + assert sample(registry, DURATION, "update") == 1

107 +

108 +

109 +HOLD_OPERATION = """

110 +import sys

111 +from vllm.entrypoints.serve.dev.rlhf.metrics import weight_operation_metrics

112 +with weight_operation_metrics().record("update"):

113 + print("entered", flush=True)

114 + sys.stdin.readline()

115 +"""

116 +

117 +

118 +def test_in_flight_is_summed_across_api_server_processes(tmp_path, monkeypatch):

119 + """Two API-server processes each hold an operation; /metrics sums them."""

120 + monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path))

121 + servers = [

122 + subprocess.Popen(

123 + [sys.executable, "-c", HOLD_OPERATION],

124 + stdin=subprocess.PIPE,

125 + stdout=subprocess.PIPE,

126 + text=True,

127 + env=os.environ.copy(),

128 + )

129 + for _ in range(2)

130 + ]

131 + try:

132 + for server in servers:

133 + assert server.stdout is not None

134 + assert server.stdout.readline().strip() == "entered"

135 + assert sample(get_prometheus_registry(), IN_FLIGHT, "update") == 2

136 + finally:

137 + for server in servers:

138 + server.communicate("\n", timeout=60)

139 +

140 + scraped = get_prometheus_registry()

141 + assert sample(scraped, IN_FLIGHT, "update") == 0

142 + assert sample(scraped, DURATION, "update") == 2

143 +

144 +

145 +def test_rejected_request_is_not_recorded(registry):

146 + app = FastAPI()

147 + app.include_router(api_router.router)

148 + init_exception_handler(app)

149 + app.state.engine_client = Engine()

150 + app.state.args = SimpleNamespace(log_error_stack=False)

151 + response = TestClient(app).post("/update_weights", json={})

152 + assert response.status_code == 400

153 + assert sample(registry, DURATION, "update") is None

154 +

155 +

156 +def test_default_recorder_is_created_lazily_on_the_default_registry():

157 + """The default recorder registers on the process registry only when first used.

158 +

159 + Runs in a fresh interpreter: the registry is process-global, so a route test

160 + elsewhere in the session may already have registered these collectors, and

161 + "not registered yet" would then be unobservable.

162 + """

163 + probe = """

164 +from prometheus_client import REGISTRY

165 +

166 +from vllm.entrypoints.serve.dev.rlhf import metrics as rlhf_metrics

167 +

168 +NAMES = [

169 + "vllm:rl_weight_update_operation_duration_seconds",

170 + "vllm:rl_weight_update_operations_in_flight",

171 +]

172 +assert rlhf_metrics._metrics is None

173 +assert not set(NAMES) & set(REGISTRY._names_to_collectors)

174 +

175 +metrics = rlhf_metrics.weight_operation_metrics()

176 +assert REGISTRY._names_to_collectors[NAMES[0]] is metrics.duration

177 +assert REGISTRY._names_to_collectors[NAMES[1]] is metrics.in_flight

178 +assert rlhf_metrics.weight_operation_metrics() is metrics

179 +print("lazy registration ok")

180 +"""

181 + result = subprocess.run(

182 + [sys.executable, "-c", probe],

183 + cwd=Path(__file__).parents[3],

184 + capture_output=True,

185 + text=True,

186 + )

187 + assert result.returncode == 0, result.stderr

188 + assert "lazy registration ok" in result.stdout

vllm/entrypoints/serve/dev/rlhf/api_router.py

+15 -9

@@ -14,6 +14,7 @@ from vllm.distributed.weight_transfer.base import (

14 WeightTransferUpdateRequest,

15 )

16 from vllm.engine.protocol import EngineClient

17 +from vllm.entrypoints.serve.dev.rlhf.metrics import weight_operation_metrics

18 from vllm.logger import init_logger

19 from vllm.v1.engine import PauseMode

@@ -163,21 +164,24 @@ async def init_weight_transfer_engine(

164 raw_request: Request,

165 init_info: Annotated[dict[str, Any], Body(embed=True)],

166 ):

166 - await engine_client(raw_request).init_weight_transfer_engine(

167 - WeightTransferInitRequest(init_info=init_info)

168 - )

167 + with weight_operation_metrics().record("init"):

168 + await engine_client(raw_request).init_weight_transfer_engine(

169 + WeightTransferInitRequest(init_info=init_info)

170 + )

171 return JSONResponse(content={"message": "Weight transfer initialized"})

172

173

174 @router.post("/start_weight_update")

175 async def start_weight_update(raw_request: Request):

174 - await engine_client(raw_request).start_weight_update()

176 + with weight_operation_metrics().record("start"):

177 + await engine_client(raw_request).start_weight_update()

178 return JSONResponse(content={"message": "Weight update started"})

179

180

181 @router.post("/start_draft_weight_update")

182 async def start_draft_weight_update(raw_request: Request):

180 - await engine_client(raw_request).start_draft_weight_update()

183 + with weight_operation_metrics().record("start_draft"):

184 + await engine_client(raw_request).start_draft_weight_update()

185 return JSONResponse(content={"message": "Draft weight update started"})

186

187

@@ -186,9 +190,10 @@ async def update_weights(

190 raw_request: Request,

191 update_info: Annotated[WeightTransferUpdatePayload, Body(embed=True)],

192 ):

189 - await engine_client(raw_request).update_weights(

190 - request=WeightTransferUpdateRequest(update_info=update_info)

191 - )

193 + with weight_operation_metrics().record("update"):

194 + await engine_client(raw_request).update_weights(

195 + request=WeightTransferUpdateRequest(update_info=update_info)

196 + )

197 return JSONResponse(content={"message": "Weights updated"})

198

199

@@ -197,7 +202,8 @@ async def finish_weight_update(

202 raw_request: Request,

203 weight_version: Annotated[str | None, Body(embed=True)] = None,

204 ):

200 - await engine_client(raw_request).finish_weight_update(weight_version)

205 + with weight_operation_metrics().record("finish"):

206 + await engine_client(raw_request).finish_weight_update(weight_version)

207 return JSONResponse(content={"message": "Weight update finished"})

208

209

vllm/entrypoints/serve/dev/rlhf/metrics.py new

+56

@@ -0,0 +1,56 @@

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