frozenleaves/vllm
1
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2025-10-20 14:53:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

[ci][bugfix] fix kernel tests (#10431)

Browse Source 
Signed-off-by: youkaichao <youkaichao@gmail.com>
This commit is contained in:
youkaichao
2024-11-18 15:29:37 -08:00
committed by GitHub
parent a03ea40792
commit 2298e69b5f
1 changed files with 13 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
vllm/plugins/__init__.py
View File

@ -6,9 +6,6 @@ import vllm.envs as envs
if TYPE_CHECKING:
from vllm.config import CompilationConfig, VllmConfig
else:
CompilationConfig = None
VllmConfig = None
logger = logging.getLogger(__name__)
@ -50,23 +47,23 @@ def load_general_plugins():
logger.exception("Failed to load plugin %s", plugin.name)
_compilation_config: Optional[CompilationConfig] = None
_compilation_config: Optional["CompilationConfig"] = None
def set_compilation_config(config: Optional[CompilationConfig]):
def set_compilation_config(config: Optional["CompilationConfig"]):
global _compilation_config
_compilation_config = config
def get_compilation_config() -> Optional[CompilationConfig]:
def get_compilation_config() -> Optional["CompilationConfig"]:
return _compilation_config
_current_vllm_config: Optional[VllmConfig] = None
_current_vllm_config: Optional["VllmConfig"] = None
@contextmanager
def set_current_vllm_config(vllm_config: VllmConfig):
def set_current_vllm_config(vllm_config: "VllmConfig"):
"""
Temporarily set the current VLLM config.
Used during model initialization.
@ -87,6 +84,12 @@ def set_current_vllm_config(vllm_config: VllmConfig):
_current_vllm_config = old_vllm_config
def get_current_vllm_config() -> VllmConfig:
assert _current_vllm_config is not None, "Current VLLM config is not set."
def get_current_vllm_config() -> "VllmConfig":
if _current_vllm_config is None:
# in ci, usually when we test custom ops/modules directly,
# we don't set the vllm config. In that case, we set a default
# config.
logger.warning("Current VLLM config is not set.")
from vllm.config import VllmConfig
return VllmConfig()
return _current_vllm_config