Skip to content

Commit

Permalink
[Bugfix] Use temporary directory in registry (vllm-project#9721)
Browse files Browse the repository at this point in the history
  • Loading branch information
DarkLight1337 authored and Linkun Chen committed Nov 4, 2024
1 parent 7be802c commit 48816c1
Showing 1 changed file with 8 additions and 3 deletions.
11 changes: 8 additions & 3 deletions vllm/model_executor/models/registry.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import importlib
import os
import pickle
import subprocess
import sys
Expand Down Expand Up @@ -423,9 +424,13 @@ def is_attention_free_model(self, architectures: Union[str,


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
Expand All @@ -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)


Expand Down

0 comments on commit 48816c1

Please sign in to comment.