main
py 35 lines 1.28 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from vllm.utils.network_utils import get_file_store_init_method
5 from vllm.v1.executor import UniProcExecutor
6 from vllm.v1.worker.worker_base import WorkerWrapperBase
7
8
9 # This is a dummy executor for patching in test_runai_model_streamer_s3.py.
10 # We cannot use vllm_runner fixture here, because it spawns worker process.
11 # The worker process reimports the patched entities, and the patch is not applied.
12 class RunaiDummyExecutor(UniProcExecutor):
13 def _init_executor(self) -> None:
14 distributed_init_method = get_file_store_init_method()
15
16 local_rank = 0
17 rank = 0
18 is_driver_worker = True
19
20 device_info = self.vllm_config.device_config.device.__str__().split(":")
21 if len(device_info) > 1:
22 local_rank = int(device_info[1])
23
24 worker_rpc_kwargs = dict(
25 vllm_config=self.vllm_config,
26 local_rank=local_rank,
27 rank=rank,
28 distributed_init_method=distributed_init_method,
29 is_driver_worker=is_driver_worker,
30 )
31
32 self.driver_worker = WorkerWrapperBase()
33
34 self.collective_rpc("init_worker", args=([worker_rpc_kwargs],))
35 self.collective_rpc("init_device")