Re-adds check_coldkey_swap command (#2126)

* Inital commit: Readds check-coldkey-swap

* Adds check_coldkey_swap command

bittensor/__init__.py

@@ -229,6 +229,37 @@ def debug(on: bool = True):
         "SubnetRegistrationRuntimeApi": {
             "methods": {"get_network_registration_cost": {"params": [], "type": "u64"}}
         },
+        "ColdkeySwapRuntimeApi": {
+            "methods": {
+                "get_scheduled_coldkey_swap": {
+                    "params": [
+                        {
+                            "name": "coldkey_account_vec",
+                            "type": "Vec<u8>",
+                        },
+                    ],
+                    "type": "Vec<u8>",
+                },
+                "get_remaining_arbitration_period": {
+                    "params": [
+                        {
+                            "name": "coldkey_account_vec",
+                            "type": "Vec<u8>",
+                        },
+                    ],
+                    "type": "Vec<u8>",
+                },
+                "get_coldkey_swap_destinations": {
+                    "params": [
+                        {
+                            "name": "coldkey_account_vec",
+                            "type": "Vec<u8>",
+                        },
+                    ],
+                    "type": "Vec<u8>",
+                },
+            }
+        },
     },
 }
 

bittensor/chain_data.py

@@ -191,6 +191,14 @@
                 ["liquid_alpha_enabled", "bool"],
             ],
         },
+        "ScheduledColdkeySwapInfo": {
+            "type": "struct",
+            "type_mapping": [
+                ["old_coldkey", "AccountId"],
+                ["new_coldkey", "AccountId"],
+                ["arbitration_block", "Compact<u64>"],
+            ],
+        },
     }
 }
 
@@ -324,6 +332,7 @@ class ChainDataType(Enum):
     StakeInfo = 6
     IPInfo = 7
     SubnetHyperparameters = 8
+    ScheduledColdkeySwapInfo = 9
 
 
 def from_scale_encoding(
@@ -1140,3 +1149,54 @@ class ProposalVoteData(TypedDict):
 
 
 ProposalCallData = GenericCall
+
+
+@dataclass
+class ScheduledColdkeySwapInfo:
+    """Dataclass for scheduled coldkey swap information."""
+
+    old_coldkey: str
+    new_coldkey: str
+    arbitration_block: int
+
+    @classmethod
+    def fix_decoded_values(cls, decoded: Any) -> "ScheduledColdkeySwapInfo":
+        """Fixes the decoded values."""
+        return cls(
+            old_coldkey=ss58_encode(decoded["old_coldkey"], bittensor.__ss58_format__),
+            new_coldkey=ss58_encode(decoded["new_coldkey"], bittensor.__ss58_format__),
+            arbitration_block=decoded["arbitration_block"],
+        )
+
+    @classmethod
+    def from_vec_u8(cls, vec_u8: List[int]) -> Optional["ScheduledColdkeySwapInfo"]:
+        """Returns a ScheduledColdkeySwapInfo object from a ``vec_u8``."""
+        if len(vec_u8) == 0:
+            return None
+
+        decoded = from_scale_encoding(vec_u8, ChainDataType.ScheduledColdkeySwapInfo)
+        if decoded is None:
+            return None
+
+        return ScheduledColdkeySwapInfo.fix_decoded_values(decoded)
+
+    @classmethod
+    def list_from_vec_u8(cls, vec_u8: List[int]) -> List["ScheduledColdkeySwapInfo"]:
+        """Returns a list of ScheduledColdkeySwapInfo objects from a ``vec_u8``."""
+        decoded = from_scale_encoding(
+            vec_u8, ChainDataType.ScheduledColdkeySwapInfo, is_vec=True
+        )
+        if decoded is None:
+            return []
+
+        return [ScheduledColdkeySwapInfo.fix_decoded_values(d) for d in decoded]
+
+    @classmethod
+    def decode_account_id_list(cls, vec_u8: List[int]) -> Optional[List[str]]:
+        """Decodes a list of AccountIds from vec_u8."""
+        decoded = from_scale_encoding(vec_u8, ChainDataType.AccountId, is_vec=True)
+        if decoded is None:
+            return None
+        return [
+            ss58_encode(account_id, bittensor.__ss58_format__) for account_id in decoded
+        ]

bittensor/cli.py

@@ -69,6 +69,7 @@
     WalletCreateCommand,
     CommitWeightCommand,
     RevealWeightCommand,
+    CheckColdKeySwapCommand
 )
 
 # Create a console instance for CLI display.
@@ -157,6 +158,7 @@
             "set_identity": SetIdentityCommand,
             "get_identity": GetIdentityCommand,
             "history": GetWalletHistoryCommand,
+            "check_coldkey_swap": CheckColdKeySwapCommand,
         },
     },
     "stake": {

bittensor/commands/__init__.py

@@ -121,3 +121,4 @@
     RootSetSlashCommand,
 )
 from .identity import GetIdentityCommand, SetIdentityCommand
+from .check_coldkey_swap import CheckColdKeySwapCommand

bittensor/commands/check_coldkey_swap.py

@@ -0,0 +1,126 @@
+# The MIT License (MIT)
+# Copyright © 2021 Yuma Rao
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
+# documentation files (the “Software”), to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software,
+# and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all copies or substantial portions of
+# the Software.
+
+# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
+# THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+import argparse
+
+from rich.prompt import Prompt
+
+import bittensor
+from bittensor.utils.formatting import convert_blocks_to_time
+from . import defaults
+
+console = bittensor.__console__
+
+
+def fetch_arbitration_stats(subtensor, wallet):
+    """
+    Performs a check of the current arbitration data (if any), and displays it through the bittensor console.
+    """
+    arbitration_check = len(subtensor.check_in_arbitration(wallet.coldkey.ss58_address))
+    if arbitration_check == 0:
+        bittensor.__console__.print(
+            "[green]There has been no previous key swap initiated for your coldkey.[/green]"
+        )
+    if arbitration_check == 1:
+        arbitration_remaining = subtensor.get_remaining_arbitration_period(
+            wallet.coldkey.ss58_address
+        )
+        hours, minutes, seconds = convert_blocks_to_time(arbitration_remaining)
+        bittensor.__console__.print(
+            "[yellow]There has been 1 swap request made for this coldkey already."
+            " By adding another swap request, the key will enter arbitration."
+            f" Your key swap is scheduled for {hours} hours, {minutes} minutes, {seconds} seconds"
+            " from now.[/yellow]"
+        )
+    if arbitration_check > 1:
+        bittensor.__console__.print(
+            f"[red]This coldkey is currently in arbitration with a total swaps of {arbitration_check}.[/red]"
+        )
+
+
+class CheckColdKeySwapCommand:
+    """
+    Executes the ``check_coldkey_swap`` command to check swap status of a coldkey in the Bittensor network.
+    Usage:
+        Users need to specify the wallet they want to check the swap status of.
+    Example usage::
+        btcli wallet check_coldkey_swap
+    Note:
+        This command is important for users who wish check if swap requests were made against their coldkey.
+    """
+
+    @staticmethod
+    def run(cli: "bittensor.cli"):
+        """
+        Runs the check coldkey swap command.
+        Args:
+            cli (bittensor.cli): The CLI object containing configuration and command-line interface utilities.
+        """
+        try:
+            config = cli.config.copy()
+            subtensor: "bittensor.subtensor" = bittensor.subtensor(
+                config=config, log_verbose=False
+            )
+            CheckColdKeySwapCommand._run(cli, subtensor)
+        except Exception as e:
+            bittensor.logging.warning(f"Failed to get swap status: {e}")
+        finally:
+            if "subtensor" in locals():
+                subtensor.close()
+                bittensor.logging.debug("closing subtensor connection")
+
+    @staticmethod
+    def _run(cli: "bittensor.cli", subtensor: "bittensor.subtensor"):
+        """
+        Internal method to check coldkey swap status.
+        Args:
+            cli (bittensor.cli): The CLI object containing configuration and command-line interface utilities.
+            subtensor (bittensor.subtensor): The subtensor object for blockchain interactions.
+        """
+        config = cli.config.copy()
+        wallet = bittensor.wallet(config=config)
+
+        fetch_arbitration_stats(subtensor, wallet)
+
+    @classmethod
+    def check_config(cls, config: "bittensor.config"):
+        """
+        Checks and prompts for necessary configuration settings.
+        Args:
+            config (bittensor.config): The configuration object.
+        Prompts the user for wallet name if not set in the config.
+        """
+        if not config.is_set("wallet.name") and not config.no_prompt:
+            wallet_name: str = Prompt.ask(
+                "Enter wallet name", default=defaults.wallet.name
+            )
+            config.wallet.name = str(wallet_name)
+
+    @staticmethod
+    def add_args(command_parser: argparse.ArgumentParser):
+        """
+        Adds arguments to the command parser.
+        Args:
+            command_parser (argparse.ArgumentParser): The command parser to add arguments to.
+        """
+        swap_parser = command_parser.add_parser(
+            "check_coldkey_swap",
+            help="""Check the status of swap requests for a coldkey on the Bittensor network.
+            Adding more than one swap request will make the key go into arbitration mode.""",
+        )
+        bittensor.wallet.add_args(swap_parser)
+        bittensor.subtensor.add_args(swap_parser)

bittensor/subtensor.py

@@ -2297,6 +2297,46 @@ def make_substrate_call_with_retry():
 
         return make_substrate_call_with_retry()
 
+    ##################
+    # Coldkey Swap   #
+    ##################
+
+    def check_in_arbitration(self, ss58_address: str) -> int:
+        """
+        Checks storage function to see if the provided coldkey is in arbitration.
+        If 0, `swap` has not been called on this key. If 1, swap has been called once, so
+        the key is not in arbitration. If >1, `swap` has been called with multiple destinations, and
+        the key is thus in arbitration.
+        """
+        return self.query_module(
+            "SubtensorModule", "ColdkeySwapDestinations", params=[ss58_address]
+        ).decode()
+
+    def get_remaining_arbitration_period(
+        self, coldkey_ss58: str, block: Optional[int] = None
+    ) -> Optional[int]:
+        """
+        Retrieves the remaining arbitration period for a given coldkey.
+        Args:
+            coldkey_ss58 (str): The SS58 address of the coldkey.
+            block (Optional[int], optional): The block number to query. If None, uses the latest block.
+        Returns:
+            Optional[int]: The remaining arbitration period in blocks, or 0 if not found.
+        """
+        arbitration_block = self.query_subtensor(
+            name="ColdkeyArbitrationBlock",
+            block=block,
+            params=[coldkey_ss58],
+        )
+
+        if block is None:
+            block = self.block
+
+        if arbitration_block.value > block:
+            return arbitration_block.value - block
+        else:
+            return 0
+
     ##########
     # Senate #
     ##########

bittensor/utils/formatting.py

@@ -20,3 +20,17 @@ def millify(n: int):
     )
 
     return "{:.2f}{}".format(n / 10 ** (3 * millidx), millnames[millidx])
+
+
+def convert_blocks_to_time(blocks: int, block_time: int = 12) -> tuple[int, int, int]:
+    """
+    Converts number of blocks into number of hours, minutes, seconds.
+    :param blocks: number of blocks
+    :param block_time: time per block, by default this is 12
+    :return: tuple containing number of hours, number of minutes, number of seconds
+    """
+    seconds = blocks * block_time
+    hours = seconds // 3600
+    minutes = (seconds % 3600) // 60
+    remaining_seconds = seconds % 60
+    return hours, minutes, remaining_seconds

tests/unit_tests/test_subtensor.py

@@ -2277,3 +2277,41 @@ def test_get_delegate_take_no_data(mocker, subtensor):
     subtensor.query_subtensor.assert_called_once_with("Delegates", block, [hotkey_ss58])
     spy_u16_normalized_float.assert_not_called()
     assert result is None
+
+
+def test_get_remaining_arbitration_period(subtensor, mocker):
+    """Tests successful retrieval of total stake for hotkey."""
+    # Prep
+    subtensor.query_subtensor = mocker.MagicMock(return_value=mocker.MagicMock(value=0))
+    fake_ss58_address = "12bzRJfh7arnnfPPUZHeJUaE62QLEwhK48QnH9LXeK2m1iZU"
+
+    # Call
+    result = subtensor.get_remaining_arbitration_period(coldkey_ss58=fake_ss58_address)
+
+    # Assertions
+    subtensor.query_subtensor.assert_called_once_with(
+        name="ColdkeyArbitrationBlock", block=None, params=[fake_ss58_address]
+    )
+    # if we change the methods logic in the future we have to be make sure the returned type is correct
+    assert result == 0
+
+
+def test_get_remaining_arbitration_period_happy(subtensor, mocker):
+    """Tests successful retrieval of total stake for hotkey."""
+    # Prep
+    subtensor.query_subtensor = mocker.MagicMock(
+        return_value=mocker.MagicMock(value=2000)
+    )
+    fake_ss58_address = "12bzRJfh7arnnfPPUZHeJUaE62QLEwhK48QnH9LXeK2m1iZU"
+
+    # Call
+    result = subtensor.get_remaining_arbitration_period(
+        coldkey_ss58=fake_ss58_address, block=200
+    )
+
+    # Assertions
+    subtensor.query_subtensor.assert_called_once_with(
+        name="ColdkeyArbitrationBlock", block=200, params=[fake_ss58_address]
+    )
+    # if we change the methods logic in the future we have to be make sure the returned type is correct
+    assert result == 1800  # 2000 - 200