Implement mmalloc optimization

Author: dyyoungg

.pre-commit-config.yaml

@@ -1,12 +1,12 @@
 repos:
--   repo: https://github.com/psf/black
+-   repo: git@github.com:psf/black.git
     rev: 21.12b0
     hooks:
     -   id: black
         language_version: python3
         args: [--line-length=120]
         additional_dependencies: ['click==8.0.4']
--   repo: https://github.com/pycqa/flake8
+-   repo: git@github.com:pycqa/flake8.git
     rev: 3.9.0
     hooks:
       - id: flake8

lightllm/server/api_cli.py

@@ -296,6 +296,10 @@ def make_argument_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--cache_capacity", type=int, default=200, help="cache server capacity for multimodal resources"
     )
+    parser.add_argument(
+        "--enable_concurrent_alloc", action="store_true", help="alloc multimodal resources in threadpool to save time"
+    )
+    parser.add_argument("--concurrent_alloc_workers", type=int, default=4, help="max concurrent threadpool workers")
     parser.add_argument(
         "--data_type",
         type=str,

lightllm/server/audioserver/manager.py

@@ -96,7 +96,9 @@ async def loop_for_fwd(self):
                     multimodal_params = group_req_indexes.multimodal_params
 
                     audio_uuids = [audio.uuid for audio in multimodal_params.audios]
-                    ready_audio = obtain(self.cache_client.root.get_items_embed(audio_uuids))
+                    audio_uuids = pickle.dumps(audio_uuids)
+                    ready_audio = self.cache_client.root.get_items_embed_v2(audio_uuids)
+                    ready_audio = pickle.loads(ready_audio)
 
                     for audio, ready in zip(multimodal_params.audios, ready_audio):
                         if not ready:

lightllm/server/embed_cache/manager.py

@@ -5,6 +5,7 @@
 from lightllm.utils.graceful_utils import graceful_registry
 from lightllm.server.embed_cache.impl.naive_memory_cache import InMemoryCache
 from rpyc.utils.classic import obtain
+import pickle
 
 
 class CacheServer(rpyc.Service):
@@ -48,6 +49,47 @@ def exposed_get_items_embed(self, ids: list[int]) -> list[bool]:
         ids = obtain(ids)
         return self._impl.get_items_embed(ids)
 
+    def exposed_alloc_v2(self, batch_md5_token_nums: bytes) -> bytes:
+        """
+        batch_md5_token_nums: pickle.dumps([(md5sum, token_num), ...])
+        返回: pickle.dumps(records)
+        """
+        batch_requests = pickle.loads(batch_md5_token_nums)
+        md5sum_list = [obtain(md5) for md5, num in batch_requests]
+        token_num_list = [obtain(num) for md5, num in batch_requests]
+        record = self._impl.alloc(md5sum_list, token_num_list)
+        return pickle.dumps(record)
+
+    def exposed_release_v2(self, ids_blob: bytes) -> None:
+        ids = pickle.loads(ids_blob)
+        ids = [obtain(id) for id in ids]
+        return self._impl.release(ids)
+
+    def exposed_set_items_data_v2(self, ids_blob: bytes) -> bytes:
+        ids = pickle.loads(ids_blob)
+        ids = [obtain(id) for id in ids]
+        status_list = self._impl.set_items_data(ids)
+        return pickle.dumps(status_list)
+
+    def exposed_get_items_data_v2(self, ids_blob: bytes) -> bytes:
+        ids = pickle.loads(ids_blob)
+        ids = [obtain(id) for id in ids]
+        status_list = self._impl.get_items_data(ids)
+        return pickle.dumps(status_list)
+
+    def exposed_set_items_embed_v2(self, ids_blob: bytes) -> None:
+
+        ids = pickle.loads(ids_blob)
+        ids = [obtain(id) for id in ids]
+        status_list = self._impl.set_items_embed(ids)
+        return pickle.dumps(status_list)
+
+    def exposed_get_items_embed_v2(self, ids_blob: bytes) -> bytes:
+        ids = pickle.loads(ids_blob)
+        ids = [obtain(id) for id in ids]
+        status_list = self._impl.get_items_embed(ids)
+        return pickle.dumps(status_list)
+
 
 def start_cache_manager(port: int, args, pipe_writer):
     # 注册graceful 退出的处理

lightllm/server/httpserver/manager.py

@@ -10,6 +10,7 @@
 import datetime
 import pickle
 from frozendict import frozendict
+import concurrent.futures
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 from typing import Union, List, Tuple, Dict, Optional
@@ -32,6 +33,7 @@
 from lightllm.utils.statics_utils import MovingAverage
 from lightllm.utils.config_utils import get_vocab_size
 from lightllm.utils.envs_utils import get_unique_server_name
+from lightllm.utils.infer_utils import calculate_cpu_time_async, calculate_cpu_time_sync
 from rpyc.utils.classic import obtain
 
 logger = init_logger(__name__)
@@ -112,13 +114,19 @@ def __init__(
         # If the timemark is not updated for a pre-set time, a prob request will be sent to the backend.
         self.latest_success_infer_time_mark = SharedInt(f"{get_unique_server_name()}_latest_success_infer_time_mark")
         self.latest_success_infer_time_mark.set_value(int(time.time()))
+
+        # 线程池用于创建multimodal resource alloc
+        self.enable_concurrent_alloc = self.args.enable_concurrent_alloc
+        self.max_concurrent = self.args.concurrent_alloc_workers * 48
+        if self.enable_concurrent_alloc:
+            self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=self.args.concurrent_alloc_workers)
         return
 
     async def _alloc_resource(self, items, md5sums, token_nums, datas):
-
         while True:
+            t1 = time.time()
             records = obtain(self.cache_client.root.alloc(md5sums, token_nums))
-
+            logger.info(f"cache manager batch alloc time: {(time.time() - t1)*1000} ms")
             if records is None:
                 await asyncio.sleep(0.1)
                 continue
@@ -142,37 +150,139 @@ async def _alloc_resource(self, items, md5sums, token_nums, datas):
                 self.cache_client.root.set_items_data(update_data_ids)
             return
 
+    async def _alloc_resource_v2(self, items, md5sums, token_nums, datas):
+        batch_requests = [(md5sum, token_num) for md5sum, token_num in zip(md5sums, token_nums)]
+        while True:
+            t1 = time.time()
+            req_blob = pickle.dumps(batch_requests)
+            res_blob = self.cache_client.root.alloc_v2(req_blob)
+            records = pickle.loads(res_blob)
+            logger.info(f"cache manager batch alloc time: {(time.time() - t1)*1000} ms")
+            if records is None:
+                await asyncio.sleep(0.1)
+                continue
+
+            uid_list = []
+            for item, rec in zip(items, records):
+                item.uuid = rec["id"]
+                item.token_id = rec["token_id"]
+                item.token_num = rec["token_num"]
+                uid_list.append(rec["id"])
+
+            uid_blob = pickle.dumps(uid_list)
+            ready_flags = self.cache_client.root.get_items_data_v2(uid_blob)
+            ready_flags = pickle.loads(ready_flags)
+
+            max_concurrent_shm = min(len(items), self.max_concurrent)  # 限制最大并发
+            semaphore = asyncio.Semaphore(max_concurrent_shm)
+
+            async def create_shm_with_limit(uid, data):
+                async with semaphore:
+                    loop = asyncio.get_event_loop()
+                    return await loop.run_in_executor(self.executor, create_shm, get_shm_name_data(uid), data)
+
+            update_data_ids = []
+            shm_tasks = []
+            for uid, ready, data in zip(uid_list, ready_flags, datas):
+                if not ready:
+                    task = create_shm_with_limit(uid, data)
+                    shm_tasks.append(task)
+                    update_data_ids.append(uid)
+
+            if len(shm_tasks):
+                t_shm = time.time()
+                await asyncio.gather(*shm_tasks)
+                logger.info(f"concurrent create shm time: {(time.time() - t_shm)*1000} ms")
+
+            if update_data_ids:
+                update_dataids_blob = pickle.dumps(update_data_ids)
+                self.cache_client.root.set_items_data_v2(update_dataids_blob)
+            return
+
+    @calculate_cpu_time_async(show=True)
     async def _alloc_multimodal_resources(self, multimodal_params: MultimodalParams, sampling_params: SamplingParams):
         # 只有 P 和 NORMAL 节点需要真的管理多模态资源
         if self.pd_mode.is_P_or_NORMAL():
             # 这里的锁是为了 防止多个含有多张图片的请求 同时申请的record数量 大于cache_capacity，从而造成死锁的问题。
             # 如果不加任何锁，假如请求1和请求2都有6张图片，而cache_capacity为10，
             # 那么如果某一时刻shm中存在请求1的5张图和请求2的5张图，将会资源竞争产生死锁。
             async with self._resource_lock:
-                items, md5sums, tokens_nums, datas = [], [], [], []
-                for img in multimodal_params.images:
-                    self.tokenizer.init_imageitem_extral_params(img, multimodal_params, sampling_params)
-                    data = img.read()
-                    # must after init_imageitem_extral_params
-                    token_num = self.tokenizer.get_image_token_length(img)
-                    md5sum = hashlib.md5(data).hexdigest() + "_" + str(hash(frozendict(img.extra_params)))
-                    md5sums.append(md5sum)
-                    tokens_nums.append(token_num)
-                    datas.append(data)
-                    items.append(img)
-                for audio in multimodal_params.audios:
-                    self.tokenizer.init_audioitem_extral_params(audio, multimodal_params, sampling_params)
-                    data = audio.read()
-                    token_num = self.tokenizer.get_audio_token_length(audio)
-                    md5sum = hashlib.md5(data).hexdigest() + "_" + str(hash(frozendict(audio.extra_params)))
-                    md5sums.append(md5sum)
-                    tokens_nums.append(token_num)
-                    datas.append(data)
-                    items.append(audio)
-
-                await self._alloc_resource(items, md5sums, tokens_nums, datas)
+                if self.enable_concurrent_alloc:
+                    await self._alloc_multimodal_resources_v2(multimodal_params, sampling_params)
+                else:
+                    await self._alloc_multimodal_resources_v1(multimodal_params, sampling_params)
+
         return
 
+    async def _alloc_multimodal_resources_v1(
+        self, multimodal_params: MultimodalParams, sampling_params: SamplingParams
+    ):
+        items, md5sums, tokens_nums, datas = [], [], [], []
+        for img in multimodal_params.images:
+            self.tokenizer.init_imageitem_extral_params(img, multimodal_params, sampling_params)
+            data = img.read()
+            # must after init_imageitem_extral_params
+            token_num = self.tokenizer.get_image_token_length(img)
+            md5sum = hashlib.md5(data).hexdigest() + "_" + str(hash(frozendict(img.extra_params)))
+            md5sums.append(md5sum)
+            tokens_nums.append(token_num)
+            datas.append(data)
+            items.append(img)
+        for audio in multimodal_params.audios:
+            self.tokenizer.init_audioitem_extral_params(audio, multimodal_params, sampling_params)
+            data = audio.read()
+            token_num = self.tokenizer.get_audio_token_length(audio)
+            md5sum = hashlib.md5(data).hexdigest() + "_" + str(hash(frozendict(audio.extra_params)))
+            md5sums.append(md5sum)
+            tokens_nums.append(token_num)
+            datas.append(data)
+            items.append(audio)
+
+        await self._alloc_resource(items, md5sums, tokens_nums, datas)
+
+    async def _alloc_multimodal_resources_v2(
+        self, multimodal_params: MultimodalParams, sampling_params: SamplingParams
+    ):
+        all_items = multimodal_params.images + multimodal_params.audios
+        if not all_items:
+            return
+        loop = asyncio.get_event_loop()
+
+        def _process_item(item, multimodal_params, sampling_params):
+            """初始化item参数、读取数据并计算MD5"""
+            if isinstance(item, ImageItem):  # 图片
+                self.tokenizer.init_imageitem_extral_params(item, multimodal_params, sampling_params)
+            elif isinstance(item, AudioItem):
+                self.tokenizer.init_audioitem_extral_params(item, multimodal_params, sampling_params)
+
+            data = item.read()
+            md5sum = hashlib.md5(data).hexdigest() + "_" + str(hash(frozendict(item.extra_params)))
+            return data, md5sum
+
+        chunk_size = self.max_concurrent  # 可以根据需要调整
+        for i in range(0, len(all_items), chunk_size):
+            chunk = all_items[i : i + chunk_size]
+
+            # 并发处理chunk内的所有item
+            process_tasks = [
+                loop.run_in_executor(self.executor, _process_item, item, multimodal_params, sampling_params)
+                for item in chunk
+            ]
+            chunk_results = await asyncio.gather(*process_tasks)
+            chunk_items, chunk_md5sums, chunk_tokens_nums, chunk_datas = [], [], [], []
+            for j, item in enumerate(chunk):
+                data, md5sum = chunk_results[j]
+                if isinstance(item, ImageItem):
+                    token_num = self.tokenizer.get_image_token_length(item)
+                elif isinstance(item, AudioItem):
+                    token_num = self.tokenizer.get_audio_token_length(item)
+                chunk_items.append(item)
+                chunk_md5sums.append(md5sum)
+                chunk_tokens_nums.append(token_num)
+                chunk_datas.append(data)
+
+            await self._alloc_resource_v2(chunk_items, chunk_md5sums, chunk_tokens_nums, chunk_datas)
+
     async def _release_multimodal_resources(self, multimodal_params: MultimodalParams):
         # 只有 P 和 NORMAL 节点需要真的管理多模态资源
         if self.pd_mode.is_P_or_NORMAL():
@@ -193,7 +303,11 @@ async def _release_multimodal_resources(self, multimodal_params: MultimodalParam
                         audio.token_id = None
                         audio.token_num = None
                 if ids_to_release:
-                    self.cache_client.root.release(ids_to_release)
+                    if self.enable_concurrent_alloc:
+                        release_id_blobs = pickle.dumps(ids_to_release)
+                        self.cache_client.root.release_v2(release_id_blobs)
+                    else:
+                        self.cache_client.root.release(ids_to_release)
         return
 
     def tokens(self, prompt, multimodal_params, samping_params: SamplingParams, kwargs=None):
@@ -341,7 +455,6 @@ async def generate(
         return
 
     async def _log_req_header(self, request_headers, group_request_id: int):
-
         x_request_id = request_headers.get("X-Request-Id", "")
         x_session_id = request_headers.get("X-Session-Id", "")
 
@@ -436,7 +549,6 @@ async def transfer_to_next_module(
         self,
         group_req_objs: Optional[GroupReqObjs] = None,
     ):
-
         if self.pd_mode == NodeRole.P:
             if self.enable_multimodal:
                 self.send_to_visual.send_pyobj(
@@ -483,7 +595,6 @@ async def _wait_to_token_package(
         req_status: "ReqStatus",
         request: Request,
     ):
-
         event = req_status.event
         unfinished_count = sampling_params.best_of
         out_token_counter = 0
@@ -589,7 +700,6 @@ async def recycle_resource_loop(self):
         pre_time_mark = time.time()
 
         while True:
-
             try:
                 await asyncio.wait_for(self.recycle_event.wait(), timeout=0.02)
             except asyncio.TimeoutError:
@@ -660,7 +770,6 @@ async def handle_loop(self):
 
                         for _ in range(read_token_count):
                             if not req.out_tokens_queue.is_empty():
-
                                 text, src_index, special, count_output_tokens = req.out_tokens_queue.peek()
                                 req.cumlogprob += float(req.shm_logprobs.arr[src_index])
                                 metadata = {

lightllm/server/visualserver/manager.py

@@ -122,7 +122,9 @@ async def loop_for_fwd(self):
                     multimodal_params = group_req_indexes.multimodal_params
 
                     img_uuids = [img.uuid for img in multimodal_params.images]
-                    ready_image = obtain(self.cache_client.root.get_items_embed(img_uuids))
+                    img_uuids = pickle.dumps(img_uuids)
+                    ready_image = self.cache_client.root.get_items_embed_v2(img_uuids)
+                    ready_image = pickle.loads(ready_image)
 
                     for img, ready in zip(multimodal_params.images, ready_image):
                         if not ready: