Use spawn instead of fork method for vllm

pytorch · Sep 11, 2024 · 67e30b0 · 67e30b0
1 parent d6ea6e7
commit 67e30b0
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/ts/torch_handler/vllm_handler.py b/ts/torch_handler/vllm_handler.py
@@ -1,5 +1,6 @@
 import asyncio
 import logging
+import os
 import pathlib
 import time
 from unittest.mock import MagicMock
@@ -40,6 +41,8 @@ def initialize(self, ctx):
         vllm_engine_config = self._get_vllm_engine_config(
             ctx.model_yaml_config.get("handler", {})
         )
+
+        os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
         self.vllm_engine = AsyncLLMEngine.from_engine_args(vllm_engine_config)