main
py 281 lines 9 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import datetime
5 import json
6 import logging
7 import os
8 import platform
9 import time
10 from enum import Enum
11 from pathlib import Path
12 from threading import Thread
13 from typing import Any
14 from uuid import uuid4
15
16 import cpuinfo
17 import psutil
18 import requests
19 import torch
20
21 import vllm.envs as envs
22 from vllm.logger import init_logger
23 from vllm.utils.platform_utils import cuda_get_device_properties
24 from vllm.version import __version__ as VLLM_VERSION
25
26 logger = init_logger(__name__)
27
28 _config_home = envs.VLLM_CONFIG_ROOT
29 _USAGE_STATS_JSON_PATH = os.path.join(_config_home, "usage_stats.json")
30 _USAGE_STATS_DO_NOT_TRACK_PATH = os.path.join(_config_home, "do_not_track")
31 _USAGE_STATS_ENABLED = None
32 _USAGE_STATS_SERVER = envs.VLLM_USAGE_STATS_SERVER
33
34 _GLOBAL_RUNTIME_DATA = dict[str, str | int | bool]()
35
36 _USAGE_ENV_VARS_TO_COLLECT = [
37 "VLLM_USE_MODELSCOPE",
38 "VLLM_USE_FLASHINFER_SAMPLER",
39 "VLLM_PP_LAYER_PARTITION",
40 "VLLM_USE_TRITON_AWQ",
41 "VLLM_ENABLE_V1_MULTIPROCESSING",
42 ]
43
44
45 def set_runtime_usage_data(key: str, value: str | int | bool) -> None:
46 """Set global usage data that will be sent with every usage heartbeat."""
47 _GLOBAL_RUNTIME_DATA[key] = value
48
49
50 def is_usage_stats_enabled():
51 """Determine whether or not we can send usage stats to the server.
52 The logic is as follows:
53 - By default, it should be enabled.
54 - Three environment variables can disable it:
55 - VLLM_DO_NOT_TRACK=1
56 - DO_NOT_TRACK=1
57 - VLLM_NO_USAGE_STATS=1
58 - A file in the home directory can disable it if it exists:
59 - $HOME/.config/vllm/do_not_track
60 """
61 global _USAGE_STATS_ENABLED
62 if _USAGE_STATS_ENABLED is None:
63 do_not_track = envs.VLLM_DO_NOT_TRACK
64 no_usage_stats = envs.VLLM_NO_USAGE_STATS
65 do_not_track_file = os.path.exists(_USAGE_STATS_DO_NOT_TRACK_PATH)
66
67 _USAGE_STATS_ENABLED = not (do_not_track or no_usage_stats or do_not_track_file)
68 return _USAGE_STATS_ENABLED
69
70
71 def _get_current_timestamp_ns() -> int:
72 return int(datetime.datetime.now(datetime.timezone.utc).timestamp() * 1e9)
73
74
75 def _detect_cloud_provider() -> str:
76 # Try detecting through vendor file
77 vendor_files = [
78 "/sys/class/dmi/id/product_version",
79 "/sys/class/dmi/id/bios_vendor",
80 "/sys/class/dmi/id/product_name",
81 "/sys/class/dmi/id/chassis_asset_tag",
82 "/sys/class/dmi/id/sys_vendor",
83 ]
84 # Mapping of identifiable strings to cloud providers
85 cloud_identifiers = {
86 "amazon": "AWS",
87 "microsoft corporation": "AZURE",
88 "google": "GCP",
89 "oraclecloud": "OCI",
90 }
91
92 for vendor_file in vendor_files:
93 path = Path(vendor_file)
94 if path.is_file():
95 file_content = path.read_text().lower()
96 for identifier, provider in cloud_identifiers.items():
97 if identifier in file_content:
98 return provider
99
100 # Try detecting through environment variables
101 env_to_cloud_provider = {
102 "RUNPOD_DC_ID": "RUNPOD",
103 }
104 for env_var, provider in env_to_cloud_provider.items():
105 if os.environ.get(env_var):
106 return provider
107
108 return "UNKNOWN"
109
110
111 class UsageContext(str, Enum):
112 UNKNOWN_CONTEXT = "UNKNOWN_CONTEXT"
113 LLM_CLASS = "LLM_CLASS"
114 API_SERVER = "API_SERVER"
115 OPENAI_API_SERVER = "OPENAI_API_SERVER"
116 OPENAI_BATCH_RUNNER = "OPENAI_BATCH_RUNNER"
117 ENGINE_CONTEXT = "ENGINE_CONTEXT"
118
119
120 class UsageMessage:
121 """Collect platform information and send it to the usage stats server."""
122
123 def __init__(self) -> None:
124 # NOTE: vLLM's server _only_ support flat KV pair.
125 # Do not use nested fields.
126
127 self.uuid = str(uuid4())
128
129 # Environment Information
130 self.provider: str | None = None
131 self.num_cpu: int | None = None
132 self.cpu_type: str | None = None
133 self.cpu_family_model_stepping: str | None = None
134 self.total_memory: int | None = None
135 self.architecture: str | None = None
136 self.platform: str | None = None
137 self.xpu_runtime: str | None = None
138 self.cuda_runtime: str | None = None
139 self.gpu_count: int | None = None
140 self.gpu_type: str | None = None
141 self.gpu_memory_per_device: int | None = None
142 self.env_var_json: str | None = None
143
144 # vLLM Information
145 self.model_architecture: str | None = None
146 self.vllm_version: str | None = None
147 self.context: str | None = None
148
149 # Metadata
150 self.log_time: int | None = None
151 self.source: str | None = None
152
153 def report_usage(
154 self,
155 model_architecture: str,
156 usage_context: UsageContext,
157 extra_kvs: dict[str, Any] | None = None,
158 ) -> None:
159 t = Thread(
160 target=self._report_usage_worker,
161 args=(model_architecture, usage_context, extra_kvs or {}),
162 daemon=True,
163 )
164 t.start()
165
166 def _report_usage_worker(
167 self,
168 model_architecture: str,
169 usage_context: UsageContext,
170 extra_kvs: dict[str, Any],
171 ) -> None:
172 self._report_usage_once(model_architecture, usage_context, extra_kvs)
173 self._report_continuous_usage()
174
175 def _report_tpu_inference_usage(self) -> bool:
176 try:
177 from tpu_inference import tpu_info, utils
178
179 self.gpu_count = tpu_info.get_num_chips()
180 self.gpu_type = tpu_info.get_tpu_type()
181 self.gpu_memory_per_device = utils.get_device_hbm_limit()
182 self.cuda_runtime = "tpu_inference"
183 return True
184 except Exception:
185 return False
186
187 def _report_usage_once(
188 self,
189 model_architecture: str,
190 usage_context: UsageContext,
191 extra_kvs: dict[str, Any],
192 ) -> None:
193 # Platform information
194 from vllm.platforms import current_platform
195
196 if current_platform.is_cuda_alike():
197 self.gpu_count = current_platform.device_count()
198 self.gpu_type, self.gpu_memory_per_device = cuda_get_device_properties(
199 0, ("name", "total_memory")
200 )
201 if current_platform.is_cuda():
202 self.cuda_runtime = torch.version.cuda
203 if current_platform.is_xpu():
204 self.xpu_runtime = torch.version.xpu
205 self.gpu_count = torch.xpu.device_count()
206 self.gpu_type = torch.xpu.get_device_name(0)
207 self.gpu_memory_per_device = torch.xpu.get_device_properties(0).total_memory
208 if current_platform.is_tpu(): # noqa: SIM102
209 if not self._report_tpu_inference_usage():
210 logger.exception("Failed to collect TPU information")
211 self.provider = _detect_cloud_provider()
212 self.architecture = platform.machine()
213 self.platform = platform.platform()
214 self.total_memory = psutil.virtual_memory().total
215
216 info = cpuinfo.get_cpu_info()
217 self.num_cpu = info.get("count", None)
218 self.cpu_type = info.get("brand_raw", "")
219 self.cpu_family_model_stepping = ",".join(
220 [
221 str(info.get("family", "")),
222 str(info.get("model", "")),
223 str(info.get("stepping", "")),
224 ]
225 )
226
227 # vLLM information
228 self.context = usage_context.value
229 self.vllm_version = VLLM_VERSION
230 self.model_architecture = model_architecture
231
232 # Environment variables
233 self.env_var_json = json.dumps(
234 {env_var: getattr(envs, env_var) for env_var in _USAGE_ENV_VARS_TO_COLLECT}
235 )
236
237 # Metadata
238 self.log_time = _get_current_timestamp_ns()
239 self.source = envs.VLLM_USAGE_SOURCE
240
241 data = vars(self)
242 if extra_kvs:
243 data.update(extra_kvs)
244
245 self._write_to_file(data)
246 self._send_to_server(data)
247
248 def _report_continuous_usage(self):
249 """Report usage every 10 minutes.
250
251 This helps us to collect more data points for uptime of vLLM usages.
252 This function can also help send over performance metrics over time.
253 """
254 while True:
255 time.sleep(600)
256 data = {
257 "uuid": self.uuid,
258 "log_time": _get_current_timestamp_ns(),
259 }
260 data.update(_GLOBAL_RUNTIME_DATA)
261
262 self._write_to_file(data)
263 self._send_to_server(data)
264
265 def _send_to_server(self, data: dict[str, Any]) -> None:
266 try:
267 with requests.Session() as client:
268 client.post(_USAGE_STATS_SERVER, json=data)
269 except requests.exceptions.RequestException:
270 # silently ignore unless we are using debug log
271 logging.debug("Failed to send usage data to server")
272
273 def _write_to_file(self, data: dict[str, Any]) -> None:
274 os.makedirs(os.path.dirname(_USAGE_STATS_JSON_PATH), exist_ok=True)
275 Path(_USAGE_STATS_JSON_PATH).touch(exist_ok=True)
276 with open(_USAGE_STATS_JSON_PATH, "a") as f:
277 json.dump(data, f)
278 f.write("\n")
279
280
281 usage_message = UsageMessage()