From c9677f1dc1a3b3d4cb7d32aade201b24130c6839 Mon Sep 17 00:00:00 2001
From: hscott <hscott@foundrydigital.com>
Date: Fri, 6 Dec 2024 11:57:08 -0500
Subject: [PATCH] moved error handling to weight_setter

---
 precog/validators/validator.py     | 15 ---------------
 precog/validators/weight_setter.py | 24 +++++++++++++++++++-----
 2 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/precog/validators/validator.py b/precog/validators/validator.py
index 79aae3c..b8bf9c3 100755
--- a/precog/validators/validator.py
+++ b/precog/validators/validator.py
@@ -1,9 +1,6 @@
 import asyncio
 from pathlib import Path
 
-import bittensor as bt
-import websocket
-
 from precog.utils.classes import Config
 from precog.utils.general import parse_arguments
 from precog.validators.weight_setter import weight_setter
@@ -23,18 +20,6 @@ def __init__(self):
     async def main(self):
         loop = asyncio.get_event_loop()
         self.weight_setter = weight_setter(config=self.config, loop=loop)
-        try:
-            loop.run_forever()
-        except BrokenPipeError:
-            bt.logging.error("Recieved a Broken Pipe substrate error")
-            asyncio.run(self.reset_instance())
-        except websocket._exceptions.WebSocketConnectionClosedException:
-            bt.logging.error("Recieved a websocket closed error, restarting validator")
-            asyncio.run(self.reset_instance())
-        except Exception as e:
-            bt.logging.error(f"Unhandled exception: {e}")
-        finally:
-            bt.logging.info("Exiting Validator")
 
     async def reset_instance(self):
         self.__init__()
diff --git a/precog/validators/weight_setter.py b/precog/validators/weight_setter.py
index 8df0458..7f2b3b9 100755
--- a/precog/validators/weight_setter.py
+++ b/precog/validators/weight_setter.py
@@ -56,6 +56,11 @@ def __init__(self, config=None, loop=None):
         )
         self.loop.create_task(loop_handler(self, self.resync_metagraph, sleep_time=self.resync_metagraph_rate))
         self.loop.create_task(loop_handler(self, self.set_weights, sleep_time=self.set_weights_rate))
+        try:
+            self.loop.run_forever()
+        except Exception as e:
+            bt.logging.error(f"Error on loop: {e}")
+            self.__reset_instance__()
 
     def __exit__(self, exc_type, exc_value, traceback):
         self.save_state()
@@ -63,10 +68,15 @@ def __exit__(self, exc_type, exc_value, traceback):
             pending = asyncio.all_tasks(self.loop)
             for task in pending:
                 task.cancel()
-            asyncio.gather(*pending, return_exceptions=True)
         except Exception as e:
             bt.logging.error(f"Error on __exit__ function: {e}")
-        self.loop.stop()
+        finally:
+            asyncio.gather(*pending, return_exceptions=True)
+            self.loop.stop()
+
+    def __reset_instance__(self):
+        self.__exit__(None, None, None)
+        self.__init__(self.config, self.loop)
 
     async def get_available_uids(self):
         miner_uids = []
@@ -119,9 +129,7 @@ def node_query(self, module, method, params):
     async def set_weights(self):
         try:
             self.current_block = self.subtensor.get_current_block()
-            self.blocks_since_last_update = (
-                self.current_block - self.node_query("SubtensorModule", "LastUpdate", [self.config.netuid])[self.my_uid]
-            )
+            self.blocks_since_last_update = self.current_block - self.last_update
         except Exception:
             bt.logging.error("Failed to get current block, skipping block update")
         if self.blocks_since_last_update >= self.set_weights_rate:
@@ -147,6 +155,12 @@ async def set_weights(self):
             )
             if result:
                 bt.logging.success("✅ Set Weights on chain successfully!")
+                try:
+                    self.last_update = self.node_query("SubtensorModule", "LastUpdate", [self.config.netuid])[
+                        self.my_uid
+                    ]
+                except Exception:
+                    pass
             else:
                 bt.logging.debug(
                     "Failed to set weights this iteration with message:",