diff --git a/vllm/envs.py b/vllm/envs.py index 689428ec59..0833949b52 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -64,6 +64,7 @@ if TYPE_CHECKING: VLLM_XLA_USE_SPMD: bool = False VLLM_WORKER_MULTIPROC_METHOD: Literal["fork", "spawn"] = "fork" VLLM_ASSETS_CACHE: str = os.path.join(VLLM_CACHE_ROOT, "assets") + VLLM_ASSETS_CACHE_MODEL_CLEAN: bool = False VLLM_IMAGE_FETCH_TIMEOUT: int = 5 VLLM_VIDEO_FETCH_TIMEOUT: int = 30 VLLM_AUDIO_FETCH_TIMEOUT: int = 10 @@ -699,6 +700,11 @@ environment_variables: dict[str, Callable[[], Any]] = { os.path.join(get_default_cache_root(), "vllm", "assets"), )), + # If the env var is set, we will clean model file in + # this path $VLLM_ASSETS_CACHE/model_streamer/$model_name + "VLLM_ASSETS_CACHE_MODEL_CLEAN": + lambda: bool(int(os.getenv("VLLM_ASSETS_CACHE_MODEL_CLEAN", "0"))), + # Timeout for fetching images when serving multimodal models # Default is 5 seconds "VLLM_IMAGE_FETCH_TIMEOUT": diff --git a/vllm/transformers_utils/runai_utils.py b/vllm/transformers_utils/runai_utils.py index 08466ca19b..355fd60e8d 100644 --- a/vllm/transformers_utils/runai_utils.py +++ b/vllm/transformers_utils/runai_utils.py @@ -5,9 +5,10 @@ import hashlib import os import shutil import signal -import tempfile from typing import Optional +from vllm import envs +from vllm.assets.base import get_cache_dir from vllm.logger import init_logger from vllm.utils import PlaceholderModule @@ -58,20 +59,19 @@ class ObjectStorageModel: """ def __init__(self, url: str) -> None: - for sig in (signal.SIGINT, signal.SIGTERM): - existing_handler = signal.getsignal(sig) - signal.signal(sig, self._close_by_signal(existing_handler)) + if envs.VLLM_ASSETS_CACHE_MODEL_CLEAN: + for sig in (signal.SIGINT, signal.SIGTERM): + existing_handler = signal.getsignal(sig) + signal.signal(sig, self._close_by_signal(existing_handler)) dir_name = os.path.join( - tempfile.gettempdir(), + get_cache_dir(), "model_streamer", hashlib.sha256(str(url).encode()).hexdigest()[:8]) if os.path.exists(dir_name): shutil.rmtree(dir_name) os.makedirs(dir_name) self.dir = dir_name - - def __del__(self): - self._close() + logger.debug("Init object storage, model cache path is: %s", dir_name) def _close(self) -> None: if os.path.exists(self.dir):