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() self._verify_args()
if self.use_beam_search: 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: if not envs.VLLM_NO_DEPRECATION_WARNING:
logger.warning( logger.warning(
"[IMPORTANT] We plan to discontinue the support for beam " "[IMPORTANT] We plan to discontinue the support for beam "
@ -196,6 +200,7 @@ class SamplingParams:
"https://github.com/vllm-project/vllm/issues/6226 for " "https://github.com/vllm-project/vllm/issues/6226 for "
"more information. Set VLLM_NO_DEPRECATION_WARNING=1 to " "more information. Set VLLM_NO_DEPRECATION_WARNING=1 to "
"suppress this warning.") "suppress this warning.")
self._verify_beam_search() self._verify_beam_search()
else: else:
self._verify_non_beam_search() self._verify_non_beam_search()

View File

@ -7,7 +7,7 @@ import time
from enum import Enum from enum import Enum
from pathlib import Path from pathlib import Path
from threading import Thread from threading import Thread
from typing import Any, Dict, Optional from typing import Any, Dict, Optional, Union
from uuid import uuid4 from uuid import uuid4
import cpuinfo import cpuinfo
@ -25,6 +25,13 @@ _USAGE_STATS_DO_NOT_TRACK_PATH = os.path.join(_config_home,
_USAGE_STATS_ENABLED = None _USAGE_STATS_ENABLED = None
_USAGE_STATS_SERVER = envs.VLLM_USAGE_STATS_SERVER _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(): def is_usage_stats_enabled():
"""Determine whether or not we can send usage stats to the server. """Determine whether or not we can send usage stats to the server.
@ -187,7 +194,11 @@ class UsageMessage:
""" """
while True: while True:
time.sleep(600) 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._write_to_file(data)
self._send_to_server(data) self._send_to_server(data)