Dummied up an abstact save_checkpoint

bd3114457fdceae262f9c60e0b1898af9f032cdb · 2024-12-10 13:21:29 -0800 · Nel Nibcord

Files touched

M exo/inference/inference_engine.py
M exo/inference/tinygrad/inference.py

Diff

commit bd3114457fdceae262f9c60e0b1898af9f032cdb
Author: Nel Nibcord <blindcrone@tuta.io>
Date:   Tue Dec 10 13:21:29 2024 -0800

    Dummied up an abstact save_checkpoint
---
 exo/inference/inference_engine.py   | 3 +++
 exo/inference/tinygrad/inference.py | 3 +++
 2 files changed, 6 insertions(+)

diff --git a/exo/inference/inference_engine.py b/exo/inference/inference_engine.py
index 3a867466..85f1e14c 100644
--- a/exo/inference/inference_engine.py
+++ b/exo/inference/inference_engine.py
@@ -29,6 +29,9 @@ class InferenceEngine(ABC):
   @abstractmethod
   async def load_checkpoint(self, shard: Shard, path: str):
     pass
+
+  async def save_checkpoint(self, shard: Shard, path: str):
+    pass
   
   async def save_session(self, key, value):
     self.session[key] = value
diff --git a/exo/inference/tinygrad/inference.py b/exo/inference/tinygrad/inference.py
index edf08cea..e4b76a13 100644
--- a/exo/inference/tinygrad/inference.py
+++ b/exo/inference/tinygrad/inference.py
@@ -95,6 +95,9 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
   async def load_checkpoint(self, shard: Shard, path: str):
     await self.ensure_shard(shard)
   
+  async def save_checkpoint(self, shard: Shard, path: str):
+    await self.ensure_shard(shard)
+  
   async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray) -> np.ndarray:
     await self.ensure_shard(shard)
     def wrap_infer():

← cc66a0b7 Missed one · back to Exo · embed fix 7f0c12a9 →