- refactor root register and root set weights

cybertensor/cwtensor.py

@@ -254,10 +254,11 @@ def __repr__(self) -> str:
         return self.__str__()
 
     # TODO check decorator and improve error handling
-    @retry(delay=2, tries=3, backoff=2, max_delay=4)
+    @retry(delay=3, tries=3, backoff=2, max_delay=8)
     def make_call_with_retry(self, wait_for_finalization: bool,
                              msg: dict,
-                             signer_wallet: LocalWallet, error,
+                             signer_wallet: LocalWallet,
+                             error,
                              gas: Optional[int] = cybertensor.__default_gas__,
                              funds: Optional[str] = None) -> Optional[bool]:
         if not wait_for_finalization:
@@ -275,7 +276,49 @@ def make_call_with_retry(self, wait_for_finalization: bool,
             except Exception as e:
                 raise error(e.__str__())
 
-    @retry(delay=2, tries=3, backoff=2, max_delay=4)
+    def _execute_contract(self, wait_for_finalization: bool,
+                          msg: dict,
+                          wallet: Wallet,
+                          error,
+                          logging_prefix: str,
+                          use_hotkey: bool = True,
+                          success_text: str = 'Finalized',
+                          exception_text: str = 'Failed',
+                          gas: Optional[int] = cybertensor.__default_gas__,
+                          funds: Optional[str] = None):
+        try:
+            _private_key = wallet.hotkey.private_key if use_hotkey else wallet.coldkey.private_key
+            signer_wallet = LocalWallet(
+                PrivateKey(_private_key), self.address_prefix
+            )
+            res = self.make_call_with_retry(
+                wait_for_finalization=wait_for_finalization,
+                msg=msg,
+                signer_wallet=signer_wallet,
+                error=error,
+                gas=gas,
+                funds=funds)
+            if res is True:
+                console.print(
+                    f":white_heavy_check_mark: [green]{success_text}[/green]"
+                )
+                cybertensor.logging.success(
+                        prefix=logging_prefix,
+                        sufix=f"<green>{success_text}</green>",
+                    )
+                return True
+        except Exception as e:
+            console.print(
+                f":cross_mark: [red]{exception_text}[/red]: error:{e}"
+            )
+            cybertensor.logging.warning(
+                prefix=logging_prefix,
+                sufix=f"[red]{exception_text}[/red]: error:{e}",
+            )
+
+        return False
+
+    @retry(delay=3, tries=3, backoff=2, max_delay=8)
     def make_call_with_retry_2(self, wait_for_finalization: bool,
                                msg: dict, signer_wallet: LocalWallet, gas: Optional[int] = cybertensor.__default_gas__,
                                funds: Optional[str] = None) -> [bool, Optional[str]]:
@@ -432,7 +475,8 @@ def _do_set_weights(
         netuid: int,
         version_key: int = cybertensor.__version_as_int__,
         wait_for_finalization: bool = True,
-    ) -> Tuple[bool, Optional[str]]:  # (success, error_message)
+    ) -> bool:
+
         set_weights_msg = {
             "set_weights": {
                 "netuid": netuid,
@@ -441,14 +485,14 @@ def _do_set_weights(
                 "version_key": version_key,
             }
         }
-        signer_wallet = LocalWallet(
-            PrivateKey(wallet.hotkey.private_key), self.address_prefix
-        )
 
-        return self.make_call_with_retry_2(
+        return self._execute_contract(
             wait_for_finalization=wait_for_finalization,
             msg=set_weights_msg,
-            signer_wallet=signer_wallet)
+            wallet=wallet,
+            error=NotSetWeightError,
+            logging_prefix='Set weights'
+        )
 
     ######################
     #### Registration ####
@@ -927,16 +971,18 @@ def _do_root_register(
         self,
         wallet: "Wallet",
         wait_for_finalization: bool = True,
-    ) -> Tuple[bool, Optional[str]]:
+    ) -> bool:
+
         root_register_msg = {"root_register": {"hotkey": wallet.hotkey.address}}
-        signer_wallet = LocalWallet(
-            PrivateKey(wallet.coldkey.private_key), self.address_prefix
-        )
 
-        return self.make_call_with_retry_2(
+        return self._execute_contract(
             wait_for_finalization=wait_for_finalization,
             msg=root_register_msg,
-            signer_wallet=signer_wallet)
+            wallet=wallet,
+            error=RegistrationError,
+            use_hotkey=False,
+            logging_prefix='root register',
+            exception_text='Neuron was not registered in root')
 
     def root_set_weights(
         self,

cybertensor/errors.py

@@ -77,6 +77,11 @@ class NotDelegateError(StakeError):
     pass
 
 
+class NotSetWeightError(ChainTransactionError):
+    r"""Error raised when a hotkey you are trying to set weight is not set."""
+    pass
+
+
 class KeyFileError(Exception):
     r"""Error thrown when the keyfile is corrupt, non-writable, non-readable or the password used to decrypt is invalid."""
     pass

cybertensor/messages/root.py

@@ -70,25 +70,20 @@ def root_register_message(
             return False
 
     with console.status(":satellite: Registering to root network..."):
-        success, err_msg = cwtensor._do_root_register(
+        success = cwtensor._do_root_register(
             wallet=wallet,
             wait_for_finalization=wait_for_finalization,
         )
-
-        if success != True or success is False:
-            console.print(
-                f":cross_mark: [red]Failed[/red]: error:{err_msg}"
-            )
-            time.sleep(0.5)
+        time.sleep(0.5)
 
         # Successful registration, final check for neuron and pubkey
-        else:
+        if success is True:
             is_registered = cwtensor.is_hotkey_registered(
                 netuid=0, hotkey=wallet.hotkey.address
             )
             if is_registered:
                 console.print(
-                    ":white_heavy_check_mark: [green]Registered[/green]"
+                    ":white_heavy_check_mark: [green]Registered in root[/green]"
                 )
                 return True
             else:
@@ -97,6 +92,8 @@ def root_register_message(
                     ":cross_mark: [red]Unknown error. Neuron not found.[/red]"
                 )
 
+        return False
+
 
 def set_root_weights_message(
     cwtensor: "cybertensor.cwtensor",
@@ -148,12 +145,8 @@ def set_root_weights_message(
         )
 
     # Normalize the weights to max value.
-    formatted_weights = cybertensor.utils.weight_utils.normalize_max_weight(
-        x=weights, limit=max_weight_limit
-    )
-    console.print(
-        f"\nNormalized weights: \n\t{weights} -> {formatted_weights}\n"
-    )
+    formatted_weights = normalize_max_weight(x=weights, limit=max_weight_limit)
+    console.print(f"\nNormalized weights: \n\t{weights} -> {formatted_weights}\n")
 
     # Ask before moving on.
     if prompt:
@@ -167,49 +160,14 @@ def set_root_weights_message(
     with console.status(
         f":satellite: Setting root weights on [white]{cwtensor.network}[/white] ..."
     ):
-        try:
-            weight_uids, weight_vals = weight_utils.convert_weights_and_uids_for_emit(
-                netuids, weights
-            )
-            success, error_message = cwtensor._do_set_weights(
-                wallet=wallet,
-                netuid=0,
-                uids=weight_uids,
-                vals=weight_vals,
-                version_key=version_key,
-                wait_for_finalization=wait_for_finalization,
-            )
-
-            console.print(success, error_message)
-
-            if not wait_for_finalization:
-                return True
-
-            if success is True:
-                console.print(
-                    ":white_heavy_check_mark: [green]Finalized[/green]"
-                )
-                cybertensor.logging.success(
-                    prefix="Set weights",
-                    sufix="<green>Finalized: </green>" + str(success),
-                )
-                return True
-            else:
-                console.print(
-                    f":cross_mark: [red]Failed[/red]: error:{error_message}"
-                )
-                cybertensor.logging.warning(
-                    prefix="Set weights",
-                    sufix=f"<red>Failed: </red>{error_message}",
-                )
-                return False
-
-        except Exception as e:
-            # TODO( devs ): lets remove all of the cybertensor.__console__ calls and replace with loguru.
-            console.print(
-                f":cross_mark: [red]Failed[/red]: error:{e}"
-            )
-            cybertensor.logging.warning(
-                prefix="Set weights", sufix=f"<red>Failed: </red>{e}"
-            )
-            return False
+        weight_uids, weight_vals = convert_weights_and_uids_for_emit(
+            netuids, weights
+        )
+        return cwtensor._do_set_weights(
+            wallet=wallet,
+            netuid=0,
+            uids=weight_uids,
+            vals=weight_vals,
+            version_key=version_key,
+            wait_for_finalization=wait_for_finalization,
+        )