Report usage for beam search (#6404)

This commit is contained in:
Simon Mo 2024-07-14 19:37:35 -07:00 committed by GitHub
parent ccb20db8bd
commit 32c9d7f765
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 18 additions and 2 deletions

View File

@ -189,6 +189,10 @@ class SamplingParams:
self._verify_args()
if self.use_beam_search:
# Lazy import to avoid circular imports.
from vllm.usage.usage_lib import set_runtime_usage_data
set_runtime_usage_data("use_beam_search", True)
if not envs.VLLM_NO_DEPRECATION_WARNING:
logger.warning(
"[IMPORTANT] We plan to discontinue the support for beam "
@ -196,6 +200,7 @@ class SamplingParams:
"https://github.com/vllm-project/vllm/issues/6226 for "
"more information. Set VLLM_NO_DEPRECATION_WARNING=1 to "
"suppress this warning.")
self._verify_beam_search()
else:
self._verify_non_beam_search()

View File

@ -7,7 +7,7 @@ import time
from enum import Enum
from pathlib import Path
from threading import Thread
from typing import Any, Dict, Optional
from typing import Any, Dict, Optional, Union
from uuid import uuid4
import cpuinfo
@ -25,6 +25,13 @@ _USAGE_STATS_DO_NOT_TRACK_PATH = os.path.join(_config_home,
_USAGE_STATS_ENABLED = None
_USAGE_STATS_SERVER = envs.VLLM_USAGE_STATS_SERVER
_GLOBAL_RUNTIME_DATA: Dict[str, Union[str, int, bool]] = {}
def set_runtime_usage_data(key: str, value: Union[str, int, bool]) -> None:
"""Set global usage data that will be sent with every usage heartbeat."""
_GLOBAL_RUNTIME_DATA[key] = value
def is_usage_stats_enabled():
"""Determine whether or not we can send usage stats to the server.
@ -187,7 +194,11 @@ class UsageMessage:
"""
while True:
time.sleep(600)
data = {"uuid": self.uuid, "log_time": _get_current_timestamp_ns()}
data = {
"uuid": self.uuid,
"log_time": _get_current_timestamp_ns(),
}
data.update(_GLOBAL_RUNTIME_DATA)
self._write_to_file(data)
self._send_to_server(data)