[BugFix] Fix multiprocessing shutdown errors

vllm/engine/llm_engine.py

@@ -260,98 +260,113 @@ def __init__(
             prompt_adapter_config=prompt_adapter_config,
         )
 
-        if not self.model_config.embedding_mode:
-            self._initialize_kv_caches()
-
-        # If usage stat is enabled, collect relevant info.
-        if is_usage_stats_enabled():
-            from vllm.model_executor.model_loader import (
-                get_architecture_class_name)
-            usage_message.report_usage(
-                get_architecture_class_name(model_config),
-                usage_context,
-                extra_kvs={
-                    # Common configuration
-                    "dtype":
-                    str(model_config.dtype),
-                    "tensor_parallel_size":
-                    parallel_config.tensor_parallel_size,
-                    "block_size":
-                    cache_config.block_size,
-                    "gpu_memory_utilization":
-                    cache_config.gpu_memory_utilization,
-
-                    # Quantization
-                    "quantization":
-                    model_config.quantization,
-                    "kv_cache_dtype":
-                    str(cache_config.cache_dtype),
-
-                    # Feature flags
-                    "enable_lora":
-                    bool(lora_config),
-                    "enable_prompt_adapter":
-                    bool(prompt_adapter_config),
-                    "enable_prefix_caching":
-                    cache_config.enable_prefix_caching,
-                    "enforce_eager":
-                    model_config.enforce_eager,
-                    "disable_custom_all_reduce":
-                    parallel_config.disable_custom_all_reduce,
-                })
-
-        if self.tokenizer:
-            # Ping the tokenizer to ensure liveness if it runs in a
-            # different process.
-            self.tokenizer.ping()
-
-        # Create the scheduler.
-        # NOTE: the cache_config here have been updated with the numbers of
-        # GPU and CPU blocks, which are profiled in the distributed executor.
-        self.scheduler = [
-            Scheduler(scheduler_config, cache_config, lora_config,
-                      parallel_config.pipeline_parallel_size)
-            for _ in range(parallel_config.pipeline_parallel_size)
-        ]
-
-        # Metric Logging.
-        if self.log_stats:
-            if stat_loggers is not None:
-                self.stat_loggers = stat_loggers
-            else:
-                self.stat_loggers = {
-                    "logging":
-                    LoggingStatLogger(
-                        local_interval=_LOCAL_LOGGING_INTERVAL_SEC),
-                    "prometheus":
-                    PrometheusStatLogger(
-                        local_interval=_LOCAL_LOGGING_INTERVAL_SEC,
-                        labels=dict(model_name=model_config.served_model_name),
-                        max_model_len=self.model_config.max_model_len),
-                }
-                self.stat_loggers["prometheus"].info("cache_config",
-                                                     self.cache_config)
-
-        self.tracer = None
-        if self.observability_config.otlp_traces_endpoint:
-            self.tracer = init_tracer(
-                "vllm.llm_engine",
-                self.observability_config.otlp_traces_endpoint)
-
-        # Create sequence output processor, e.g. for beam search or
-        # speculative decoding.
-        self.output_processor = (
-            SequenceGroupOutputProcessor.create_output_processor(
-                self.scheduler_config,
-                self.detokenizer,
-                self.scheduler,
-                self.seq_counter,
-                self.get_tokenizer_for_seq,
-                stop_checker=StopChecker(
-                    self.scheduler_config.max_model_len,
-                    self.get_tokenizer_for_seq,
-                ),
-            ))
+        init_success = False
+        try:
+            if not self.model_config.embedding_mode:
+                self._initialize_kv_caches()
+
+            # If usage stat is enabled, collect relevant info.
+            if is_usage_stats_enabled():
+                from vllm.model_executor.model_loader import (
+                    get_architecture_class_name)
+                usage_message.report_usage(
+                    get_architecture_class_name(model_config),
+                    usage_context,
+                    extra_kvs={
+                        # Common configuration
+                        "dtype":
+                        str(model_config.dtype),
+                        "tensor_parallel_size":
+                        parallel_config.tensor_parallel_size,
+                        "block_size":
+                        cache_config.block_size,
+                        "gpu_memory_utilization":
+                        cache_config.gpu_memory_utilization,
+
+                        # Quantization
+                        "quantization":
+                        model_config.quantization,
+                        "kv_cache_dtype":
+                        str(cache_config.cache_dtype),
+
+                        # Feature flags
+                        "enable_lora":
+                        bool(lora_config),
+                        "enable_prompt_adapter":
+                        bool(prompt_adapter_config),
+                        "enable_prefix_caching":
+                        cache_config.enable_prefix_caching,
+                        "enforce_eager":
+                        model_config.enforce_eager,
+                        "disable_custom_all_reduce":
+                        parallel_config.disable_custom_all_reduce,
+                    })
+
+            if self.tokenizer:
+                # Ping the tokenizer to ensure liveness if it runs in a
+                # different process.
+                self.tokenizer.ping()
+
+            # Create the scheduler.
+            # NOTE: the cache_config here have been updated with the numbers of
+            # GPU and CPU blocks, which are profiled in the distributed executor.
+            self.scheduler = [
+                Scheduler(scheduler_config, cache_config, lora_config,
+                          parallel_config.pipeline_parallel_size)
+                for _ in range(parallel_config.pipeline_parallel_size)
+            ]
+
+            # Metric Logging.
+            if self.log_stats:
+                if stat_loggers is not None:
+                    self.stat_loggers = stat_loggers
+                else:
+                    self.stat_loggers = {
+                        "logging":
+                        LoggingStatLogger(
+                            local_interval=_LOCAL_LOGGING_INTERVAL_SEC),
+                        "prometheus":
+                        PrometheusStatLogger(
+                            local_interval=_LOCAL_LOGGING_INTERVAL_SEC,
+                            labels=dict(
+                                model_name=model_config.served_model_name),
+                            max_model_len=self.model_config.max_model_len),
+                    }
+                    self.stat_loggers["prometheus"].info(
+                        "cache_config", self.cache_config)
+
+            self.tracer = None
+            if self.observability_config.otlp_traces_endpoint:
+                self.tracer = init_tracer(
+                    "vllm.llm_engine",
+                    self.observability_config.otlp_traces_endpoint)
+
+            tokenizer_group = self.get_tokenizer_group()
+
+            def get_tokenizer_for_seq(sequence: Sequence) -> AnyTokenizer:
+                return tokenizer_group.get_lora_tokenizer(
+                    sequence.lora_request)
+
+            # Create sequence output processor, e.g. for beam search or
+            # speculative decoding.
+            self.output_processor = (
+                SequenceGroupOutputProcessor.create_output_processor(
+                    self.scheduler_config,
+                    self.detokenizer,
+                    self.scheduler,
+                    self.seq_counter,
+                    get_tokenizer_for_seq,
+                    stop_checker=StopChecker(
+                        self.scheduler_config.max_model_len,
+                        get_tokenizer_for_seq,
+                    ),
+                ))
+            init_success = True
+        finally:
+            if not init_success:
+                # Ensure that model_executor is shut down if LLMEngine init
+                # failed
+                self.model_executor.shutdown()
 
     def _initialize_kv_caches(self) -> None:
         """Initialize the KV cache in the worker(s).
@@ -481,10 +496,6 @@ def get_tokenizer(
     ) -> AnyTokenizer:
         return self.get_tokenizer_group().get_lora_tokenizer(lora_request)
 
-    def get_tokenizer_for_seq(self, sequence: Sequence) -> AnyTokenizer:
-        return self.get_tokenizer_group().get_lora_tokenizer(
-            sequence.lora_request)
-
     def _init_tokenizer(self, **tokenizer_init_kwargs) -> BaseTokenizerGroup:
         init_kwargs = dict(
             tokenizer_id=self.model_config.tokenizer,

vllm/executor/multiproc_gpu_executor.py

@@ -1,8 +1,5 @@
 import asyncio
 import os
-import signal
-import threading
-import weakref
 from functools import partial
 from typing import Any, List, Optional
 
@@ -121,20 +118,6 @@ def _init_executor(self) -> None:
             result_handler.start()
             self.worker_monitor.start()
 
-        # Set up signal handlers to shutdown the executor cleanly
-        # sometimes gc does not work well
-
-        # Use weakref to avoid holding a reference to self
-        ref = weakref.ref(self)
-
-        def shutdown(signum, frame):
-            if executor := ref():
-                executor.shutdown()
-
-        if threading.current_thread() is threading.main_thread():
-            signal.signal(signal.SIGINT, shutdown)
-            signal.signal(signal.SIGTERM, shutdown)
-
         self.driver_worker = self._create_worker(
             distributed_init_method=distributed_init_method)
         self._run_workers("init_device")

vllm/executor/multiproc_worker_utils.py

@@ -76,7 +76,7 @@ class ResultHandler(threading.Thread):
     """Handle results from all workers (in background thread)"""
 
     def __init__(self) -> None:
-        super().__init__(daemon=True)
+        super().__init__(daemon=False)
         self.result_queue = mp.Queue()
         self.tasks: Dict[uuid.UUID, Union[ResultFuture, asyncio.Future]] = {}
 
@@ -100,7 +100,7 @@ class WorkerMonitor(threading.Thread):
 
     def __init__(self, workers: List['ProcessWorkerWrapper'],
                  result_handler: ResultHandler):
-        super().__init__(daemon=True)
+        super().__init__(daemon=False)
         self.workers = workers
         self.result_handler = result_handler
         self._close = False
@@ -111,16 +111,23 @@ def run(self) -> None:
         if not self._close:
             self._close = True
 
-            # Kill / cleanup all workers
-            for worker in self.workers:
-                process = worker.process
-                if process.sentinel in dead_sentinels:
-                    process.join(JOIN_TIMEOUT_S)
-                if process.exitcode is not None and process.exitcode != 0:
-                    logger.error("Worker %s pid %s died, exit code: %s",
-                                 process.name, process.pid, process.exitcode)
+            if not sys.is_finalizing():
+                # Kill / cleanup all workers
+                died_count = 0
+                for worker in self.workers:
+                    process = worker.process
+                    if process.sentinel in dead_sentinels:
+                        process.join(JOIN_TIMEOUT_S)
+                    if process.exitcode is not None and process.exitcode != 0:
+                        died_count += 1
+                        logger.error("Worker %s pid %s died, exit code: %s",
+                                     process.name, process.pid,
+                                     process.exitcode)
+                if died_count < len(self.workers):
+                    logger.info(
+                        "Killing remaining local vLLM worker processes")
+
             # Cleanup any remaining workers
-            logger.info("Killing local vLLM worker processes")
             for worker in self.workers:
                 worker.kill_worker()
             # Must be done after worker task queues are all closed