Merge pull request #3062 from opentensor/feat/roman/autostaking

set auto hotkey per subnet

Author: basfroman

bittensor/core/async_subtensor.py

@@ -73,6 +73,7 @@
 from bittensor.core.extrinsics.asyncex.staking import (
     add_stake_extrinsic,
     add_stake_multiple_extrinsic,
+    set_auto_stake_extrinsic,
 )
 from bittensor.core.extrinsics.asyncex.start_call import start_call_extrinsic
 from bittensor.core.extrinsics.asyncex.mechanism import (
@@ -1201,6 +1202,42 @@ async def get_all_subnets_info(
 
         return SubnetInfo.list_from_dicts(result)
 
+    async def get_auto_stakes(
+        self,
+        coldkey_ss58: str,
+        block: Optional[int] = None,
+        block_hash: Optional[str] = None,
+        reuse_block: bool = False,
+    ) -> dict[int, str]:
+        """Fetches auto stake destinations for a given wallet across all subnets.
+
+        Parameters:
+            coldkey_ss58: Coldkey ss58 address.
+            block: The block number for the query.
+            block_hash: The block hash for the query.
+            reuse_block: Whether to reuse the last-used block hash.
+
+        Returns:
+            dict[int, str]:
+                - netuid: The unique identifier of the subnet.
+                - hotkey: The hotkey of the wallet.
+        """
+        block_hash = await self.determine_block_hash(block, block_hash, reuse_block)
+        query = await self.substrate.query_map(
+            module="SubtensorModule",
+            storage_function="AutoStakeDestination",
+            params=[coldkey_ss58],
+            block_hash=block_hash,
+        )
+
+        pairs = {}
+        async for netuid, destination in query:
+            hotkey_ss58 = decode_account_id(destination.value[0])
+            if hotkey_ss58:
+                pairs[int(netuid)] = hotkey_ss58
+
+        return pairs
+
     async def get_balance(
         self,
         address: str,
@@ -5170,6 +5207,46 @@ async def root_set_weights(
             period=period,
         )
 
+    async def set_auto_stake(
+        self,
+        wallet: "Wallet",
+        netuid: int,
+        hotkey_ss58: str,
+        period: Optional[int] = None,
+        raise_error: bool = False,
+        wait_for_inclusion: bool = True,
+        wait_for_finalization: bool = True,
+    ) -> tuple[bool, str]:
+        """Sets the coldkey to automatically stake to the hotkey within specific subnet mechanism.
+
+        Parameters:
+            wallet: Bittensor Wallet instance.
+            netuid: The subnet unique identifier.
+            hotkey_ss58: The SS58 address of the validator's hotkey to which the miner automatically stakes all rewards
+                received from the specified subnet immediately upon receipt.
+            period: The number of blocks during which the transaction will remain valid after it's submitted. If the
+                transaction is not included in a block within that number of blocks, it will expire and be rejected. You
+                can think of it as an expiration date for the transaction.
+            raise_error: Raises a relevant exception rather than returning `False` if unsuccessful.
+            wait_for_inclusion: Whether to wait for the inclusion of the transaction.
+            wait_for_finalization: Whether to wait for the finalization of the transaction.
+
+        Returns:
+            tuple[bool, str]:
+                `True` if the extrinsic executed successfully, `False` otherwise.
+                `message` is a string value describing the success or potential error.
+        """
+        return await set_auto_stake_extrinsic(
+            subtensor=self,
+            wallet=wallet,
+            netuid=netuid,
+            hotkey_ss58=hotkey_ss58,
+            period=period,
+            raise_error=raise_error,
+            wait_for_inclusion=wait_for_inclusion,
+            wait_for_finalization=wait_for_finalization,
+        )
+
     async def set_children(
         self,
         wallet: "Wallet",

bittensor/core/extrinsics/asyncex/staking.py

@@ -426,3 +426,71 @@ async def add_stake_multiple_extrinsic(
         return True
 
     return False
+
+
+async def set_auto_stake_extrinsic(
+    subtensor: "AsyncSubtensor",
+    wallet: "Wallet",
+    netuid: int,
+    hotkey_ss58: str,
+    period: Optional[int] = None,
+    raise_error: bool = False,
+    wait_for_inclusion: bool = True,
+    wait_for_finalization: bool = True,
+) -> tuple[bool, str]:
+    """Sets the coldkey to automatically stake to the hotkey within specific subnet mechanism.
+
+    Parameters:
+        subtensor: AsyncSubtensor instance.
+        wallet: Bittensor Wallet instance.
+        netuid: The subnet unique identifier.
+        hotkey_ss58: The SS58 address of the validator's hotkey to which the miner automatically stakes all rewards
+            received from the specified subnet immediately upon receipt.
+        period: The number of blocks during which the transaction will remain valid after it's submitted. If the
+            transaction is not included in a block within that number of blocks, it will expire and be rejected. You can
+            think of it as an expiration date for the transaction.
+        raise_error: Raises a relevant exception rather than returning `False` if unsuccessful.
+        wait_for_inclusion: Whether to wait for the inclusion of the transaction.
+        wait_for_finalization: Whether to wait for the finalization of the transaction.
+
+    Returns:
+        tuple[bool, str]:
+            `True` if the extrinsic executed successfully, `False` otherwise.
+            `message` is a string value describing the success or potential error.
+    """
+    try:
+        unlock = unlock_key(wallet, raise_error=raise_error)
+        if not unlock.success:
+            logging.error(unlock.message)
+            return False, unlock.message
+
+        call = await subtensor.substrate.compose_call(
+            call_module="SubtensorModule",
+            call_function="set_coldkey_auto_stake_hotkey",
+            call_params={
+                "netuid": netuid,
+                "hotkey": hotkey_ss58,
+            },
+        )
+        success, message = await subtensor.sign_and_send_extrinsic(
+            call=call,
+            wallet=wallet,
+            period=period,
+            raise_error=raise_error,
+            wait_for_inclusion=wait_for_inclusion,
+            wait_for_finalization=wait_for_finalization,
+        )
+
+        if success:
+            logging.debug(message)
+            return True, message
+
+        logging.error(message)
+        return False, message
+
+    except Exception as error:
+        if raise_error:
+            raise error
+        logging.error(str(error))
+
+        return False, str(error)

bittensor/core/extrinsics/staking.py

@@ -405,3 +405,71 @@ def add_stake_multiple_extrinsic(
         return True
 
     return False
+
+
+def set_auto_stake_extrinsic(
+    subtensor: "Subtensor",
+    wallet: "Wallet",
+    netuid: int,
+    hotkey_ss58: str,
+    period: Optional[int] = None,
+    raise_error: bool = False,
+    wait_for_inclusion: bool = True,
+    wait_for_finalization: bool = True,
+) -> tuple[bool, str]:
+    """Sets the coldkey to automatically stake to the hotkey within specific subnet mechanism.
+
+    Parameters:
+        subtensor: AsyncSubtensor instance.
+        wallet: Bittensor Wallet instance.
+        netuid: The subnet unique identifier.
+        hotkey_ss58: The SS58 address of the validator's hotkey to which the miner automatically stakes all rewards
+            received from the specified subnet immediately upon receipt.
+        period: The number of blocks during which the transaction will remain valid after it's submitted. If the
+            transaction is not included in a block within that number of blocks, it will expire and be rejected. You can
+            think of it as an expiration date for the transaction.
+        raise_error: Raises a relevant exception rather than returning `False` if unsuccessful.
+        wait_for_inclusion: Whether to wait for the inclusion of the transaction.
+        wait_for_finalization: Whether to wait for the finalization of the transaction.
+
+    Returns:
+        tuple[bool, str]:
+            `True` if the extrinsic executed successfully, `False` otherwise.
+            `message` is a string value describing the success or potential error.
+    """
+    try:
+        unlock = unlock_key(wallet, raise_error=raise_error)
+        if not unlock.success:
+            logging.error(unlock.message)
+            return False, unlock.message
+
+        call = subtensor.substrate.compose_call(
+            call_module="SubtensorModule",
+            call_function="set_coldkey_auto_stake_hotkey",
+            call_params={
+                "netuid": netuid,
+                "hotkey": hotkey_ss58,
+            },
+        )
+        success, message = subtensor.sign_and_send_extrinsic(
+            call=call,
+            wallet=wallet,
+            period=period,
+            raise_error=raise_error,
+            wait_for_inclusion=wait_for_inclusion,
+            wait_for_finalization=wait_for_finalization,
+        )
+
+        if success:
+            logging.debug(message)
+            return True, message
+
+        logging.error(message)
+        return False, message
+
+    except Exception as error:
+        if raise_error:
+            raise error
+        logging.error(str(error))
+
+        return False, str(error)

bittensor/core/subtensor.py

@@ -77,6 +77,7 @@
 from bittensor.core.extrinsics.staking import (
     add_stake_extrinsic,
     add_stake_multiple_extrinsic,
+    set_auto_stake_extrinsic,
 )
 from bittensor.core.extrinsics.start_call import start_call_extrinsic
 from bittensor.core.extrinsics.take import (
@@ -687,6 +688,38 @@ def get_all_subnets_info(self, block: Optional[int] = None) -> list["SubnetInfo"
 
         return SubnetInfo.list_from_dicts(result)
 
+    def get_auto_stakes(
+        self,
+        coldkey_ss58: str,
+        block: Optional[int] = None,
+    ) -> dict[int, str]:
+        """Fetches auto stake destinations for a given wallet across all subnets.
+
+        Parameters:
+            coldkey_ss58: Coldkey ss58 address.
+            block: Subnet unique identifier.
+
+        Returns:
+            dict[int, str]:
+                - netuid: The unique identifier of the subnet.
+                - hotkey: The hotkey of the wallet.
+        """
+        block_hash = self.determine_block_hash(block=block)
+        query = self.substrate.query_map(
+            module="SubtensorModule",
+            storage_function="AutoStakeDestination",
+            params=[coldkey_ss58],
+            block_hash=block_hash,
+        )
+
+        pairs = {}
+        for netuid, destination in query:
+            hotkey_ss58 = decode_account_id(destination.value[0])
+            if hotkey_ss58:
+                pairs[int(netuid)] = hotkey_ss58
+
+        return pairs
+
     def get_balance(self, address: str, block: Optional[int] = None) -> Balance:
         """
         Retrieves the balance for given coldkey. Always in TAO.
@@ -4005,6 +4038,46 @@ def root_set_weights(
             period=period,
         )
 
+    def set_auto_stake(
+        self,
+        wallet: "Wallet",
+        netuid: int,
+        hotkey_ss58: str,
+        period: Optional[int] = None,
+        raise_error: bool = False,
+        wait_for_inclusion: bool = True,
+        wait_for_finalization: bool = True,
+    ) -> tuple[bool, str]:
+        """Sets the coldkey to automatically stake to the hotkey within specific subnet mechanism.
+
+        Parameters:
+            wallet: Bittensor Wallet instance.
+            netuid: The subnet unique identifier.
+            hotkey_ss58: The SS58 address of the validator's hotkey to which the miner automatically stakes all rewards
+                received from the specified subnet immediately upon receipt.
+            period: The number of blocks during which the transaction will remain valid after it's submitted. If the
+                transaction is not included in a block within that number of blocks, it will expire and be rejected. You
+                can think of it as an expiration date for the transaction.
+            raise_error: Raises a relevant exception rather than returning `False` if unsuccessful.
+            wait_for_inclusion: Whether to wait for the inclusion of the transaction.
+            wait_for_finalization: Whether to wait for the finalization of the transaction.
+
+        Returns:
+            tuple[bool, str]:
+                `True` if the extrinsic executed successfully, `False` otherwise.
+                `message` is a string value describing the success or potential error.
+        """
+        return set_auto_stake_extrinsic(
+            subtensor=self,
+            wallet=wallet,
+            netuid=netuid,
+            hotkey_ss58=hotkey_ss58,
+            period=period,
+            raise_error=raise_error,
+            wait_for_inclusion=wait_for_inclusion,
+            wait_for_finalization=wait_for_finalization,
+        )
+
     def set_children(
         self,
         wallet: "Wallet",

bittensor/core/subtensor_api/staking.py

@@ -9,6 +9,7 @@ class Staking:
     def __init__(self, subtensor: Union["_Subtensor", "_AsyncSubtensor"]):
         self.add_stake = subtensor.add_stake
         self.add_stake_multiple = subtensor.add_stake_multiple
+        self.get_auto_stakes = subtensor.get_auto_stakes
         self.get_hotkey_stake = subtensor.get_hotkey_stake
         self.get_minimum_required_stake = subtensor.get_minimum_required_stake
         self.get_stake = subtensor.get_stake
@@ -23,6 +24,7 @@ def __init__(self, subtensor: Union["_Subtensor", "_AsyncSubtensor"]):
         self.get_stake_weight = subtensor.get_stake_weight
         self.get_unstake_fee = subtensor.get_unstake_fee
         self.move_stake = subtensor.move_stake
+        self.set_auto_stake = subtensor.set_auto_stake
         self.unstake = subtensor.unstake
         self.unstake_all = subtensor.unstake_all
         self.unstake_multiple = subtensor.unstake_multiple

bittensor/core/subtensor_api/utils.py

@@ -35,6 +35,7 @@ def add_legacy_methods(subtensor: "SubtensorApi"):
         subtensor._subtensor.get_all_revealed_commitments
     )
     subtensor.get_all_subnets_info = subtensor._subtensor.get_all_subnets_info
+    subtensor.get_auto_stakes = subtensor._subtensor.get_auto_stakes
     subtensor.get_balance = subtensor._subtensor.get_balance
     subtensor.get_balances = subtensor._subtensor.get_balances
     subtensor.get_block_hash = subtensor._subtensor.get_block_hash
@@ -161,6 +162,7 @@ def add_legacy_methods(subtensor: "SubtensorApi"):
     )
     subtensor.root_set_weights = subtensor._subtensor.root_set_weights
     subtensor.serve_axon = subtensor._subtensor.serve_axon
+    subtensor.set_auto_stake = subtensor._subtensor.set_auto_stake
     subtensor.set_children = subtensor._subtensor.set_children
     subtensor.set_commitment = subtensor._subtensor.set_commitment
     subtensor.set_delegate_take = subtensor._subtensor.set_delegate_take