Usage Stats Collection

requirements.txt

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

vllm/engine/async_llm_engine.py

@@ -12,6 +12,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__)
 
@@ -611,9 +612,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.UNKNOWN_CONTEXT
+    ) -> "AsyncLLMEngine":
         """Creates an async LLM engine from the engine arguments."""
         # Create the engine configs.
         engine_configs = engine_args.create_engine_configs()
@@ -629,7 +633,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

@@ -21,7 +21,7 @@
 from vllm.transformers_utils.tokenizer import (detokenize_incrementally,
                                                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
 
@@ -63,18 +63,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,
+            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.UNKNOWN_CONTEXT
     ) -> None:
         logger.info(
             "Initializing an LLM engine with config: "
@@ -109,6 +111,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.
@@ -360,7 +366,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.UNKNOWN_CONTEXT
+    ) -> "LLMEngine":
         """Creates an LLM engine from the engine arguments."""
         # Create the engine configs.
         engine_configs = engine_args.create_engine_configs()
@@ -370,7 +380,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

@@ -10,6 +10,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()
@@ -85,9 +86,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)
         self.request_counter = Counter()
 
     def get_tokenizer(

vllm/entrypoints/openai/api_server.py

@@ -23,6 +23,7 @@
 from vllm.logger import init_logger
 from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
 from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion
+from vllm.usage.usage_lib import UsageContext
 
 TIMEOUT_KEEP_ALIVE = 5  # seconds
 
@@ -212,9 +213,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.chat_template)

vllm/usage/usage_lib.py

@@ -0,0 +1,126 @@
+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
+
+_USAGE_STATS_FILE = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)),
+    'usage_stats.json')  #File path to store usage data locally
+_USAGE_STATS_ENABLED = None
+_USAGE_STATS_SEVER = os.environ.get('VLLM_USAGE_STATS_SERVER',
+                                    'https://stats.vllm.ai')
+_USAGE_STATS_URL = "https://vector-dev-server-uzyrqjjayq-uc.a.run.app"  #Placeholder for sending usage data to vector.dev http server
+
+
+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'
+    ]
+    for vendor_file in vendor_files:
+        path = Path(vendor_file)
+        if path.is_file():
+            if 'amazon' in path.read_text().lower():
+                return "AWS"
+            elif 'Microsoft Corporation' in path.read_text():
+                return "AZURE"
+            elif 'Google' in path.read_text():
+                return "GCP"
+            elif 'OracleCloud' in path.read_text():
+                return "OCI"
+    return "UNKNOWN"
+
+
+class UsageContext(Enum):
+    UNKNOWN_CONTEXT = "UNKNOWN_CONTEXT"
+    LLM = "LLM"
+    API_SERVER = "API_SERVER"
+    OPENAI_API_SERVER = "OPENAI_API_SERVER"
+
+
+class UsageMessage:
+
+    def __init__(self) -> None:
+        self.gpu: 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.total_memory: Optional[int] = 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 = dict()
+        for i in range(torch.cuda.device_count()):
+            k = torch.cuda.get_device_properties(i).name
+            if k in self.gpu:
+                self.gpu[k] += 1
+            else:
+                self.gpu[k] = 1
+        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()
+        self.total_memory = psutil.virtual_memory().total
+        self._write_to_file()
+        headers = {'Content-type': 'application/json'}
+        payload = json.dumps(vars(self))
+        try:
+            requests.post(_USAGE_STATS_URL, data=payload, headers=headers)
+        except requests.exceptions.RequestException:
+            print("Usage Log Request Failed")
+
+    def _write_to_file(self):
+        with open(_USAGE_STATS_FILE, "w") as outfile:
+            json.dump(vars(self), outfile)
+
+
+usage_message = UsageMessage()