xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-01-06 04:21:49 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Bugfix] Use temporary directory in registry (#9721)

Browse Source
This commit is contained in:
Cyrus Leung 2024-10-29 13:08:20 +08:00 committed by GitHub
parent e74f2d448c
commit eae3d48181
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 8 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
vllm/model_executor/models/registry.py
View File

@ -1,4 +1,5 @@
import importlib
import os
import pickle
import subprocess
import sys
@ -423,9 +424,13 @@ _T = TypeVar("_T")
def _run_in_subprocess(fn: Callable[[], _T]) -> _T:
with tempfile.NamedTemporaryFile() as output_file:
# NOTE: We use a temporary directory instead of a temporary file to avoid
# issues like https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file
with tempfile.TemporaryDirectory() as tempdir:
output_filepath = os.path.join(tempdir, "registry_output.tmp")
# `cloudpickle` allows pickling lambda functions directly
input_bytes = cloudpickle.dumps((fn, output_file.name))
input_bytes = cloudpickle.dumps((fn, output_filepath))
# cannot use `sys.executable __file__` here because the script
# contains relative imports
@ -442,7 +447,7 @@ def _run_in_subprocess(fn: Callable[[], _T]) -> _T:
raise RuntimeError(f"Error raised in subprocess:\n"
f"{returned.stderr.decode()}") from e
with open(output_file.name, "rb") as f:
with open(output_filepath, "rb") as f:
return pickle.load(f)