move EventLoopThread from async_engine to pipeline

Author: lvhan028

lmdeploy/pipeline.py

@@ -1,13 +1,19 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import asyncio
+import atexit
+import concurrent.futures
 import os
 from contextlib import closing
+from functools import partial
+from queue import Queue
+from threading import Thread
 from typing import TYPE_CHECKING, Dict, Iterator, List, Tuple
 
 import tqdm
 from typing_extensions import deprecated
 
 from .archs import autoget_backend_config, get_task
-from .messages import GenerationConfig, PytorchEngineConfig, SpeculativeConfig, TurbomindEngineConfig
+from .messages import GenerationConfig, PytorchEngineConfig, Response, SpeculativeConfig, TurbomindEngineConfig
 from .model import ChatTemplateConfig
 from .serve.processors import MultimodalProcessor
 from .utils import get_logger, get_model
@@ -67,8 +73,11 @@ def __init__(self,
                                            max_log_len=max_log_len,
                                            speculative_config=speculative_config,
                                            **kwargs)
+        self.internal_thread = _EventLoopThread(daemon=True)
+        self.limiter: asyncio.Semaphore = None
         self.session_mgr = self.async_engine.session_mgr
         self.backend_config = self.async_engine.backend_config
+        self.async_engine.start_loop(self.internal_thread.loop, use_async_api=False)
 
     def infer(self,
               prompts: List[str] | str | List[Dict] | List[List[Dict]] | Tuple | List[Tuple],
@@ -100,7 +109,7 @@ def infer(self,
                                                adapter_name=adapter_name,
                                                stream_response=False,
                                                **kwargs)
-            for g in self.async_engine._infer(requests, multiplex=False, pbar=pbar):
+            for g in self._infer(requests, multiplex=False, pbar=pbar):
                 res = None
                 for out in g:
                     res = res.extend(out) if res else out
@@ -149,10 +158,11 @@ def stream_infer(self,
                                            adapter_name=adapter_name,
                                            stream_response=stream_response,
                                            **kwargs)
-        return self.async_engine._infer(requests, multiplex=True)
+        return self._infer(requests, multiplex=True)
 
     def close(self):
         """Close the pipeline."""
+        self.internal_thread.close()
         self.async_engine.close()
 
     def chat(self,
@@ -197,7 +207,7 @@ def _gen():
                     resp = resp.extend(out) if resp else out
                     yield out
             except:  # noqa
-                self.async_engine._run(coro=session.async_abort())
+                self._run(coro=session.async_abort())
                 raise
             else:
                 session.response = resp
@@ -295,3 +305,116 @@ def _request_generator(self,
             # Since AsyncEngine.generate defines session_id in the argument lists, here we
             # use session_id to pass the session to the AsyncEngine.generate. It's
             yield dict(session_id=session, messages=prompt, gen_config=gen_cfg, **kwargs)
+
+    def _get_limiter(self):
+        if not self.limiter:
+            self.limiter = asyncio.Semaphore(self.backend_config.max_batch_size)
+        return self.limiter
+
+    def _infer(self, requests: Iterator[Dict], multiplex: bool, pbar=None, loop=None) -> Iterator[Iterator[Response]]:
+
+        async def _sync_resp(g, que: Queue, idx: int, sem: asyncio.Semaphore):
+            async for out in g:
+                que.put(out.to_response(idx))
+            sem.release()
+            if not multiplex:
+                que.put(None)  # sentinel of inner generator
+            if pbar:
+                pbar.update(1)
+
+        que = Queue()
+
+        async def _infer():
+            sem = self._get_limiter()
+            tasks = []
+            for idx, req in enumerate(requests):
+                await sem.acquire()
+                gen = self.async_engine.generate(**req)
+                dst = que if multiplex else Queue()
+                if not multiplex:
+                    que.put(iter(dst.get, None))
+                # create a task to send the responses
+                task = asyncio.create_task(_sync_resp(gen, dst, idx, sem))
+                tasks.append(task)
+            if not multiplex:  # sentinel of outer generator
+                que.put(None)
+            await asyncio.gather(*tasks)
+            if multiplex:
+                que.put(None)  # sentinel of inner generator
+
+        loop = loop or self.internal_thread.loop
+        # submit the coroutine to async world
+        asyncio.run_coroutine_threadsafe(_infer(),
+                                         loop).add_done_callback(lambda f: None if f.cancelled() else f.result())
+
+        return iter(que.get, None)
+
+    def _run(self, fn=None, coro=None):
+        assert (fn or coro) and not (fn and coro)
+        loop = self.internal_thread.loop
+        if fn:
+
+            async def _coro():
+                return fn()
+
+            coro = _coro()
+        return asyncio.run_coroutine_threadsafe(coro, loop)
+
+
+class _EventLoopThread:
+
+    def __init__(self, daemon=False):
+        fut = concurrent.futures.Future()
+        self.thread = Thread(target=partial(self._thread_entry, fut), daemon=daemon)
+        self.thread.start()
+        self.loop: asyncio.AbstractEventLoop = fut.result()
+        self.closed = False
+        if daemon:
+            atexit.register(self.close)
+
+    def _thread_entry(self, fut):
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        fut.set_result(loop)
+        try:
+            loop.run_forever()
+        except BaseException as e:
+            logger.error(f'[internal_thread] {type(e).__name__} {e}')
+        finally:
+            try:
+                self._cancel_all_tasks()
+                loop.run_until_complete(loop.shutdown_asyncgens())
+            finally:
+                asyncio.set_event_loop(None)
+                loop.close()
+
+    def _cancel_all_tasks(self):
+        """Modified from asyncio/runners.py."""
+        to_cancel = asyncio.all_tasks(self.loop)
+        if not to_cancel:
+            return
+
+        for task in to_cancel:
+            task.cancel()
+
+        async def _gather():
+            await asyncio.gather(*to_cancel, return_exceptions=True)
+
+        self.loop.run_until_complete(_gather())
+
+        for task in to_cancel:
+            if task.cancelled():
+                continue
+            if task.exception() is not None:
+                self.loop.call_exception_handler({
+                    'message': 'unhandled exception during worker thread shutdown',
+                    'exception': task.exception(),
+                    'task': task,
+                })
+
+    def close(self):
+        if self.closed:
+            return
+        self.closed = True
+        self.loop.call_soon_threadsafe(self.loop.stop)
+        self.thread.join()

lmdeploy/serve/core/async_engine.py

@@ -1,16 +1,12 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 
 import asyncio
-import atexit
 import concurrent.futures
 import dataclasses
 import random
 from contextlib import asynccontextmanager
 from copy import deepcopy
-from functools import partial
-from queue import Queue
-from threading import Thread
-from typing import Any, Dict, Iterator, List, Literal
+from typing import Any, Dict, List, Literal
 
 from lmdeploy.archs import get_model_arch
 from lmdeploy.logger import RequestLogger
@@ -65,65 +61,6 @@ def to_response(self, index: int = 0) -> Response:
                         index=index)
 
 
-class _EventLoopThread:
-
-    def __init__(self, daemon=False):
-        fut = concurrent.futures.Future()
-        self.thread = Thread(target=partial(self._thread_entry, fut), daemon=daemon)
-        self.thread.start()
-        self.loop: asyncio.AbstractEventLoop = fut.result()
-        self.closed = False
-        if daemon:
-            atexit.register(self.close)
-
-    def _thread_entry(self, fut):
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
-        fut.set_result(loop)
-        try:
-            loop.run_forever()
-        except BaseException as e:
-            logger.error(f'[internal_thread] {type(e).__name__} {e}')
-        finally:
-            try:
-                self._cancel_all_tasks()
-                loop.run_until_complete(loop.shutdown_asyncgens())
-            finally:
-                asyncio.set_event_loop(None)
-                loop.close()
-
-    def _cancel_all_tasks(self):
-        """Modified from asyncio/runners.py."""
-        to_cancel = asyncio.all_tasks(self.loop)
-        if not to_cancel:
-            return
-
-        for task in to_cancel:
-            task.cancel()
-
-        async def _gather():
-            await asyncio.gather(*to_cancel, return_exceptions=True)
-
-        self.loop.run_until_complete(_gather())
-
-        for task in to_cancel:
-            if task.cancelled():
-                continue
-            if task.exception() is not None:
-                self.loop.call_exception_handler({
-                    'message': 'unhandled exception during worker thread shutdown',
-                    'exception': task.exception(),
-                    'task': task,
-                })
-
-    def close(self):
-        if self.closed:
-            return
-        self.closed = True
-        self.loop.call_soon_threadsafe(self.loop.stop)
-        self.thread.join()
-
-
 class AsyncEngine(LogitsMixin):
     """Async inference engine. Maintaining a bunch of tm_model instances.
 
@@ -199,21 +136,18 @@ def __init__(self,
             self.stop_words = self.stop_words[0][0].tolist()
         self.backend = backend
         self.request_logger = RequestLogger(max_log_len)
-        self.internal_thread = _EventLoopThread(daemon=True)
-        self.limiter: asyncio.Semaphore = None
+
         self.num_spec_token = 0 if backend == 'turbomind' or speculative_config is None \
             else speculative_config.num_speculative_tokens
 
         self.session_mgr = SessionManager()
-        self.session_mgr.attach_event_loop(self.internal_thread.loop)
         self.session_mgr.build_request_handle_pool(self.engine, self.backend_config.max_batch_size)
 
         # build stat loggers
         self._build_stat_loggers()
         self.epoch = 0
 
     def close(self):
-        self.internal_thread.close()
         self.session_mgr.clear()
         self.engine.close()
 
@@ -303,49 +237,6 @@ def wakeup(self, tags: List[str] | None = None):
         self.sleeping_tags = self.sleeping_tags - set(tags)
         self.is_sleeping = bool(self.sleeping_tags)
 
-    def _get_limiter(self):
-        if not self.limiter:
-            self.limiter = asyncio.Semaphore(self.backend_config.max_batch_size)
-        return self.limiter
-
-    def _infer(self, requests: Iterator[Dict], multiplex: bool, pbar=None, loop=None) -> Iterator[Iterator[Response]]:
-
-        async def _sync_resp(g, que: Queue, idx: int, sem: asyncio.Semaphore):
-            async for out in g:
-                que.put(out.to_response(idx))
-            sem.release()
-            if not multiplex:
-                que.put(None)  # sentinel of inner generator
-            if pbar:
-                pbar.update(1)
-
-        que = Queue()
-
-        async def _infer():
-            sem = self._get_limiter()
-            tasks = []
-            for idx, req in enumerate(requests):
-                await sem.acquire()
-                gen = self.generate(**req)
-                dst = que if multiplex else Queue()
-                if not multiplex:
-                    que.put(iter(dst.get, None))
-                # create a task to send the responses
-                task = asyncio.create_task(_sync_resp(gen, dst, idx, sem))
-                tasks.append(task)
-            if not multiplex:  # sentinel of outer generator
-                que.put(None)
-            await asyncio.gather(*tasks)
-            if multiplex:
-                que.put(None)  # sentinel of inner generator
-
-        loop = loop or self.internal_thread.loop
-        # submit the coroutine to async world
-        asyncio.run_coroutine_threadsafe(_infer(),
-                                         loop).add_done_callback(lambda f: None if f.cancelled() else f.result())
-
-        return iter(que.get, None)
-
     def _determine_gen_config(self, session, input_ids, gen_config: GenerationConfig | None = None) -> GenerationConfig:
         """Determine the generation configuration."""
         gen_config = deepcopy(gen_config) or GenerationConfig()
@@ -640,18 +531,7 @@ def is_error(status):
         #         await session.async_close()
         #     self.session_mgr.remove(session)
 
-    def _run(self, fn=None, coro=None, loop=None):
-        assert (fn or coro) and not (fn and coro)
-        loop = loop or self.internal_thread.loop
-        if fn:
-
-            async def _coro():
-                return fn()
-
-            coro = _coro()
-        return asyncio.run_coroutine_threadsafe(coro, loop)
-
-    def start_loop(self, use_async_api=False):
+    def start_loop(self, loop, use_async_api=False):
         """Start engine loop.
 
         When using pytorch backend with dp > 1, all dp_rank should receive at least one request before it can start
@@ -661,6 +541,7 @@ def start_loop(self, use_async_api=False):
         The purpose of this function is to allow users to choose whether to use the synchronous interface or the
         asynchronous interface for the pipeline.
         """
+        self.session_mgr.attach_event_loop(loop)
         if hasattr(self.engine, 'start_loop'):
             if use_async_api:
                 return self.engine.start_loop()
@@ -671,7 +552,7 @@ def _start_loop(fut):
                     res = self.engine.start_loop()
                     fut.set_result(res)
 
-                self.internal_thread.loop.call_soon_threadsafe(_start_loop, fut)
+                loop.call_soon_threadsafe(_start_loop, fut)
                 return fut.result()
         else:
             return True

lmdeploy/serve/openai/api_server.py

@@ -1263,7 +1263,7 @@ def dummy_get_device_id():
 @router.on_event('startup')
 async def startup_event():
     async_engine = VariableInterface.async_engine
-    async_engine.start_loop(use_async_api=True)
+    async_engine.start_loop(asyncio.get_running_loop(), use_async_api=True)
 
     if VariableInterface.proxy_url is None:
         return