@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 119 lines 4.42 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Utility functions that create ORCA endpoint load report response headers."""
4
5 import json
6 from collections.abc import Mapping
7
8 from vllm.logger import init_logger
9 from vllm.v1.metrics.reader import Gauge, get_metrics_snapshot
10
11 logger = init_logger(__name__)
12
13
14 def create_orca_header(
15 metrics_format: str, named_metrics: list[tuple[str, float]]
16 ) -> Mapping[str, str] | None:
17 """Creates ORCA headers named 'endpoint-load-metrics' in the specified format
18 and adds custom metrics to named_metrics.
19 ORCA headers format description: https://docs.google.com/document/d/1C1ybMmDKJIVlrbOLbywhu9iRYo4rilR-cT50OTtOFTs/edit?tab=t.0
20 ORCA proto https://github.com/cncf/xds/blob/main/xds/data/orca/v3/orca_load_report.proto
21
22 Args:
23 metrics_format (str): The format of the header ('TEXT', 'JSON').
24 named_metrics (List[Tuple[str, float]]): List of tuples with metric names
25 and their corresponding double values.
26
27 Returns:
28 - Optional[Mapping[str,str]]: A dictionary with header key as
29 'endpoint-load-metrics' and values as the ORCA header strings with
30 format prefix and data in with named_metrics in.
31
32 """
33 if metrics_format.lower() not in ["text", "json"]:
34 logger.warning(
35 "Warning: `%s` format is not supported in the ORCA response header",
36 format,
37 )
38 return None
39
40 header = {}
41 orca_report = {
42 "named_metrics": {
43 metric_name: value
44 for metric_name, value in named_metrics
45 if isinstance(metric_name, str) and isinstance(value, float)
46 }
47 }
48 # output example:
49 # endpoint-load-metrics: TEXT named_metrics.kv_cache_utilization=0.4
50 if metrics_format.lower() == "text":
51 native_http_header = ", ".join(
52 [
53 f"named_metrics.{metric_name}={value}"
54 for metric_name, value in named_metrics
55 if isinstance(metric_name, str) and isinstance(value, float)
56 ]
57 )
58 header["endpoint-load-metrics"] = f"TEXT {native_http_header}"
59
60 # output example:
61 # endpoint-load-metrics: JSON “named_metrics”: {“custom-metric-util”: 0.4}
62 elif metrics_format.lower() == "json":
63 header["endpoint-load-metrics"] = f"JSON {json.dumps(orca_report)}"
64
65 logger.info("Created ORCA header %s", header)
66
67 return header
68
69
70 def get_named_metrics_from_prometheus() -> list[tuple[str, float]]:
71 """Collects current metrics from Prometheus and returns some of them
72 in the form of the `named_metrics` list for `create_orca_header()`.
73
74 Parameters
75 ----------
76 - None
77
78 Returns
79 -------
80 - list[tuple[str, float]]: List of tuples of metric names and their values.
81
82 """
83 named_metrics: list[tuple[str, float]] = []
84 # Map from prometheus metric names to ORCA named metrics.
85 prometheus_to_orca_metrics = {
86 "vllm:kv_cache_usage_perc": "kv_cache_usage_perc",
87 "vllm:num_requests_waiting": "num_requests_waiting",
88 }
89 metrics = get_metrics_snapshot()
90 for metric in metrics:
91 orca_name = prometheus_to_orca_metrics.get(metric.name)
92 # If this metric is mapped into ORCA, then add it to the report.
93 # Note: Only Gauge metrics are currently supported.
94 if orca_name is not None and isinstance(metric, Gauge):
95 named_metrics.append((str(orca_name), float(metric.value)))
96 return named_metrics
97
98
99 def metrics_header(metrics_format: str) -> Mapping[str, str] | None:
100 """Creates ORCA headers named 'endpoint-load-metrics' in the specified format.
101 Metrics are collected from Prometheus using `get_named_metrics_from_prometheus()`.
102
103 ORCA headers format description: https://docs.google.com/document/d/1C1ybMmDKJIVlrbOLbywhu9iRYo4rilR-cT50OTtOFTs/edit?tab=t.0
104 ORCA proto https://github.com/cncf/xds/blob/main/xds/data/orca/v3/orca_load_report.proto
105
106 Args:
107 metrics_format (str): The format of the header ('TEXT', 'JSON').
108
109 Returns:
110 - Optional[Mapping[str,str]]: A dictionary with header key as
111 'endpoint-load-metrics' and values as the ORCA header strings with
112 format prefix and data in with named_metrics in.
113
114 """
115 if not metrics_format:
116 return None
117 # Get named metrics from prometheus.
118 named_metrics = get_named_metrics_from_prometheus()
119 return create_orca_header(metrics_format, named_metrics)