Add swap_hotkey command to wallet

bittensor/cli.py

@@ -96,6 +96,7 @@
             "regen_hotkey": RegenHotkeyCommand,
             "faucet": RunFaucetCommand,
             "update": UpdateWalletCommand,
+            "swap_hotkey": SwapHotkeyCommand,
             "set_identity": SetIdentityCommand,
             "get_identity": GetIdentityCommand,
         },

bittensor/commands/__init__.py

@@ -65,7 +65,12 @@
 from .stake import StakeCommand, StakeShow
 from .unstake import UnStakeCommand
 from .overview import OverviewCommand
-from .register import PowRegisterCommand, RegisterCommand, RunFaucetCommand
+from .register import (
+    PowRegisterCommand,
+    RegisterCommand,
+    RunFaucetCommand,
+    SwapHotkeyCommand,
+)
 from .delegates import (
     NominateCommand,
     ListDelegatesCommand,

bittensor/commands/register.py

@@ -20,6 +20,7 @@
 import bittensor
 from rich.prompt import Prompt, Confirm
 from .utils import check_netuid_set, check_for_cuda_reg_config
+from copy import deepcopy
 
 from . import defaults
 
@@ -465,3 +466,66 @@ def check_config(config: "bittensor.config"):
             config.wallet.name = str(wallet_name)
         if not config.no_prompt:
             check_for_cuda_reg_config(config)
+
+
+class SwapHotkeyCommand:
+    @staticmethod
+    def run(cli):
+        r"""Swap your hotkey for all registered axons on the network."""
+        wallet = bittensor.wallet(config=cli.config)
+        subtensor = bittensor.subtensor(config=cli.config)
+
+        # This creates an unnecessary amount of extra data, but simplifies implementation.
+        new_config = deepcopy(cli.config)
+        new_config.wallet.hotkey = new_config.wallet.hotkey_b
+        new_wallet = bittensor.wallet(config=new_config)
+
+        subtensor.swap_hotkey(
+            wallet=wallet,
+            new_wallet=new_wallet,
+            wait_for_finalization=False,
+            wait_for_inclusion=True,
+            prompt=False,
+        )
+
+    @staticmethod
+    def add_args(parser: argparse.ArgumentParser):
+        swap_hotkey_parser = parser.add_parser(
+            "swap_hotkey", help="""Swap your associated hotkey."""
+        )
+
+        swap_hotkey_parser.add_argument(
+            "--wallet.hotkey_b",
+            type=str,
+            default=defaults.wallet.hotkey,
+            help="""Name of the new hotkey""",
+            required=False,
+        )
+
+        bittensor.wallet.add_args(swap_hotkey_parser)
+        bittensor.subtensor.add_args(swap_hotkey_parser)
+
+    @staticmethod
+    def check_config(config: "bittensor.config"):
+        if (
+            not config.is_set("subtensor.network")
+            and not config.is_set("subtensor.chain_endpoint")
+            and not config.no_prompt
+        ):
+            config.subtensor.network = Prompt.ask(
+                "Enter subtensor network",
+                choices=bittensor.__networks__,
+                default=defaults.subtensor.network,
+            )
+
+        if not config.is_set("wallet.name") and not config.no_prompt:
+            wallet_name = Prompt.ask("Enter wallet name", default=defaults.wallet.name)
+            config.wallet.name = str(wallet_name)
+
+        if not config.is_set("wallet.hotkey") and not config.no_prompt:
+            hotkey = Prompt.ask("Enter old hotkey name", default=defaults.wallet.hotkey)
+            config.wallet.hotkey = str(hotkey)
+
+        if not config.is_set("wallet.hotkey_b") and not config.no_prompt:
+            hotkey = Prompt.ask("Enter new hotkey name", default=defaults.wallet.hotkey)
+            config.wallet.hotkey_b = str(hotkey)

bittensor/extrinsics/registration.py

@@ -474,3 +474,39 @@ def run_faucet_extrinsic(
 
         except MaxAttemptedException:
             return False, f"Max attempts reached: {max_allowed_attempts}"
+
+
+def swap_hotkey_extrinsic(
+    subtensor: "bittensor.subtensor",
+    wallet: "bittensor.wallet",
+    new_wallet: "bittensor.wallet",
+    wait_for_inclusion: bool = False,
+    wait_for_finalization: bool = True,
+    prompt: bool = False,
+) -> bool:
+    wallet.coldkey  # unlock coldkey
+    if prompt:
+        # Prompt user for confirmation.
+        if not Confirm.ask(
+            f"Swap {wallet.hotkey} for new hotkey: {new_wallet.hotkey}?"
+        ):
+            return False
+
+    with bittensor.__console__.status(":satellite: Swapping hotkeys..."):
+        success, err_msg = subtensor._do_swap_hotkey(
+            wallet=wallet,
+            new_wallet=new_wallet,
+            wait_for_inclusion=wait_for_inclusion,
+            wait_for_finalization=wait_for_finalization,
+        )
+
+        if success != True or success == False:
+            bittensor.__console__.print(
+                ":cross_mark: [red]Failed[/red]: error:{}".format(err_msg)
+            )
+            time.sleep(0.5)
+
+        else:
+            bittensor.__console__.print(
+                f"Hotkey {wallet.hotkey} swapped for new hotkey: {new_wallet.hotkey}"
+            )

bittensor/subtensor.py

@@ -57,6 +57,7 @@
     register_extrinsic,
     burned_register_extrinsic,
     run_faucet_extrinsic,
+    swap_hotkey_extrinsic,
 )
 from .extrinsics.transfer import transfer_extrinsic
 from .extrinsics.set_weights import set_weights_extrinsic
@@ -483,6 +484,24 @@ def register(
             log_verbose=log_verbose,
         )
 
+    def swap_hotkey(
+        self,
+        wallet: "bittensor.wallet",
+        new_wallet: "bittensor.wallet",
+        wait_for_inclusion: bool = False,
+        wait_for_finalization: bool = True,
+        prompt: bool = False,
+    ) -> bool:
+        """Swaps an old hotkey to a new hotkey."""
+        return swap_hotkey_extrinsic(
+            subtensor=self,
+            wallet=wallet,
+            new_wallet=new_wallet,
+            wait_for_inclusion=wait_for_inclusion,
+            wait_for_finalization=wait_for_finalization,
+            prompt=prompt,
+        )
+
     def run_faucet(
         self,
         wallet: "bittensor.wallet",
@@ -634,6 +653,48 @@ def make_substrate_call_with_retry():
 
         return make_substrate_call_with_retry()
 
+    def _do_swap_hotkey(
+        self,
+        wallet: "bittensor.wallet",
+        new_wallet: "bittensor.wallet",
+        wait_for_inclusion: bool = False,
+        wait_for_finalization: bool = True,
+    ) -> Tuple[bool, Optional[str]]:
+        @retry(delay=2, tries=3, backoff=2, max_delay=4)
+        def make_substrate_call_with_retry():
+            with self.substrate as substrate:
+                # create extrinsic call
+                call = substrate.compose_call(
+                    call_module="SubtensorModule",
+                    call_function="swap_hotkey",
+                    call_params={
+                        "hotkey": wallet.hotkey.ss58_address,
+                        "new_hotkey": new_wallet.hotkey.ss58_address,
+                    },
+                )
+                extrinsic = substrate.create_signed_extrinsic(
+                    call=call, keypair=wallet.coldkey
+                )
+                response = substrate.submit_extrinsic(
+                    extrinsic,
+                    wait_for_inclusion=wait_for_inclusion,
+                    wait_for_finalization=wait_for_finalization,
+                )
+
+                # We only wait here if we expect finalization.
+                if not wait_for_finalization and not wait_for_inclusion:
+                    return True
+
+                # process if registration successful, try again if pow is still valid
+                response.process_events()
+                if not response.is_success:
+                    return False, response.error_message
+                # Successful registration
+                else:
+                    return True, None
+
+        return make_substrate_call_with_retry()
+
     ##################
     #### Transfer ####
     ##################