Usage Stats Collection (#2852)

vllm/usage/usage_lib.py

@@ -0,0 +1,207 @@
+import datetime
+import json
+import logging
+import os
+import platform
+import time
+from enum import Enum
+from pathlib import Path
+from threading import Thread
+from typing import Dict, Optional
+from uuid import uuid4
+
+import cpuinfo
+import pkg_resources
+import psutil
+import requests
+import torch
+
+_config_home = os.getenv("XDG_CONFIG_HOME", os.path.expanduser("~/.config"))
+_USAGE_STATS_JSON_PATH = os.path.join(_config_home, "vllm/usage_stats.json")
+_USAGE_STATS_DO_NOT_TRACK_PATH = os.path.join(_config_home,
+                                              "vllm/do_not_track")
+_USAGE_STATS_ENABLED = None
+_USAGE_STATS_SERVER = os.environ.get("VLLM_USAGE_STATS_SERVER",
+                                     "https://stats.vllm.ai")
+
+
+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.
+    - Two environment variables can disable it:
+        - 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:
+        do_not_track = os.environ.get("DO_NOT_TRACK", "0") == "1"
+        no_usage_stats = os.environ.get("VLLM_NO_USAGE_STATS", "0") == "1"
+        do_not_track_file = os.path.exists(_USAGE_STATS_DO_NOT_TRACK_PATH)
+
+        _USAGE_STATS_ENABLED = not (do_not_track or no_usage_stats
+                                    or do_not_track_file)
+    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 = [
+        "/sys/class/dmi/id/product_version", "/sys/class/dmi/id/bios_vendor",
+        "/sys/class/dmi/id/product_name",
+        "/sys/class/dmi/id/chassis_asset_tag", "/sys/class/dmi/id/sys_vendor"
+    ]
+    # 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"
+    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
+        self.provider: Optional[str] = None
+        self.num_cpu: Optional[int] = None
+        self.cpu_type: Optional[str] = None
+        self.cpu_family_model_stepping: Optional[str] = None
+        self.total_memory: Optional[int] = None
+        self.architecture: Optional[str] = None
+        self.platform: Optional[str] = None
+        self.gpu_count: Optional[int] = None
+        self.gpu_type: Optional[str] = None
+        self.gpu_memory_per_device: Optional[int] = None
+
+        # vLLM Information
+        self.model_architecture: Optional[str] = None
+        self.vllm_version: Optional[str] = None
+        self.context: Optional[str] = None
+
+        # Metadata
+        self.log_time: Optional[int] = None
+        self.source: Optional[str] = None
+
+    def report_usage(self,
+                     model_architecture: str,
+                     usage_context: UsageContext,
+                     extra_kvs: Dict[str, any] = None) -> None:
+        t = Thread(target=self._report_usage_worker,
+                   args=(model_architecture, usage_context, extra_kvs or {}),
+                   daemon=True)
+        t.start()
+
+    def _report_usage_worker(self, model_architecture: str,
+                             usage_context: UsageContext,
+                             extra_kvs: Dict[str, any]) -> None:
+        self._report_usage_once(model_architecture, usage_context, extra_kvs)
+        self._report_continous_usage()
+
+    def _report_usage_once(self, model_architecture: str,
+                           usage_context: UsageContext,
+                           extra_kvs: Dict[str, any]) -> None:
+        # Platform information
+        if torch.cuda.is_available():
+            device_property = torch.cuda.get_device_properties(0)
+            self.gpu_count = torch.cuda.device_count()
+            self.gpu_type = device_property.name
+            self.gpu_memory_per_device = device_property.total_memory
+        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", "")
+        self.cpu_family_model_stepping = ",".join([
+            str(info.get("family", "")),
+            str(info.get("model", "")),
+            str(info.get("stepping", ""))
+        ])
+
+        # vLLM information
+        self.context = usage_context.value
+        self.vllm_version = pkg_resources.get_distribution("vllm").version
+        self.model_architecture = model_architecture
+
+        # Metadata
+        self.log_time = _get_current_timestamp_ns()
+        self.source = os.environ.get("VLLM_USAGE_SOURCE", "production")
+
+        data = vars(self)
+        if extra_kvs:
+            data.update(extra_kvs)
+
+        self._write_to_file(data)
+        self._send_to_server(data)
+
+    def _report_continous_usage(self):
+        """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)
+            data = {"uuid": self.uuid, "log_time": _get_current_timestamp_ns()}
+
+            self._write_to_file(data)
+            self._send_to_server(data)
+
+    def _send_to_server(self, data):
+        try:
+            requests.post(_USAGE_STATS_SERVER, json=data)
+        except requests.exceptions.RequestException:
+            # silently ignore unless we are using debug log
+            logging.debug("Failed to send usage data to server")
+
+    def _write_to_file(self, data):
+        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()