Usage Stats Collection

.buildkite/test-template.j2

@@ -45,6 +45,8 @@ steps:
                     nvidia.com/gpu: "{{ step.num_gpus or default_num_gpu }}"
                 {% endif %}
                 env:
+                  - name: VLLM_USAGE_SOURCE
+                    value: ci-test
                   - name: HF_TOKEN
                     valueFrom:
                       secretKeyRef:

requirements.txt

@@ -4,6 +4,8 @@ ray >= 2.9
 sentencepiece  # Required for LLaMA tokenizer.
 numpy
 torch == 2.1.2
+requests
+psutil
 transformers >= 4.38.0  # Required for Gemma.
 xformers == 0.0.23.post1  # Required for CUDA 12.1.
 fastapi

vllm/engine/async_llm_engine.py

@@ -13,6 +13,7 @@
 from vllm.logger import init_logger
 from vllm.outputs import RequestOutput
 from vllm.sampling_params import SamplingParams
+from vllm.usage.usage_lib import UsageContext
 
 logger = init_logger(__name__)
 ENGINE_ITERATION_TIMEOUT_S = int(
@@ -666,9 +667,12 @@ async def get_model_config(self) -> ModelConfig:
             return self.engine.get_model_config()
 
     @classmethod
-    def from_engine_args(cls,
-                         engine_args: AsyncEngineArgs,
-                         start_engine_loop: bool = True) -> "AsyncLLMEngine":
+    def from_engine_args(
+        cls,
+        engine_args: AsyncEngineArgs,
+        start_engine_loop: bool = True,
+        usage_context: UsageContext = UsageContext.ENGINE_CONTEXT
+    ) -> "AsyncLLMEngine":
         """Creates an async LLM engine from the engine arguments."""
         # Create the engine configs.
         engine_configs = engine_args.create_engine_configs()
@@ -684,7 +688,8 @@ def from_engine_args(cls,
                      log_requests=not engine_args.disable_log_requests,
                      log_stats=not engine_args.disable_log_stats,
                      max_log_len=engine_args.max_log_len,
-                     start_engine_loop=start_engine_loop)
+                     start_engine_loop=start_engine_loop,
+                     usage_context=usage_context)
         return engine
 
     async def do_log_stats(self) -> None:

vllm/engine/llm_engine.py

@@ -24,7 +24,7 @@
                                                TokenizerGroup)
 from vllm.utils import (Counter, set_cuda_visible_devices, get_ip,
                         get_open_port, get_distributed_init_method)
-
+from vllm.usage.usage_lib import UsageContext, is_usage_stats_enabled, usage_message
 if ray:
     from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 
@@ -72,19 +72,20 @@ class LLMEngine:
         placement_group: Ray placement group for distributed execution.
             Required for distributed execution.
         log_stats: Whether to log statistics.
+        usage_context: Specified entry point, used for usage info collection
     """
 
     def __init__(
-        self,
-        model_config: ModelConfig,
-        cache_config: CacheConfig,
-        parallel_config: ParallelConfig,
-        scheduler_config: SchedulerConfig,
-        device_config: DeviceConfig,
-        lora_config: Optional[LoRAConfig],
-        placement_group: Optional["PlacementGroup"],
-        log_stats: bool,
-    ) -> None:
+            self,
+            model_config: ModelConfig,
+            cache_config: CacheConfig,
+            parallel_config: ParallelConfig,
+            scheduler_config: SchedulerConfig,
+            device_config: DeviceConfig,
+            lora_config: Optional[LoRAConfig],
+            placement_group: Optional["PlacementGroup"],
+            log_stats: bool,
+            usage_context: UsageContext = UsageContext.ENGINE_CONTEXT) -> None:
         logger.info(
             f"Initializing an LLM engine (v{vllm.__version__}) with config: "
             f"model={model_config.model!r}, "
@@ -118,6 +119,10 @@ def __init__(
         self._init_tokenizer()
         self.seq_counter = Counter()
 
+        #If usage stat is enabled, collect relevant info.
+        if is_usage_stats_enabled():
+            usage_message.report_usage(model_config.model, usage_context)
+
         # Create the parallel GPU workers.
         if self.parallel_config.worker_use_ray:
             # Disable Ray usage stats collection.
@@ -394,7 +399,11 @@ def _init_cache(self) -> None:
         self._run_workers("warm_up_model")
 
     @classmethod
-    def from_engine_args(cls, engine_args: EngineArgs) -> "LLMEngine":
+    def from_engine_args(
+        cls,
+        engine_args: EngineArgs,
+        usage_context: UsageContext = UsageContext.ENGINE_CONTEXT
+    ) -> "LLMEngine":
         """Creates an LLM engine from the engine arguments."""
         # Create the engine configs.
         engine_configs = engine_args.create_engine_configs()
@@ -404,7 +413,8 @@ def from_engine_args(cls, engine_args: EngineArgs) -> "LLMEngine":
         # Create the LLM engine.
         engine = cls(*engine_configs,
                      placement_group,
-                     log_stats=not engine_args.disable_log_stats)
+                     log_stats=not engine_args.disable_log_stats,
+                     usage_context=usage_context)
         return engine
 
     def encode_request(

vllm/entrypoints/api_server.py

@@ -16,6 +16,7 @@
 from vllm.engine.async_llm_engine import AsyncLLMEngine
 from vllm.sampling_params import SamplingParams
 from vllm.utils import random_uuid
+from vllm.usage.usage_lib import UsageContext
 
 TIMEOUT_KEEP_ALIVE = 5  # seconds.
 app = FastAPI()
@@ -87,9 +88,9 @@ async def stream_results() -> AsyncGenerator[bytes, None]:
         help="FastAPI root_path when app is behind a path based routing proxy")
     parser = AsyncEngineArgs.add_cli_args(parser)
     args = parser.parse_args()
-
     engine_args = AsyncEngineArgs.from_cli_args(args)
-    engine = AsyncLLMEngine.from_engine_args(engine_args)
+    engine = AsyncLLMEngine.from_engine_args(
+        engine_args, usage_context=UsageContext.API_SERVER)
 
     app.root_path = args.root_path
     uvicorn.run(app,

vllm/entrypoints/llm.py

@@ -9,6 +9,7 @@
 from vllm.outputs import RequestOutput
 from vllm.sampling_params import SamplingParams
 from vllm.utils import Counter
+from vllm.usage.usage_lib import UsageContext
 
 
 class LLM:
@@ -106,7 +107,8 @@ def __init__(
             disable_custom_all_reduce=disable_custom_all_reduce,
             **kwargs,
         )
-        self.llm_engine = LLMEngine.from_engine_args(engine_args)
+        self.llm_engine = LLMEngine.from_engine_args(
+            engine_args, usage_context=UsageContext.LLM_CLASS)
         self.request_counter = Counter()
 
     def get_tokenizer(

vllm/entrypoints/openai/api_server.py

@@ -23,6 +23,7 @@
 from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
 from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion
 from vllm.entrypoints.openai.serving_engine import LoRA
+from vllm.usage.usage_lib import UsageContext
 
 TIMEOUT_KEEP_ALIVE = 5  # seconds
 
@@ -245,9 +246,9 @@ async def authentication(request: Request, call_next):
         served_model = args.served_model_name
     else:
         served_model = args.model
-
     engine_args = AsyncEngineArgs.from_cli_args(args)
-    engine = AsyncLLMEngine.from_engine_args(engine_args)
+    engine = AsyncLLMEngine.from_engine_args(
+        engine_args, usage_context=UsageContext.OPENAI_API_SERVER)
     openai_serving_chat = OpenAIServingChat(engine, served_model,
                                             args.response_role,
                                             args.lora_modules,

vllm/usage/usage_lib.py

@@ -0,0 +1,139 @@
+import os
+import torch
+import json
+import platform
+import pkg_resources
+import requests
+import datetime
+import psutil
+from threading import Thread
+from pathlib import Path
+from typing import Optional
+from enum import Enum
+
+_xdg_config_home = os.getenv('XDG_CONFIG_HOME',
+                             os.path.expanduser('~/.config'))
+_vllm_internal_path = 'vllm/usage_stats.json'
+
+_USAGE_STATS_FILE = os.path.join(
+    _xdg_config_home,
+    _vllm_internal_path)  #File path to store usage data locally
+_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(
+            os.path.expanduser('~/.config/vllm/do_not_track'))
+
+        _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
+    return "UNKNOWN"
+
+
+class UsageContext(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:
+
+    def __init__(self) -> None:
+        self.gpu_list: Optional[dict] = None
+        self.provider: Optional[str] = None
+        self.architecture: Optional[str] = None
+        self.platform: Optional[str] = None
+        self.model: Optional[str] = None
+        self.vllm_version: Optional[str] = None
+        self.context: Optional[str] = None
+        self.log_time: Optional[int] = None
+        #Logical CPU count
+        self.num_cpu: Optional[int] = None
+        self.cpu_type: Optional[str] = None
+        self.total_memory: Optional[int] = None
+        self.source: Optional[str] = None
+
+    def report_usage(self, model: str, context: UsageContext) -> None:
+        t = Thread(target=usage_message._report_usage, args=(model, context))
+        t.start()
+
+    def _report_usage(self, model: str, context: UsageContext) -> None:
+        self.context = context.value
+        self.gpu_list = []
+        for i in range(torch.cuda.device_count()):
+            device_property = torch.cuda.get_device_properties(i)
+            gpu_name = device_property.name
+            gpu_memory = device_property.total_memory
+            self.gpu_list.append({"name": gpu_name, "memory": gpu_memory})
+        self.provider = _detect_cloud_provider()
+        self.architecture = platform.machine()
+        self.platform = platform.platform()
+        self.vllm_version = pkg_resources.get_distribution("vllm").version
+        self.model = model
+        self.log_time = _get_current_timestamp_ns()
+        self.num_cpu = os.cpu_count()
+        #Best effort reading processor name
+        self.cpu_type = platform.processor()
+        self.total_memory = psutil.virtual_memory().total
+        self.source = os.environ.get("VLLM_USAGE_SOURCE", "production")
+        self._write_to_file()
+        headers = {'Content-type': 'application/x-ndjson'}
+        payload = json.dumps(vars(self))
+        try:
+            requests.post(_USAGE_STATS_SERVER, data=payload, headers=headers)
+        except requests.exceptions.RequestException:
+            print("Usage Log Request Failed")
+
+    def _write_to_file(self):
+        os.makedirs(os.path.dirname(_USAGE_STATS_FILE), exist_ok=True)
+        with open(_USAGE_STATS_FILE, "w+") as outfile:
+            json.dump(vars(self), outfile)
+
+
+usage_message = UsageMessage()