use a queue for non-blocking mlx inference

9db16f8dca224fc7e19837f66a4706e6ca6db555 · 2024-09-05 15:16:10 +0100 · Alex Cheema
Files touched

M exo/inference/mlx/sharded_inference_engine.py
Diff

commit 9db16f8dca224fc7e19837f66a4706e6ca6db555
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Thu Sep 5 15:16:10 2024 +0100

    use a queue for non-blocking mlx inference
---
 exo/inference/mlx/sharded_inference_engine.py | 40 +++++++++++++++++++++------
 1 file changed, 31 insertions(+), 9 deletions(-)

diff --git a/exo/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
index 1e87e818..17c5ae15 100644
--- a/exo/inference/mlx/sharded_inference_engine.py
+++ b/exo/inference/mlx/sharded_inference_engine.py
@@ -16,32 +16,54 @@ class MLXDynamicShardInferenceEngine(InferenceEngine):
     self.shard_downloader = shard_downloader
     self.model_lock = threading.Lock()
     self.executor = ThreadPoolExecutor(max_workers=1)
+    self.inference_queue = asyncio.Queue()
+    self._worker_task = None
+
+  async def _ensure_worker(self):
+    if self._worker_task is None:
+      self._worker_task = asyncio.create_task(self._inference_worker())
 
   async def infer_prompt(self, request_id: str, shard: Shard, prompt: str, image_str: Optional[str] = None, inference_state: Optional[str] = None) -> (np.ndarray, str, bool):
     await self.ensure_shard(shard)
+    await self._ensure_worker()
+
     if image_str:
       image = await get_image_from_str(image_str)
       inputs = self.tokenizer(prompt, image, return_tensors="np")
       pixel_values = mx.array(inputs["pixel_values"])
       input_ids = mx.array(inputs["input_ids"])
-      output_data = await self._run_inference(request_id, input_ids, pixel_values)
+      output_data = await self._queue_inference(request_id, input_ids, pixel_values)
     else:
       input_ids = mx.array(self.tokenizer.encode(prompt))
-      output_data = await self._run_inference(request_id, input_ids)
+      output_data = await self._queue_inference(request_id, input_ids)
     return output_data, "", output_data.size == 1 and output_data.item() == self.tokenizer.eos_token_id
 
   async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> (np.ndarray, str, bool):
     await self.ensure_shard(shard)
+    await self._ensure_worker()
+
     input_tensor = mx.array(input_data)
-    output_data = await self._run_inference(request_id, input_tensor)
+    output_data = await self._queue_inference(request_id, input_tensor)
     return output_data, "", output_data.size == 1 and output_data.item() == self.tokenizer.eos_token_id
 
-  async def _run_inference(self, request_id: str, *args):
-    with self.model_lock:
-      return await asyncio.get_event_loop().run_in_executor(
-        self.executor,
-        lambda: np.array(self.stateful_sharded_model.step(request_id, *args))
-      )
+  async def _queue_inference(self, request_id: str, *args):
+    future = asyncio.get_running_loop().create_future()
+    await self.inference_queue.put((future, request_id, args))
+    return await future
+
+  async def _inference_worker(self):
+    while True:
+      future, request_id, args = await self.inference_queue.get()
+      try:
+        result = await asyncio.get_running_loop().run_in_executor(
+          self.executor,
+          lambda: np.array(self.stateful_sharded_model.step(request_id, *args))
+        )
+        future.set_result(result)
+      except Exception as e:
+        future.set_exception(e)
+      finally:
+        self.inference_queue.task_done()
 
   async def ensure_shard(self, shard: Shard):
     if self.shard == shard:
← 0ca5c260 run mlx inference engine on a single thread too · back to Exo · simplify non-blocking mlx inference 6881722b →