[Bugfix] Fix OpenVino/Neuron driver_worker init (#10779)

Signed-off-by: NickLucche <[email protected]> Signed-off-by: Cyrus Leung <[email protected]> Co-authored-by: Cyrus Leung <[email protected]>
vllm-project · Nov 30, 2024 · 40bc242 · 40bc242
1 parent 661175b
commit 40bc242
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/vllm/executor/neuron_executor.py b/vllm/executor/neuron_executor.py
@@ -29,11 +29,13 @@ def _init_worker(self):
         wrapper = WorkerWrapperBase(vllm_config=self.vllm_config)
         distributed_init_method = get_distributed_init_method(
             get_ip(), get_open_port())
-        self.driver_worker = wrapper.init_worker(
+        wrapper.init_worker(
             vllm_config=self.vllm_config,
             local_rank=0,
             rank=0,
-            distributed_init_method=distributed_init_method)
+            distributed_init_method=distributed_init_method,
+        )
+        self.driver_worker = wrapper.worker
         self.driver_worker.init_device()
         self.driver_worker.load_model()
 

diff --git a/vllm/executor/openvino_executor.py b/vllm/executor/openvino_executor.py
@@ -36,7 +36,7 @@ def _init_worker(self):
 
         distributed_init_method = get_distributed_init_method(
             get_ip(), get_open_port())
-        self.driver_worker = wrapper.init_worker(
+        wrapper.init_worker(
             ov_core=ov.Core(),
             vllm_config=self.vllm_config,
             local_rank=0,
@@ -45,6 +45,7 @@ def _init_worker(self):
             kv_cache_dtype=self.cache_config.cache_dtype,
             is_driver_worker=True,
         )
+        self.driver_worker = wrapper.worker
         self.driver_worker.init_device()
         self.driver_worker.load_model()