Skip to content

Commit

Permalink
minor-fix (#2813)
Browse files Browse the repository at this point in the history
  • Loading branch information
lvhan028 authored Nov 25, 2024
1 parent 099721a commit 4e6760e
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion lmdeploy/vl/engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,20 +97,23 @@ def __init__(
self.vision_config = vision_config
self.max_batch_size = vision_config.max_batch_size
torch.cuda.empty_cache()
self._que = asyncio.Queue()
self._que: asyncio.Queue = None
self._loop_task: asyncio.Task = None
if vision_config.thread_safe:
self._create_thread_safe_task()

def _create_thread_safe_task(self):
"""thread safe loop task."""
self._loop = asyncio.new_event_loop()
self._que = asyncio.Queue()

def _work_thread():
asyncio.set_event_loop(self._loop)
self._loop.run_until_complete(self._forward_loop())

thread = Thread(target=_work_thread, daemon=True)
thread.start()
self._loop_thread = thread

def _create_event_loop_task(self):
"""event loop task."""
Expand All @@ -122,6 +125,8 @@ def _create_event_loop_task(self):
def req_que(self):
if self.vision_config.thread_safe:
return self._que
if self._que is None:
self._que = asyncio.Queue()
if self._loop_task is None:
self._create_event_loop_task()
if asyncio.get_event_loop() != self._loop:
Expand Down

0 comments on commit 4e6760e

Please sign in to comment.