[Bugfix] Add init_cached_hf_modules to RayWorkerWrapper (#4286)

2024-04-24 00:28:35 +08:00
parent d3c8180ac4
commit d87f39e9a9
2 changed files with 8 additions and 1 deletions
--- a/vllm/worker/worker_base.py
+++ b/vllm/worker/worker_base.py
@@ -103,10 +103,15 @@ class WorkerWrapperBase:

    def __init__(self,
                 worker_module_name=None,
-                 worker_class_name=None) -> None:
+                 worker_class_name=None,
+                 trust_remote_code: bool = False) -> None:
        self.worker_module_name = worker_module_name
        self.worker_class_name = worker_class_name
        self.worker = None
+        if trust_remote_code:
+            # note: lazy import to avoid importing torch before initializing
+            from vllm.utils import init_cached_hf_modules
+            init_cached_hf_modules()

    @staticmethod
    def update_environment_variables(envs: Dict[str, str]) -> None: