From d74e5c633fb0feaf14a6256d0bab347f4e8a17b2 Mon Sep 17 00:00:00 2001 From: hh-space-invader Date: Tue, 8 Apr 2025 23:50:07 +0200 Subject: [PATCH] improve: Added pool of shared memory and reuse mechanisim --- fastembed/parallel_processor.py | 67 ++++++++++++++++++++++++++------- 1 file changed, 54 insertions(+), 13 deletions(-) diff --git a/fastembed/parallel_processor.py b/fastembed/parallel_processor.py index 043262d..03d073f 100644 --- a/fastembed/parallel_processor.py +++ b/fastembed/parallel_processor.py @@ -4,7 +4,7 @@ from collections import defaultdict from copy import deepcopy from enum import Enum from dataclasses import dataclass -from multiprocessing import shared_memory +from multiprocessing import shared_memory, Manager, Lock from multiprocessing import Queue, get_context from multiprocessing.context import BaseContext from multiprocessing.process import BaseProcess @@ -46,12 +46,54 @@ class Worker: raise NotImplementedError() +class SharedMemoryPool: + def __init__(self, lock: Lock): + self._lock = lock + self._pool: dict[str, tuple[shared_memory.SharedMemory, int, np.dtype]] = {} + self._free_buffers: list[str] = [] + + def allocate(self, size: int, dtype: np.dtype) -> tuple[shared_memory.SharedMemory, str]: + best_match = None + best_size = float("inf") + for buf_name in self._free_buffers: + shm, buf_size, buf_dtype = self._pool[buf_name] + # get best match for needed size + if buf_size >= size and buf_dtype == dtype and buf_size < best_size: + best_match = buf_name + best_size = buf_size + if best_match: + self._free_buffers.remove(best_match) + return self._pool[best_match][0], best_match + shm = shared_memory.SharedMemory(create=True, size=size) + self._pool[shm.name] = (shm, size, dtype) + return shm, shm.name + # if no match found, create new buffer + shm = shared_memory.SharedMemory(create=True, size=size) + self._pool[shm.name] = (shm, size, dtype) + return shm, shm.name + + def release(self, name: str) -> None: + if name in self._pool and name not in self._free_buffers: + self._free_buffers.append(name) + + def cleanup(self) -> None: + for shm, _, _ in self._pool.values(): + shm.close() + shm.unlink() + self._pool.clear() + self._free_buffers.clear() + + def __del__(self): + self.cleanup() + + def _worker( worker_class: Type[Worker], input_queue: Queue, output_queue: Queue, num_active_workers: BaseValue, worker_id: int, + shared_pool: SharedMemoryPool, kwargs: Optional[dict[str, Any]] = None, ) -> None: """ @@ -66,11 +108,9 @@ def _worker( logging.info( f"Reader worker: {worker_id} PID: {os.getpid()} Device: {kwargs.get('device_id', 'CPU')}" ) - shm_buffers = [] try: worker = worker_class.start(**kwargs) - # Keep going until you get an item that's None. def input_queue_iterable() -> Iterable[Any]: while True: item = input_queue.get() @@ -84,19 +124,18 @@ def _worker( for field in ["model_output", "attention_mask", "input_ids"]: array = getattr(output_context, field, None) if array is not None: - # Create shared memory and copy array - shm = shared_memory.SharedMemory(create=True, size=array.nbytes) - shm_buffers.append(shm) + shm, shm_name = shared_pool.allocate(array.nbytes, array.dtype) shm_array = np.ndarray(array.shape, dtype=array.dtype, buffer=shm.buf) np.copyto(shm_array, array) output_metadata[field] = { - "name": shm.name, + "name": shm_name, "shape": array.shape, "dtype": array.dtype.str, } shm.close() output_queue.put((idx, output_metadata)) - shm_buffers.clear() + for field in output_metadata: # mark release to reuse + shared_pool.release(output_metadata[field]["name"]) except Exception as e: # pylint: disable=broad-except logging.exception(e) @@ -110,9 +149,6 @@ def _worker( # See: # https://docs.python.org/3.6/library/multiprocessing.html?highlight=process#pipes-and-queues # https://docs.python.org/3.6/library/multiprocessing.html?highlight=process#programming-guidelines - for shm in shm_buffers: - shm.close() - shm.unlink() input_queue.close() output_queue.close() input_queue.join_thread() @@ -144,6 +180,8 @@ class ParallelWorkerPool: self.device_ids = device_ids self.cuda = cuda self.num_active_workers: Optional[BaseValue] = None + self.manager = Manager() + self.shared_pool = SharedMemoryPool(self.manager.Lock()) def start(self, **kwargs: Any) -> None: self.input_queue = self.ctx.Queue(self.queue_size) @@ -169,6 +207,7 @@ class ParallelWorkerPool: self.output_queue, self.num_active_workers, worker_id, + self.shared_pool, worker_kwargs, ), ) @@ -224,7 +263,6 @@ class ParallelWorkerPool: ).copy() output_arrays[field] = array shm.close() - shm.unlink() yield (idx, OnnxOutputContext(**output_arrays)) read += 1 @@ -248,10 +286,10 @@ class ParallelWorkerPool: array = np.ndarray(meta["shape"], dtype=meta["dtype"], buffer=shm.buf).copy() output_arrays[field] = array shm.close() - shm.unlink() yield (idx, OnnxOutputContext(**output_arrays)) read += 1 finally: + self.shared_pool.cleanup() assert self.input_queue is not None, "Input queue is None" assert self.output_queue is not None, "Output queue is None" self.join() @@ -287,11 +325,13 @@ class ParallelWorkerPool: if process.is_alive(): process.terminate() self.processes.clear() + self.shared_pool.cleanup() def join(self) -> None: for process in self.processes: process.join() self.processes.clear() + self.shared_pool.cleanup() def __del__(self) -> None: """ @@ -306,3 +346,4 @@ class ParallelWorkerPool: for process in self.processes: if process.is_alive(): process.terminate() + self.shared_pool.cleanup()