[V1][Core] Use weakref.finalize instead of atexit (#11242)

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>
This commit is contained in:
Tyler Michael Smith 2024-12-17 01:11:33 -05:00 committed by GitHub
parent e88db68cf5
commit 2bfdbf2a36
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 5 additions and 18 deletions

View File

@ -1,5 +1,5 @@
import atexit
import os
import weakref
from typing import List, Optional
import msgspec
@ -165,15 +165,9 @@ class MPClient(EngineCoreClient):
ready_path=ready_path, # type: ignore[misc]
**kwargs,
)
atexit.register(self.shutdown)
self._finalizer = weakref.finalize(self, self.shutdown)
def shutdown(self):
# During final garbage collection in process shutdown, atexit may be
# None.
if atexit:
# in case shutdown gets called via __del__ first
atexit.unregister(self.shutdown)
# Shut down the zmq context.
self.ctx.destroy(linger=0)
@ -197,9 +191,6 @@ class MPClient(EngineCoreClient):
os.remove(socket_file)
self.proc_handle = None
def __del__(self):
self.shutdown()
class SyncMPClient(MPClient):
"""Synchronous client for multi-proc EngineCore."""

View File

@ -1,9 +1,9 @@
import atexit
import os
import pickle
import signal
import sys
import time
import weakref
from dataclasses import dataclass
from enum import Enum, auto
from multiprocessing.process import BaseProcess
@ -37,7 +37,7 @@ class MultiprocExecutor(Executor):
def __init__(self, vllm_config: VllmConfig) -> None:
# Call self.shutdown at exit to clean up
# and ensure workers will be terminated.
atexit.register(self.shutdown)
self._finalizer = weakref.finalize(self, self.shutdown)
self.vllm_config = vllm_config
self.parallel_config = vllm_config.parallel_config
@ -195,14 +195,10 @@ class MultiprocExecutor(Executor):
os.remove(socket_path)
def shutdown(self):
if atexit:
# in case shutdown was called explicitly, we don't need to call it
# again
atexit.unregister(self.shutdown)
"""Properly shut down the executor and its workers"""
if getattr(self, 'shutting_down', False):
self.shutting_down = True
for w in self.workers: #TODO: not sure if needed
for w in self.workers:
w.worker_response_mq = None
self._ensure_worker_termination()