[Core] Get multiprocessing context at runtime

Instead of getting the multiprocessing context at import time, get it
at runtime. This allows other code in vllm to change this env var and
have it take effect here.

Signed-off-by: Russell Bryant <[email protected]>

Author: russellb

vllm/executor/multiproc_worker_utils.py

@@ -27,9 +27,6 @@
 
 JOIN_TIMEOUT_S = 2
 
-mp_method = envs.VLLM_WORKER_MULTIPROC_METHOD
-mp = multiprocessing.get_context(mp_method)
-
 
 @dataclass
 class Result(Generic[T]):
@@ -77,7 +74,7 @@ class ResultHandler(threading.Thread):
 
     def __init__(self) -> None:
         super().__init__(daemon=True)
-        self.result_queue = mp.Queue()
+        self.result_queue = get_mp_context().Queue()
         self.tasks: Dict[uuid.UUID, Union[ResultFuture, asyncio.Future]] = {}
 
     def run(self):
@@ -147,10 +144,11 @@ class ProcessWorkerWrapper:
 
     def __init__(self, result_handler: ResultHandler,
                  worker_factory: Callable[[], Any]) -> None:
-        self._task_queue = mp.Queue()
+        self.mp = get_mp_context()
+        self._task_queue = self.mp.Queue()
         self.result_queue = result_handler.result_queue
         self.tasks = result_handler.tasks
-        self.process: BaseProcess = mp.Process(  # type: ignore[attr-defined]
+        self.process: BaseProcess = self.mp.Process(  # type: ignore[attr-defined]
             target=_run_worker_process,
             name="VllmWorkerProcess",
             kwargs=dict(
@@ -204,7 +202,7 @@ def _run_worker_process(
     """Worker process event loop"""
 
     # Add process-specific prefix to stdout and stderr
-    process_name = mp.current_process().name
+    process_name = get_mp_context().current_process().name
     pid = os.getpid()
     _add_prefix(sys.stdout, process_name, pid)
     _add_prefix(sys.stderr, process_name, pid)
@@ -269,3 +267,8 @@ def write_with_prefix(s: str):
 
     file.start_new_line = True  # type: ignore[attr-defined]
     file.write = write_with_prefix  # type: ignore[method-assign]
+
+
+def get_mp_context():
+    mp_method = envs.VLLM_WORKER_MULTIPROC_METHOD
+    return multiprocessing.get_context(mp_method)