usage_lib.py 8.85 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

yhu422's avatar
yhu422 committed
4
5
6
7
8
9
10
11
12
import datetime
import json
import logging
import os
import platform
import time
from enum import Enum
from pathlib import Path
from threading import Thread
13
from typing import Any
yhu422's avatar
yhu422 committed
14
15
16
17
18
19
20
from uuid import uuid4

import cpuinfo
import psutil
import requests
import torch

21
import vllm.envs as envs
22
from vllm.connections import global_http_connection
23
from vllm.logger import init_logger
24
from vllm.utils.platform_utils import cuda_get_device_properties
25
from vllm.utils.torch_utils import cuda_device_count_stateless
26
from vllm.version import __version__ as VLLM_VERSION
27

28
29
logger = init_logger(__name__)

30
_config_home = envs.VLLM_CONFIG_ROOT
31
32
_USAGE_STATS_JSON_PATH = os.path.join(_config_home, "usage_stats.json")
_USAGE_STATS_DO_NOT_TRACK_PATH = os.path.join(_config_home, "do_not_track")
yhu422's avatar
yhu422 committed
33
_USAGE_STATS_ENABLED = None
34
_USAGE_STATS_SERVER = envs.VLLM_USAGE_STATS_SERVER
yhu422's avatar
yhu422 committed
35

36
_GLOBAL_RUNTIME_DATA = dict[str, str | int | bool]()
37

38
39
40
41
42
43
44
45
46
_USAGE_ENV_VARS_TO_COLLECT = [
    "VLLM_USE_MODELSCOPE",
    "VLLM_ATTENTION_BACKEND",
    "VLLM_USE_FLASHINFER_SAMPLER",
    "VLLM_PP_LAYER_PARTITION",
    "VLLM_USE_TRITON_AWQ",
    "VLLM_ENABLE_V1_MULTIPROCESSING",
]

47

48
def set_runtime_usage_data(key: str, value: str | int | bool) -> None:
49
50
51
    """Set global usage data that will be sent with every usage heartbeat."""
    _GLOBAL_RUNTIME_DATA[key] = value

yhu422's avatar
yhu422 committed
52
53
54
55
56

def is_usage_stats_enabled():
    """Determine whether or not we can send usage stats to the server.
    The logic is as follows:
    - By default, it should be enabled.
57
58
    - Three environment variables can disable it:
        - VLLM_DO_NOT_TRACK=1
yhu422's avatar
yhu422 committed
59
60
61
62
63
64
65
        - DO_NOT_TRACK=1
        - VLLM_NO_USAGE_STATS=1
    - A file in the home directory can disable it if it exists:
        - $HOME/.config/vllm/do_not_track
    """
    global _USAGE_STATS_ENABLED
    if _USAGE_STATS_ENABLED is None:
66
67
        do_not_track = envs.VLLM_DO_NOT_TRACK
        no_usage_stats = envs.VLLM_NO_USAGE_STATS
yhu422's avatar
yhu422 committed
68
69
        do_not_track_file = os.path.exists(_USAGE_STATS_DO_NOT_TRACK_PATH)

70
        _USAGE_STATS_ENABLED = not (do_not_track or no_usage_stats or do_not_track_file)
yhu422's avatar
yhu422 committed
71
72
73
74
75
76
77
78
79
80
    return _USAGE_STATS_ENABLED


def _get_current_timestamp_ns() -> int:
    return int(datetime.datetime.now(datetime.timezone.utc).timestamp() * 1e9)


def _detect_cloud_provider() -> str:
    # Try detecting through vendor file
    vendor_files = [
81
82
        "/sys/class/dmi/id/product_version",
        "/sys/class/dmi/id/bios_vendor",
yhu422's avatar
yhu422 committed
83
        "/sys/class/dmi/id/product_name",
84
85
        "/sys/class/dmi/id/chassis_asset_tag",
        "/sys/class/dmi/id/sys_vendor",
yhu422's avatar
yhu422 committed
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
    ]
    # Mapping of identifiable strings to cloud providers
    cloud_identifiers = {
        "amazon": "AWS",
        "microsoft corporation": "AZURE",
        "google": "GCP",
        "oraclecloud": "OCI",
    }

    for vendor_file in vendor_files:
        path = Path(vendor_file)
        if path.is_file():
            file_content = path.read_text().lower()
            for identifier, provider in cloud_identifiers.items():
                if identifier in file_content:
                    return provider

    # Try detecting through environment variables
    env_to_cloud_provider = {
        "RUNPOD_DC_ID": "RUNPOD",
    }
    for env_var, provider in env_to_cloud_provider.items():
        if os.environ.get(env_var):
            return provider

    return "UNKNOWN"


class UsageContext(str, Enum):
    UNKNOWN_CONTEXT = "UNKNOWN_CONTEXT"
    LLM_CLASS = "LLM_CLASS"
    API_SERVER = "API_SERVER"
    OPENAI_API_SERVER = "OPENAI_API_SERVER"
119
    OPENAI_BATCH_RUNNER = "OPENAI_BATCH_RUNNER"
yhu422's avatar
yhu422 committed
120
121
122
123
124
125
126
127
128
129
130
131
132
    ENGINE_CONTEXT = "ENGINE_CONTEXT"


class UsageMessage:
    """Collect platform information and send it to the usage stats server."""

    def __init__(self) -> None:
        # NOTE: vLLM's server _only_ support flat KV pair.
        # Do not use nested fields.

        self.uuid = str(uuid4())

        # Environment Information
133
134
135
136
137
138
139
140
141
142
143
144
        self.provider: str | None = None
        self.num_cpu: int | None = None
        self.cpu_type: str | None = None
        self.cpu_family_model_stepping: str | None = None
        self.total_memory: int | None = None
        self.architecture: str | None = None
        self.platform: str | None = None
        self.cuda_runtime: str | None = None
        self.gpu_count: int | None = None
        self.gpu_type: str | None = None
        self.gpu_memory_per_device: int | None = None
        self.env_var_json: str | None = None
yhu422's avatar
yhu422 committed
145
146

        # vLLM Information
147
148
149
        self.model_architecture: str | None = None
        self.vllm_version: str | None = None
        self.context: str | None = None
yhu422's avatar
yhu422 committed
150
151

        # Metadata
152
153
        self.log_time: int | None = None
        self.source: str | None = None
yhu422's avatar
yhu422 committed
154

155
156
157
158
    def report_usage(
        self,
        model_architecture: str,
        usage_context: UsageContext,
159
        extra_kvs: dict[str, Any] | None = None,
160
161
162
163
164
165
    ) -> None:
        t = Thread(
            target=self._report_usage_worker,
            args=(model_architecture, usage_context, extra_kvs or {}),
            daemon=True,
        )
yhu422's avatar
yhu422 committed
166
167
        t.start()

168
169
170
171
172
173
    def _report_usage_worker(
        self,
        model_architecture: str,
        usage_context: UsageContext,
        extra_kvs: dict[str, Any],
    ) -> None:
yhu422's avatar
yhu422 committed
174
        self._report_usage_once(model_architecture, usage_context, extra_kvs)
Ning Xie's avatar
Ning Xie committed
175
        self._report_continuous_usage()
yhu422's avatar
yhu422 committed
176

177
178
179
180
181
182
183
184
185
186
187
188
    def _report_tpu_inference_usage(self) -> bool:
        try:
            from tpu_inference import tpu_info, utils

            self.gpu_count = tpu_info.get_num_chips()
            self.gpu_type = tpu_info.get_tpu_type()
            self.gpu_memory_per_device = utils.get_device_hbm_limit()
            self.cuda_runtime = "tpu_inference"
            return True
        except Exception:
            return False

189
190
191
192
193
194
    def _report_usage_once(
        self,
        model_architecture: str,
        usage_context: UsageContext,
        extra_kvs: dict[str, Any],
    ) -> None:
yhu422's avatar
yhu422 committed
195
        # Platform information
196
        from vllm.platforms import current_platform
197

198
        if current_platform.is_cuda_alike():
199
            self.gpu_count = cuda_device_count_stateless()
200
201
202
            self.gpu_type, self.gpu_memory_per_device = cuda_get_device_properties(
                0, ("name", "total_memory")
            )
203
204
        if current_platform.is_cuda():
            self.cuda_runtime = torch.version.cuda
205
        if current_platform.is_tpu():  # noqa: SIM102
206
            if not self._report_tpu_inference_usage():
207
                logger.exception("Failed to collect TPU information")
yhu422's avatar
yhu422 committed
208
209
210
211
212
213
214
215
        self.provider = _detect_cloud_provider()
        self.architecture = platform.machine()
        self.platform = platform.platform()
        self.total_memory = psutil.virtual_memory().total

        info = cpuinfo.get_cpu_info()
        self.num_cpu = info.get("count", None)
        self.cpu_type = info.get("brand_raw", "")
216
217
218
219
220
221
222
        self.cpu_family_model_stepping = ",".join(
            [
                str(info.get("family", "")),
                str(info.get("model", "")),
                str(info.get("stepping", "")),
            ]
        )
yhu422's avatar
yhu422 committed
223
224
225

        # vLLM information
        self.context = usage_context.value
226
        self.vllm_version = VLLM_VERSION
yhu422's avatar
yhu422 committed
227
228
        self.model_architecture = model_architecture

229
        # Environment variables
230
231
232
        self.env_var_json = json.dumps(
            {env_var: getattr(envs, env_var) for env_var in _USAGE_ENV_VARS_TO_COLLECT}
        )
233

yhu422's avatar
yhu422 committed
234
235
        # Metadata
        self.log_time = _get_current_timestamp_ns()
236
        self.source = envs.VLLM_USAGE_SOURCE
yhu422's avatar
yhu422 committed
237
238
239
240
241
242
243
244

        data = vars(self)
        if extra_kvs:
            data.update(extra_kvs)

        self._write_to_file(data)
        self._send_to_server(data)

Ning Xie's avatar
Ning Xie committed
245
    def _report_continuous_usage(self):
yhu422's avatar
yhu422 committed
246
247
248
249
250
251
252
        """Report usage every 10 minutes.

        This helps us to collect more data points for uptime of vLLM usages.
        This function can also help send over performance metrics over time.
        """
        while True:
            time.sleep(600)
253
254
255
256
257
            data = {
                "uuid": self.uuid,
                "log_time": _get_current_timestamp_ns(),
            }
            data.update(_GLOBAL_RUNTIME_DATA)
yhu422's avatar
yhu422 committed
258
259
260
261

            self._write_to_file(data)
            self._send_to_server(data)

262
    def _send_to_server(self, data: dict[str, Any]) -> None:
yhu422's avatar
yhu422 committed
263
        try:
264
265
            global_http_client = global_http_connection.get_sync_client()
            global_http_client.post(_USAGE_STATS_SERVER, json=data)
yhu422's avatar
yhu422 committed
266
267
268
269
        except requests.exceptions.RequestException:
            # silently ignore unless we are using debug log
            logging.debug("Failed to send usage data to server")

270
    def _write_to_file(self, data: dict[str, Any]) -> None:
yhu422's avatar
yhu422 committed
271
272
273
274
275
276
277
278
        os.makedirs(os.path.dirname(_USAGE_STATS_JSON_PATH), exist_ok=True)
        Path(_USAGE_STATS_JSON_PATH).touch(exist_ok=True)
        with open(_USAGE_STATS_JSON_PATH, "a") as f:
            json.dump(data, f)
            f.write("\n")


usage_message = UsageMessage()