release vision model

lmdeploy/pytorch/engine/engine.py

@@ -250,9 +250,6 @@ def close(self):
         self._seq_length_buf = None
         self._inputs = None
         torch._C._cuda_clearCublasWorkspaces()
-        torch.cuda.empty_cache()
-        import gc
-        gc.collect()
 
     def _start_loop(self):
         """start loop."""

lmdeploy/serve/async_engine.py

@@ -187,10 +187,15 @@ def __init__(self,
         self.request_logger = RequestLogger(max_log_len)
 
     def close(self):
+        self.gens_set.clear()
         if hasattr(self, 'engine'):
             if isinstance(self.backend_config, PytorchEngineConfig):
                 self.engine.close()
             del self.engine
+            import torch
+            torch.cuda.empty_cache()
+            import gc
+            gc.collect()
 
     def _build_turbomind(
             self,

lmdeploy/serve/vl_async_engine.py

@@ -32,6 +32,10 @@ def __init__(self, model_path: str, **kwargs) -> None:
         self.vl_prompt_template = get_vl_prompt_template(
             model_path, self.chat_template, self.model_name)
 
+    def close(self):
+        self.vl_encoder.close()
+        super().close()
+
     def _convert_prompts(self,
                          prompts: Union[VLPromptType, List[Dict],
                                         List[VLPromptType], List[List[Dict]]]):

lmdeploy/turbomind/turbomind.py

@@ -4,6 +4,7 @@
 import json
 import os.path as osp
 import sys
+import weakref
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import asdict
 from itertools import repeat
@@ -318,7 +319,8 @@ def create_instance(self, cuda_stream_id=0):
         Returns:
             TurboMindInstance: an instance of turbomind
         """
-        return TurboMindInstance(self, self.config, cuda_stream_id)
+        return TurboMindInstance(weakref.proxy(self), self.config,
+                                 cuda_stream_id)
 
 
 class TurboMindInstance:

lmdeploy/vl/engine.py

@@ -95,9 +95,20 @@ def __init__(self,
         torch.cuda.empty_cache()
         self._que: asyncio.Queue = None
         self._loop_task: asyncio.Task = None
+        self._stop = False
         if vision_config.thread_safe:
             self._create_thread_safe_task()
 
+    def close(self):
+        if self.model is not None:
+            self._stop = True
+            if self.vision_config.thread_safe:
+                self._loop_thread.join()
+            else:
+                if hasattr(self, '_loop'):
+                    self._loop.run_until_complete(self._loop_task)
+            self.model = None
+
     def _create_thread_safe_task(self):
         """thread safe loop task."""
         self._loop = asyncio.new_event_loop()
@@ -138,6 +149,8 @@ async def _forward_loop(self):
             while record.total == 0 or (self._que.qsize() and
                                         record.total < self.max_batch_size):
                 while self._que.qsize() == 0:
+                    if self._stop and record.total == 0:
+                        return
                     await asyncio.sleep(0.01)
                 item = await self._que.get()
                 record.enqueue(item[0], item[1], item[2])