make sure mlx stuff is on separate thread non blocking

4a5b80a958cc44e5b566e74fe3b88b0e06ebbdaf · 2025-01-28 18:56:00 +0000 · Alex Cheema

Files touched

M exo/inference/mlx/sharded_inference_engine.py

Diff

commit 4a5b80a958cc44e5b566e74fe3b88b0e06ebbdaf
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Tue Jan 28 18:56:00 2025 +0000

    make sure mlx stuff is on separate thread non blocking
---
 exo/inference/mlx/sharded_inference_engine.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/exo/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
index 54073c55..8bb94118 100644
--- a/exo/inference/mlx/sharded_inference_engine.py
+++ b/exo/inference/mlx/sharded_inference_engine.py
@@ -93,7 +93,11 @@ class MLXDynamicShardInferenceEngine(InferenceEngine):
       )
       output_data, inference_state = result
 
-    output_data = np.array(output_data, copy=False)
+    await self._eval_mlx(output_data)
+    output_data = await asyncio.get_running_loop().run_in_executor(
+      self._mlx_thread,
+      lambda: np.array(output_data, copy=False)
+    )
     return output_data, inference_state
 
   async def evaluate(self, request_id: str, shard: Shard, inputs, targets, lengths, loss: str = "length_masked_ce"):

← 6662d566 load mlx model shard on mlx thread so it doesnt block · back to Exo · propagate prompts to other nodes so they can display them, c af171f06 →