| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import pytest |
| 4 | import regex as re |
| 5 | import requests |
| 6 | |
| 7 | from tests.utils import RemoteOpenAIServer |
| 8 | |
| 9 | # Prometheus metrics utilities for testing |
| 10 | |
| 11 | |
| 12 | def get_prometheus_metrics(server: RemoteOpenAIServer) -> dict[str, dict[str, float]]: |
| 13 | """Fetch and parse Prometheus metrics from the /metrics endpoint. |
| 14 | |
| 15 | Returns: |
| 16 | Dict mapping metric names to their values grouped by labels. |
| 17 | For example: {"vllm:request_success": { |
| 18 | "engine=0": 5.0, "engine=1": 3.0} |
| 19 | } |
| 20 | |
| 21 | """ |
| 22 | try: |
| 23 | response = requests.get(server.url_for("metrics"), timeout=10) |
| 24 | response.raise_for_status() |
| 25 | |
| 26 | metrics: dict[str, dict[str, float]] = {} |
| 27 | |
| 28 | # Regex patterns for Prometheus metrics |
| 29 | metric_with_labels = re.compile( |
| 30 | r"^([a-zA-Z_:][a-zA-Z0-9_:]*)\{([^}]*)\}\s+([\d\.\-\+e]+)$" |
| 31 | ) |
| 32 | metric_simple = re.compile(r"^([a-zA-Z_:][a-zA-Z0-9_:]*)\s+([\d\.\-\+e]+)$") |
| 33 | |
| 34 | for line in response.text.split("\n"): |
| 35 | line = line.strip() |
| 36 | # Skip comments and empty lines |
| 37 | if not line or line.startswith("#"): |
| 38 | continue |
| 39 | |
| 40 | # Try to match metric with labels first |
| 41 | match = metric_with_labels.match(line) |
| 42 | if match: |
| 43 | metric_name, labels_part, value_str = match.groups() |
| 44 | try: |
| 45 | value = float(value_str) |
| 46 | if metric_name not in metrics: |
| 47 | metrics[metric_name] = {} |
| 48 | metrics[metric_name][f"{{{labels_part}}}"] = value |
| 49 | except ValueError: |
| 50 | continue |
| 51 | else: |
| 52 | # Try simple metric without labels |
| 53 | match = metric_simple.match(line) |
| 54 | if match: |
| 55 | metric_name, value_str = match.groups() |
| 56 | try: |
| 57 | value = float(value_str) |
| 58 | if metric_name not in metrics: |
| 59 | metrics[metric_name] = {} |
| 60 | metrics[metric_name][""] = value |
| 61 | except ValueError: |
| 62 | continue |
| 63 | |
| 64 | return metrics |
| 65 | except Exception as e: |
| 66 | pytest.fail(f"Failed to fetch Prometheus metrics: {e}") |
| 67 | return {} |
| 68 | |
| 69 | |
| 70 | def get_engine_request_counts(metrics: dict[str, dict[str, float]]) -> dict[str, float]: |
| 71 | """Extract request counts per engine from Prometheus metrics. |
| 72 | |
| 73 | Returns: |
| 74 | Dict mapping engine indices to request counts. |
| 75 | For example: {"0": 15.0, "1": 12.0} |
| 76 | |
| 77 | """ |
| 78 | engine_counts = {} |
| 79 | |
| 80 | # Look for request success metrics with engine labels |
| 81 | success_metrics = metrics.get("vllm:request_success_total", {}) |
| 82 | engine_pattern = re.compile(r'engine="([^"]*)"') |
| 83 | |
| 84 | for labels, count in success_metrics.items(): |
| 85 | # Extract engine ID from labels using regex |
| 86 | match = engine_pattern.search(labels) |
| 87 | if match: |
| 88 | engine_id = match.group(1) |
| 89 | if engine_id not in engine_counts: |
| 90 | engine_counts[engine_id] = 0.0 |
| 91 | engine_counts[engine_id] += count |
| 92 | |
| 93 | return engine_counts |
| 94 | |
| 95 | |
| 96 | def check_request_balancing(server: RemoteOpenAIServer, dp_size: int): |
| 97 | """Check request balancing via Prometheus metrics if dp_size > 1. |
| 98 | |
| 99 | Args: |
| 100 | server: The RemoteOpenAIServer instance |
| 101 | dp_size: Number of data parallel ranks |
| 102 | |
| 103 | """ |
| 104 | if dp_size <= 1: |
| 105 | return |
| 106 | |
| 107 | # Get metrics after all requests are completed |
| 108 | metrics = get_prometheus_metrics(server) |
| 109 | engine_counts = get_engine_request_counts(metrics) |
| 110 | |
| 111 | # Check that multiple engines received requests |
| 112 | engines_with_requests = [ |
| 113 | engine for engine, count in engine_counts.items() if count > 0 |
| 114 | ] |
| 115 | assert len(engines_with_requests) == dp_size, ( |
| 116 | f"Expected requests to be distributed across multiple engines," |
| 117 | f" but only engine(s) {engines_with_requests} received " |
| 118 | f"requests. Engine counts: {engine_counts}" |
| 119 | ) |
| 120 | |
| 121 | # Verify that the load is reasonably balanced |
| 122 | # (no engine should handle all requests) |
| 123 | total_requests = sum(engine_counts.values()) |
| 124 | |
| 125 | for count in engine_counts.values(): |
| 126 | assert count > total_requests // (dp_size + 1), ( |
| 127 | f"requests are imbalanced: {engine_counts}" |
| 128 | ) |