| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from unittest.mock import MagicMock |
| 5 | |
| 6 | import pytest |
| 7 | import ray |
| 8 | |
| 9 | from tests.utils import wait_for_memory_to_settle |
| 10 | from vllm.config.model import ModelDType |
| 11 | from vllm.platforms import current_platform |
| 12 | from vllm.sampling_params import SamplingParams |
| 13 | from vllm.v1.engine.async_llm import AsyncEngineArgs, AsyncLLM |
| 14 | from vllm.v1.metrics.ray_wrappers import ( |
| 15 | RayCounterWrapper, |
| 16 | RayGaugeWrapper, |
| 17 | RayHistogramWrapper, |
| 18 | RayPrometheusMetric, |
| 19 | RayPrometheusStatLogger, |
| 20 | ) |
| 21 | |
| 22 | MODELS = [ |
| 23 | "distilbert/distilgpt2", |
| 24 | ] |
| 25 | |
| 26 | |
| 27 | # The first .remote() call starts a local Ray cluster via ray.init(), whose |
| 28 | # GCS server occasionally fails to start within Ray's fixed 30s bootstrap |
| 29 | # window on ROCm CI (RuntimeError: "Timed out waiting for file |
| 30 | # .../gcs_server_port_..."). That timeout is not configurable, so retry the |
| 31 | # whole test: a fresh ray.init() gets a new GCS process. Scoped to that error |
| 32 | # so real failures still fail immediately. |
| 33 | @pytest.mark.flaky( |
| 34 | reruns=2, |
| 35 | reruns_delay=5, |
| 36 | only_rerun="Timed out waiting for file", |
| 37 | condition=current_platform.is_rocm(), |
| 38 | ) |
| 39 | @pytest.mark.parametrize("model", MODELS) |
| 40 | @pytest.mark.parametrize("dtype", ["half"]) |
| 41 | @pytest.mark.parametrize("max_tokens", [16]) |
| 42 | def test_engine_log_metrics_ray( |
| 43 | example_prompts, |
| 44 | model: str, |
| 45 | dtype: ModelDType, |
| 46 | max_tokens: int, |
| 47 | ) -> None: |
| 48 | """Simple smoke test, verifying this can be used without exceptions. |
| 49 | Need to start a Ray cluster in order to verify outputs.""" |
| 50 | engine_args = AsyncEngineArgs( |
| 51 | model=model, dtype=dtype, disable_log_stats=False, enforce_eager=True |
| 52 | ) |
| 53 | |
| 54 | @ray.remote(num_gpus=1) |
| 55 | class EngineTestActor: |
| 56 | async def run(self): |
| 57 | engine = AsyncLLM.from_engine_args( |
| 58 | engine_args, stat_loggers=[RayPrometheusStatLogger] |
| 59 | ) |
| 60 | |
| 61 | try: |
| 62 | for i, prompt in enumerate(example_prompts): |
| 63 | results = engine.generate( |
| 64 | request_id=f"request-id-{i}", |
| 65 | prompt=prompt, |
| 66 | sampling_params=SamplingParams(max_tokens=max_tokens), |
| 67 | ) |
| 68 | |
| 69 | async for _ in results: |
| 70 | pass |
| 71 | finally: |
| 72 | engine.shutdown() |
| 73 | |
| 74 | # Create the actor and call the async method |
| 75 | try: |
| 76 | actor = EngineTestActor.remote() # type: ignore[attr-defined] |
| 77 | ray.get(actor.run.remote()) |
| 78 | finally: |
| 79 | ray.shutdown() |
| 80 | wait_for_memory_to_settle( |
| 81 | threshold_ratio=1.0 - engine_args.gpu_memory_utilization |
| 82 | ) |
| 83 | |
| 84 | |
| 85 | def test_sanitized_opentelemetry_name(): |
| 86 | """Test the metric name sanitization logic for Ray.""" |
| 87 | # Only a-z, A-Z, 0-9, _, test valid characters are preserved |
| 88 | valid_name = "valid_metric_123_abcDEF" |
| 89 | assert ( |
| 90 | RayPrometheusMetric._get_sanitized_opentelemetry_name(valid_name) == valid_name |
| 91 | ) |
| 92 | |
| 93 | # Test dash, dot, are replaced |
| 94 | name_with_dash_dot = "metric-name.test" |
| 95 | expected = "metric_name_test" |
| 96 | assert ( |
| 97 | RayPrometheusMetric._get_sanitized_opentelemetry_name(name_with_dash_dot) |
| 98 | == expected |
| 99 | ) |
| 100 | |
| 101 | # Test colon is replaced with underscore |
| 102 | name_with_colon = "metric:name" |
| 103 | expected = "metric_name" |
| 104 | assert ( |
| 105 | RayPrometheusMetric._get_sanitized_opentelemetry_name(name_with_colon) |
| 106 | == expected |
| 107 | ) |
| 108 | |
| 109 | # Test multiple invalid characters are replaced |
| 110 | name_with_invalid = "metric:name@with#special%chars" |
| 111 | expected = "metric_name_with_special_chars" |
| 112 | assert ( |
| 113 | RayPrometheusMetric._get_sanitized_opentelemetry_name(name_with_invalid) |
| 114 | == expected |
| 115 | ) |
| 116 | |
| 117 | # Test mixed valid and invalid characters |
| 118 | complex_name = "vllm:engine_stats/time.latency_ms-99p" |
| 119 | expected = "vllm_engine_stats_time_latency_ms_99p" |
| 120 | assert ( |
| 121 | RayPrometheusMetric._get_sanitized_opentelemetry_name(complex_name) == expected |
| 122 | ) |
| 123 | |
| 124 | # Test empty string |
| 125 | assert RayPrometheusMetric._get_sanitized_opentelemetry_name("") == "" |
| 126 | |
| 127 | |
| 128 | def _install_mock_metric(wrapper: RayPrometheusMetric) -> MagicMock: |
| 129 | """Swap the wrapper's underlying Ray metric for a MagicMock while |
| 130 | preserving the real metric's ``_tag_keys`` (labels() reads them to |
| 131 | validate arity).""" |
| 132 | real_metric = wrapper.metric |
| 133 | mock = MagicMock() |
| 134 | mock._tag_keys = real_metric._tag_keys |
| 135 | wrapper.metric = mock |
| 136 | return mock |
| 137 | |
| 138 | |
| 139 | def test_ray_counter_labels_returns_independent_children(): |
| 140 | """RayCounterWrapper.labels() must return distinct labeled children that |
| 141 | each carry their own tag set.""" |
| 142 | base = RayCounterWrapper( |
| 143 | name="vllm_test_finish_reason", |
| 144 | documentation="", |
| 145 | labelnames=["reason"], |
| 146 | ) |
| 147 | |
| 148 | stop_child = base.labels("stop") |
| 149 | rep_child = base.labels("repetition") |
| 150 | |
| 151 | assert stop_child is not rep_child |
| 152 | assert stop_child._tags["reason"] == "stop" |
| 153 | assert rep_child._tags["reason"] == "repetition" |
| 154 | # Mutating one child's tags must not leak into another. |
| 155 | stop_child._tags["reason"] = "mutated" |
| 156 | assert rep_child._tags["reason"] == "repetition" |
| 157 | |
| 158 | |
| 159 | def test_ray_counter_inc_forwards_per_child_tags(): |
| 160 | """.inc() on a labeled counter must forward that child's tags to the |
| 161 | underlying Ray metric (not rely on a shared set_default_tags).""" |
| 162 | wrapper = RayCounterWrapper( |
| 163 | name="vllm_test_counter_tag_forward", |
| 164 | documentation="", |
| 165 | labelnames=["reason"], |
| 166 | ) |
| 167 | mock = _install_mock_metric(wrapper) |
| 168 | |
| 169 | wrapper.labels("stop").inc() |
| 170 | wrapper.labels("repetition").inc(3) |
| 171 | wrapper.labels("stop").inc(0) # zero increment must be a no-op. |
| 172 | |
| 173 | # The zero-increment call should not reach the underlying metric. |
| 174 | assert mock.inc.call_count == 2 |
| 175 | first, second = mock.inc.call_args_list |
| 176 | assert first.args == (1.0,) |
| 177 | assert first.kwargs["tags"]["reason"] == "stop" |
| 178 | assert second.args == (3,) |
| 179 | assert second.kwargs["tags"]["reason"] == "repetition" |
| 180 | |
| 181 | |
| 182 | def test_ray_gauge_labels_returns_independent_children_and_forwards_tags(): |
| 183 | wrapper = RayGaugeWrapper( |
| 184 | name="vllm_test_gauge_tag_forward", |
| 185 | documentation="", |
| 186 | labelnames=["kind"], |
| 187 | ) |
| 188 | mock = _install_mock_metric(wrapper) |
| 189 | |
| 190 | a = wrapper.labels("a") |
| 191 | b = wrapper.labels("b") |
| 192 | assert a is not b |
| 193 | |
| 194 | a.set(1) |
| 195 | b.set(2) |
| 196 | assert mock.set.call_args_list[0].args == (1,) |
| 197 | assert mock.set.call_args_list[0].kwargs["tags"]["kind"] == "a" |
| 198 | assert mock.set.call_args_list[1].args == (2,) |
| 199 | assert mock.set.call_args_list[1].kwargs["tags"]["kind"] == "b" |
| 200 | |
| 201 | |
| 202 | def test_ray_histogram_labels_returns_independent_children_and_forwards_tags(): |
| 203 | wrapper = RayHistogramWrapper( |
| 204 | name="vllm_test_histogram_tag_forward", |
| 205 | documentation="", |
| 206 | labelnames=["bucket"], |
| 207 | buckets=[1.0, 2.0, 5.0], |
| 208 | ) |
| 209 | mock = _install_mock_metric(wrapper) |
| 210 | |
| 211 | x = wrapper.labels("x") |
| 212 | y = wrapper.labels("y") |
| 213 | assert x is not y |
| 214 | |
| 215 | x.observe(0.5) |
| 216 | y.observe(4.0) |
| 217 | assert mock.observe.call_args_list[0].args == (0.5,) |
| 218 | assert mock.observe.call_args_list[0].kwargs["tags"]["bucket"] == "x" |
| 219 | assert mock.observe.call_args_list[1].args == (4.0,) |
| 220 | assert mock.observe.call_args_list[1].kwargs["tags"]["bucket"] == "y" |
| 221 | |
| 222 | |
| 223 | def test_ray_counter_labels_accepts_non_string_label_values(): |
| 224 | """RayPrometheusStatLogger passes ``str(idx)`` for engine indexes; this |
| 225 | covers the coercion path for any caller that passes a non-string label |
| 226 | value positionally.""" |
| 227 | wrapper = RayCounterWrapper( |
| 228 | name="vllm_test_nonstr_label", |
| 229 | documentation="", |
| 230 | labelnames=["engine", "reason"], |
| 231 | ) |
| 232 | child = wrapper.labels(0, "stop") |
| 233 | assert child._tags["engine"] == "0" |
| 234 | assert child._tags["reason"] == "stop" |
| 235 | |
| 236 | |
| 237 | def test_ray_counter_labels_arity_validation(): |
| 238 | wrapper = RayCounterWrapper( |
| 239 | name="vllm_test_arity", |
| 240 | documentation="", |
| 241 | labelnames=["a", "b"], |
| 242 | ) |
| 243 | with pytest.raises(ValueError, match="Number of labels must match"): |
| 244 | wrapper.labels("only-one") |
| 245 | |
| 246 | |
| 247 | def test_unlabeled_inc_carries_replica_id(): |
| 248 | """Recording on an unlabeled metric must still pass ReplicaId — it's a |
| 249 | declared tag_key and Ray rejects updates that omit any declared key.""" |
| 250 | wrapper = RayCounterWrapper( |
| 251 | name="vllm_test_unlabeled_replica_id", |
| 252 | documentation="", |
| 253 | labelnames=None, |
| 254 | ) |
| 255 | mock = _install_mock_metric(wrapper) |
| 256 | wrapper.inc() |
| 257 | assert mock.inc.call_args.kwargs["tags"] == {"ReplicaId": ""} |
| 258 | |
| 259 | |
| 260 | def test_double_labels_raises(): |
| 261 | """labels() on an already-labeled child should raise, mirroring the |
| 262 | prometheus_client contract.""" |
| 263 | wrapper = RayCounterWrapper( |
| 264 | name="vllm_test_double_labels", |
| 265 | documentation="", |
| 266 | labelnames=["reason"], |
| 267 | ) |
| 268 | child = wrapper.labels("stop") |
| 269 | with pytest.raises(ValueError, match="already-labeled"): |
| 270 | child.labels("repetition") |