← back to Exo

Initialize inference engine session in base class

98edb393b2c67ca3645ae16348bbc3ffc42de69e · 2024-12-08 02:34:31 -0800 · Nel Nibcord

Files touched

• M exo/inference/inference_engine.py
• M exo/inference/mlx/sharded_inference_engine.py
• M exo/inference/tinygrad/inference.py

Diff

commit 98edb393b2c67ca3645ae16348bbc3ffc42de69e
Author: Nel Nibcord <blindcrone@tuta.io>
Date:   Sun Dec 8 02:34:31 2024 -0800

    Initialize inference engine session in base class
---
 exo/inference/inference_engine.py             | 2 ++
 exo/inference/mlx/sharded_inference_engine.py | 1 -
 exo/inference/tinygrad/inference.py           | 1 -
 3 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/exo/inference/inference_engine.py b/exo/inference/inference_engine.py
index 2ab5f540..ad14bf8c 100644
--- a/exo/inference/inference_engine.py
+++ b/exo/inference/inference_engine.py
@@ -8,6 +8,8 @@ from .shard import Shard
 
 
 class InferenceEngine(ABC):
+  session = {}
+
   @abstractmethod
   async def encode(self, shard: Shard, prompt: str) -> np.ndarray:
     pass
diff --git a/exo/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
index a0fe7d2d..3ba8370e 100644
--- a/exo/inference/mlx/sharded_inference_engine.py
+++ b/exo/inference/mlx/sharded_inference_engine.py
@@ -42,7 +42,6 @@ class MLXDynamicShardInferenceEngine(InferenceEngine):
     self.shard_downloader = shard_downloader
     self.executor = ThreadPoolExecutor(max_workers=1)
     self.caches = OrderedDict()
-    self.session = {}
 
   async def poll_state(self, request_id: str, max_caches=2):
     if request_id in self.caches:
diff --git a/exo/inference/tinygrad/inference.py b/exo/inference/tinygrad/inference.py
index f8503ba0..343deec3 100644
--- a/exo/inference/tinygrad/inference.py
+++ b/exo/inference/tinygrad/inference.py
@@ -65,7 +65,6 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
     self.shard_downloader = shard_downloader
     self.executor = ThreadPoolExecutor(max_workers=1)
     self.states = OrderedDict()
-    self.session = {}
 
   def poll_state(self, x, request_id: str, max_states=2):
     if request_id not in self.states:

← bcf87e79 Okay let's turn no_grad back on. We'll worry about that when  ·  back to Exo  ·  Some session method cleanup b22c21ac →