call destroy_distributed_environment atexit

Signed-off-by: Tyler Michael Smith <[email protected]>

vllm/distributed/parallel_state.py

@@ -1169,16 +1169,6 @@ def destroy_model_parallel():
     _PP = None
 
 
-# In V1, Calling _TP.destroy() results in 2 leaked shared memory objects. This
-# is related to the torch.distributed.destroy_process_group calls. However, not
-# cleaning up its mq_broadcaster results in 1 leaked shm object.
-# TODO: Fix up this hack
-def destroy_tp_mq_broadcaster():
-    global _TP
-    if _TP:
-        _TP.mq_broadcaster = None
-
-
 def destroy_distributed_environment():
     global _WORLD
     if _WORLD:

vllm/v1/worker/gpu_worker.py

@@ -1,4 +1,5 @@
 """A GPU worker class."""
+import atexit
 import gc
 import os
 import pickle
@@ -13,7 +14,7 @@
 
 import vllm.envs as envs
 from vllm.config import CacheConfig, ModelConfig, ParallelConfig, VllmConfig
-from vllm.distributed import (destroy_tp_mq_broadcaster,
+from vllm.distributed import (destroy_distributed_environment,
                               ensure_model_parallel_initialized,
                               init_distributed_environment,
                               set_custom_all_reduce)
@@ -392,6 +393,7 @@ def make_worker_process(
     @staticmethod
     def run_worker(*args, **kwargs):
         """Launch Worker busy loop in background process."""
+        atexit.register(destroy_distributed_environment)
 
         try:
             worker = WorkerProc(*args, **kwargs)
@@ -403,7 +405,6 @@ def run_worker(*args, **kwargs):
 
             # Clean up once worker exits busy loop
             worker = None
-            destroy_tp_mq_broadcaster()
 
         except KeyboardInterrupt:
             logger.debug("Worker interrupted.")