Feat/coldkey swap

bittensor/cli.py

@@ -60,6 +60,7 @@
     SubnetLockCostCommand,
     SubnetSudoCommand,
     SwapHotkeyCommand,
+    SwapColdkeyCommand,
     TransferCommand,
     UnStakeCommand,
     UpdateCommand,
@@ -154,6 +155,7 @@
             "faucet": RunFaucetCommand,
             "update": UpdateWalletCommand,
             "swap_hotkey": SwapHotkeyCommand,
+            "swap_coldkey": SwapColdkeyCommand,
             "set_identity": SetIdentityCommand,
             "get_identity": GetIdentityCommand,
             "history": GetWalletHistoryCommand,

bittensor/commands/__init__.py

@@ -70,6 +70,7 @@
     RegisterCommand,
     RunFaucetCommand,
     SwapHotkeyCommand,
+    SwapColdkeyCommand,
 )
 from .delegates import (
     NominateCommand,

bittensor/commands/register.py

@@ -611,3 +611,99 @@ def check_config(config: "bittensor.config"):
         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)
+
+
+class SwapColdkeyCommand:
+    """
+    Executes the `swap_coldkey` command to swap the coldkeys for a neuron on the network.
+
+    This command allows users to change the coldkey associated with their wallet.
+
+    Usage:
+        The command is used to swap the coldkey of a wallet for another coldkey.
+
+    Optional arguments:
+        - `--wallet.name` (str): Specifies the wallet for which the coldkey is to be swapped.
+        - `--wallet.coldkey` (str): The original coldkey name that is getting swapped out.
+        - `--wallet.coldkey_b` (str): The new coldkey name for which the old is getting swapped out.
+
+    Example usage:
+        btcli wallet swap_coldkey --wallet.name your_wallet_name --wallet.coldkey original_coldkey --wallet.coldkey_b new_coldkey
+    """
+
+    @staticmethod
+    def run(cli: "bittensor.cli"):
+        try:
+            subtensor: "bittensor.subtensor" = bittensor.subtensor(
+                config=cli.config, log_verbose=False
+            )
+            SwapColdkeyCommand._run(cli, subtensor)
+        finally:
+            if "subtensor" in locals():
+                subtensor.close()
+                bittensor.logging.debug("closing subtensor connection")
+
+    @staticmethod
+    def _run(cli: "bittensor.cli", subtensor: "bittensor.subtensor"):
+        """Swap your coldkey for all registered neurons on the network."""
+        wallet = bittensor.wallet(config=cli.config)
+
+        # Create a new wallet with the new coldkey
+        new_config = deepcopy(cli.config)
+        new_config.wallet.coldkey = new_config.wallet.coldkey_b
+        new_wallet = bittensor.wallet(config=new_config)
+
+        # Perform the coldkey swap
+        subtensor.swap_coldkey(
+            wallet=wallet,
+            new_wallet=new_wallet,
+            wait_for_finalization=False,
+            wait_for_inclusion=True,
+            prompt=False,
+        )
+
+    @staticmethod
+    def add_args(parser: argparse.ArgumentParser):
+        swap_coldkey_parser = parser.add_parser(
+            "swap_coldkey", help="""Swap your associated coldkey."""
+        )
+
+        swap_coldkey_parser.add_argument(
+            "--wallet.coldkey_b",
+            type=str,
+            default=defaults.wallet.coldkey,
+            help="""Name of the new coldkey""",
+            required=False,
+        )
+
+        bittensor.wallet.add_args(swap_coldkey_parser)
+        bittensor.subtensor.add_args(swap_coldkey_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,
+            )
+            _, endpoint = bittensor.subtensor.determine_chain_endpoint_and_network(
+                config.subtensor.network
+            )
+            config.subtensor.chain_endpoint = endpoint
+
+        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.coldkey") and not config.no_prompt:
+            coldkey = Prompt.ask("Enter old coldkey name", default=defaults.wallet.coldkey)
+            config.wallet.coldkey = str(coldkey)
+
+        if not config.is_set("wallet.coldkey_b") and not config.no_prompt:
+            coldkey = Prompt.ask("Enter new coldkey name", default=defaults.wallet.coldkey)
+            config.wallet.coldkey_b = str(coldkey)

bittensor/extrinsics/registration.py

@@ -523,3 +523,62 @@ def swap_hotkey_extrinsic(
                 f"Hotkey {wallet.hotkey} swapped for new hotkey: {new_wallet.hotkey}"
             )
             return True
+
+
+def swap_coldkey_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:
+    """
+    Executes a coldkey swap extrinsic on the Bittensor network.
+
+    This function swaps the coldkey associated with a wallet for a new coldkey. It can optionally prompt
+    for user confirmation before proceeding with the swap.
+
+    Args:
+        subtensor (bittensor.subtensor): The Subtensor object for blockchain interaction.
+        wallet (bittensor.wallet): The wallet containing the current coldkey to be swapped.
+        new_wallet (bittensor.wallet): The wallet containing the new coldkey to be set.
+        wait_for_inclusion (bool, optional): If True, waits for the transaction to be included in a block. 
+                                             Defaults to False.
+        wait_for_finalization (bool, optional): If True, waits for the transaction to be finalized. 
+                                                Defaults to True.
+        prompt (bool, optional): If True, prompts the user for confirmation before proceeding. 
+                                 Defaults to False.
+
+    Returns:
+        bool: True if the coldkey swap was successful, False otherwise.
+
+    Note:
+        This function unlocks the coldkey before proceeding with the swap operation.
+    """
+    wallet.coldkey  # unlock coldkey
+    if prompt:
+        # Prompt user for confirmation.
+        if not Confirm.ask(
+            f"Swap coldkey {wallet.coldkey.ss58_address} for new coldkey: {new_wallet.coldkey.ss58_address}?"
+        ):
+            return False
+
+    with bittensor.__console__.status(":satellite: Swapping coldkeys..."):
+        success, err_msg = subtensor._do_swap_coldkey(
+            wallet=wallet,
+            new_wallet=new_wallet,
+            wait_for_inclusion=wait_for_inclusion,
+            wait_for_finalization=wait_for_finalization,
+        )
+
+        if not success:
+            bittensor.__console__.print(f":cross_mark: [red]Failed[/red]: {err_msg}")
+            time.sleep(0.5)
+            return False
+
+        else:
+            bittensor.__console__.print(
+                f"Coldkey {wallet.coldkey.ss58_address} swapped for new coldkey: {new_wallet.coldkey.ss58_address}"
+            )
+            return True

bittensor/subtensor.py

@@ -1274,6 +1274,40 @@ def swap_hotkey(
             prompt=prompt,
         )
 
+    def swap_coldkey(
+        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 coldkey with a new coldkey for the specified wallet.
+
+        This method initiates an extrinsic to change the coldkey associated with a wallet to a new coldkey. It provides
+        options to wait for inclusion and finalization of the transaction, and to prompt the user for confirmation.
+
+        Args:
+            wallet (bittensor.wallet): The wallet whose coldkey is to be swapped.
+            new_wallet (bittensor.wallet): The new wallet with the coldkey to be set.
+            wait_for_inclusion (bool): Whether to wait for the transaction to be included in a block.
+                Default is `False`.
+            wait_for_finalization (bool): Whether to wait for the transaction to be finalized. Default is `True`.
+            prompt (bool): Whether to prompt the user for confirmation before proceeding. Default is `False`.
+
+        Returns:
+            bool: True if the coldkey swap was successful, False otherwise.
+        """
+        return swap_coldkey_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",
@@ -1546,6 +1580,62 @@ def make_substrate_call_with_retry():
 
         return make_substrate_call_with_retry()
 
+    def _do_swap_coldkey(
+        self,
+        wallet: "bittensor.wallet",
+        new_wallet: "bittensor.wallet",
+        wait_for_inclusion: bool = False,
+        wait_for_finalization: bool = True,
+    ) -> Tuple[bool, Optional[str]]:
+        """
+        Performs a coldkey swap extrinsic call to the Subtensor chain.
+
+        Args:
+            wallet (bittensor.wallet): The wallet whose coldkey is to be swapped.
+            new_wallet (bittensor.wallet): The wallet with the new coldkey to be set.
+            wait_for_inclusion (bool): Whether to wait for the transaction to be included in a block. Default is
+            `False`.
+            wait_for_finalization (bool): Whether to wait for the transaction to be finalized. Default is `True`.
+
+        Returns:
+            Tuple[bool, Optional[str]]: A tuple containing a boolean indicating success or failure, and an optional
+                error message.
+        """
+
+        @retry(delay=1, tries=3, backoff=2, max_delay=4, logger=_logger)
+        def make_substrate_call_with_retry():
+            # create extrinsic call
+            call = self.substrate.compose_call(
+                call_module="SubtensorModule",
+                call_function="swap_coldkey",
+                call_params={
+                    "coldkey": wallet.coldkey.ss58_address,
+                    "new_coldkey": new_wallet.coldkey.ss58_address,
+                },
+            )
+            extrinsic = self.substrate.create_signed_extrinsic(
+                call=call, keypair=wallet.coldkey
+            )
+            response = self.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, None
+
+            # process if swap successful, try again if still valid
+            response.process_events()
+            if not response.is_success:
+                return False, format_error_message(response.error_message)
+            # Successful swap
+            else:
+                return True, None
+
+        return make_substrate_call_with_retry()
+
     ############
     # Transfer #
     ############