Merge pull request #505 from opentensor/feat/swap-hotkey-netuids-2

Feat: Adds netuid support in swap_hotkeys

Author: ibraheem-abe

bittensor_cli/cli.py

@@ -1897,6 +1897,8 @@ def wallet_swap_hotkey(
         wallet_name: Optional[str] = Options.wallet_name,
         wallet_path: Optional[str] = Options.wallet_path,
         wallet_hotkey: Optional[str] = Options.wallet_hotkey,
+        netuid: Optional[int] = Options.netuid_not_req,
+        all_netuids: bool = Options.all_netuids,
         network: Optional[list[str]] = Options.network,
         destination_hotkey_name: Optional[str] = typer.Argument(
             None, help="Destination hotkey name."
@@ -1917,12 +1919,14 @@ def wallet_swap_hotkey(
 
         - Make sure that your original key pair (coldkeyA, hotkeyA) is already registered.
         - Make sure that you use a newly created hotkeyB in this command. A hotkeyB that is already registered cannot be used in this command.
+        - You can specify the netuid for which you want to swap the hotkey for. If it is not defined, the swap will be initiated for all subnets.
         - Finally, note that this command requires a fee of 1 TAO for recycling and this fee is taken from your wallet (coldkeyA).
 
         EXAMPLE
 
-        [green]$[/green] btcli wallet swap_hotkey destination_hotkey_name --wallet-name your_wallet_name --wallet-hotkey original_hotkey
+        [green]$[/green] btcli wallet swap_hotkey destination_hotkey_name --wallet-name your_wallet_name --wallet-hotkey original_hotkey --netuid 1
         """
+        netuid = get_optional_netuid(netuid, all_netuids)
         self.verbosity_handler(quiet, verbose, json_output)
         original_wallet = self.wallet_ask(
             wallet_name,
@@ -1946,7 +1950,7 @@ def wallet_swap_hotkey(
         self.initialize_chain(network)
         return self._run_command(
             wallets.swap_hotkey(
-                original_wallet, new_wallet, self.subtensor, prompt, json_output
+                original_wallet, new_wallet, self.subtensor, netuid, prompt, json_output
             )
         )
 

bittensor_cli/src/bittensor/extrinsics/registration.py

@@ -1611,7 +1611,8 @@ def _update_curr_block(
     """
     Update the current block data with the provided block information and difficulty.
 
-    This function updates the current block and its difficulty in a thread-safe manner. It sets the current block
+    This function updates the current block
+    and its difficulty in a thread-safe manner. It sets the current block
     number, hashes the block with the hotkey, updates the current block bytes, and packs the difficulty.
 
     :param curr_diff: Shared array to store the current difficulty.
@@ -1745,6 +1746,7 @@ async def swap_hotkey_extrinsic(
     subtensor: "SubtensorInterface",
     wallet: Wallet,
     new_wallet: Wallet,
+    netuid: Optional[int] = None,
     prompt: bool = False,
 ) -> bool:
     """
@@ -1756,37 +1758,76 @@ async def swap_hotkey_extrinsic(
     netuids_registered = await subtensor.get_netuids_for_hotkey(
         wallet.hotkey.ss58_address, block_hash=block_hash
     )
-    if not len(netuids_registered) > 0:
+    netuids_registered_new_hotkey = await subtensor.get_netuids_for_hotkey(
+        new_wallet.hotkey.ss58_address, block_hash=block_hash
+    )
+
+    if netuid is not None and netuid not in netuids_registered:
+        err_console.print(
+            f":cross_mark: [red]Failed[/red]: Original hotkey {wallet.hotkey.ss58_address} is not registered on subnet {netuid}"
+        )
+        return False
+
+    elif not len(netuids_registered) > 0:
         err_console.print(
-            f"Destination hotkey [dark_orange]{new_wallet.hotkey.ss58_address}[/dark_orange] is not registered. "
+            f"Original hotkey [dark_orange]{wallet.hotkey.ss58_address}[/dark_orange] is not registered on any subnet. "
             f"Please register and try again"
         )
         return False
 
+    if netuid is not None:
+        if netuid in netuids_registered_new_hotkey:
+            err_console.print(
+                f":cross_mark: [red]Failed[/red]: New hotkey {new_wallet.hotkey.ss58_address} "
+                f"is already registered on subnet {netuid}"
+            )
+            return False
+    else:
+        if len(netuids_registered_new_hotkey) > 0:
+            err_console.print(
+                f":cross_mark: [red]Failed[/red]: New hotkey {new_wallet.hotkey.ss58_address} "
+                f"is already registered on subnet(s) {netuids_registered_new_hotkey}"
+            )
+            return False
+
     if not unlock_key(wallet).success:
         return False
 
     if prompt:
         # Prompt user for confirmation.
-        if not Confirm.ask(
-            f"Do you want to swap [dark_orange]{wallet.name}[/dark_orange] hotkey \n\t"
-            f"[dark_orange]{wallet.hotkey.ss58_address}[/dark_orange] with hotkey \n\t"
-            f"[dark_orange]{new_wallet.hotkey.ss58_address}[/dark_orange]\n"
-            "This operation will cost [bold cyan]1 TAO t (recycled)[/bold cyan]"
-        ):
+        if netuid is not None:
+            confirm_message = (
+                f"Do you want to swap [dark_orange]{wallet.name}[/dark_orange] hotkey \n\t"
+                f"[dark_orange]{wallet.hotkey.ss58_address}[/dark_orange] with hotkey \n\t"
+                f"[dark_orange]{new_wallet.hotkey.ss58_address}[/dark_orange] on subnet {netuid}\n"
+                "This operation will cost [bold cyan]1 TAO (recycled)[/bold cyan]"
+            )
+        else:
+            confirm_message = (
+                f"Do you want to swap [dark_orange]{wallet.name}[/dark_orange] hotkey \n\t"
+                f"[dark_orange]{wallet.hotkey.ss58_address}[/dark_orange] with hotkey \n\t"
+                f"[dark_orange]{new_wallet.hotkey.ss58_address}[/dark_orange] on all subnets\n"
+                "This operation will cost [bold cyan]1 TAO (recycled)[/bold cyan]"
+            )
+
+        if not Confirm.ask(confirm_message):
             return False
     print_verbose(
         f"Swapping {wallet.name}'s hotkey ({wallet.hotkey.ss58_address}) with "
-        f"{new_wallet.name}s hotkey ({new_wallet.hotkey.ss58_address})"
+        f"{new_wallet.name}'s hotkey ({new_wallet.hotkey.ss58_address})"
     )
     with console.status(":satellite: Swapping hotkeys...", spinner="aesthetic"):
+        call_params = {
+            "hotkey": wallet.hotkey.ss58_address,
+            "new_hotkey": new_wallet.hotkey.ss58_address,
+        }
+        if netuid is not None:
+            call_params["netuid"] = netuid
+
         call = await subtensor.substrate.compose_call(
             call_module="SubtensorModule",
             call_function="swap_hotkey",
-            call_params={
-                "hotkey": wallet.hotkey.ss58_address,
-                "new_hotkey": new_wallet.hotkey.ss58_address,
-            },
+            call_params=call_params,
         )
         success, err_msg = await subtensor.sign_and_send_extrinsic(call, wallet)
 

bittensor_cli/src/commands/wallets.py

@@ -1632,6 +1632,7 @@ async def swap_hotkey(
     original_wallet: Wallet,
     new_wallet: Wallet,
     subtensor: SubtensorInterface,
+    netuid: Optional[int],
     prompt: bool,
     json_output: bool,
 ):
@@ -1640,6 +1641,7 @@ async def swap_hotkey(
         subtensor,
         original_wallet,
         new_wallet,
+        netuid=netuid,
         prompt=prompt,
     )
     if json_output: