Improve move_stake extrinsic (add move_all_stake parameter)

bittensor/core/async_subtensor.py

@@ -4657,10 +4657,11 @@ async def move_stake(
         origin_netuid: int,
         destination_hotkey: str,
         destination_netuid: int,
-        amount: Balance,
+        amount: Optional[Balance] = None,
         wait_for_inclusion: bool = True,
         wait_for_finalization: bool = False,
         period: Optional[int] = None,
+        move_all_stake: bool = False,
     ) -> bool:
         """
         Moves stake to a different hotkey and/or subnet.
@@ -4677,6 +4678,7 @@ async def move_stake(
             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.
+            move_all_stake: If true, moves all stake from the source hotkey to the destination hotkey.
 
         Returns:
             success: True if the stake movement was successful.
@@ -4693,6 +4695,7 @@ async def move_stake(
             wait_for_inclusion=wait_for_inclusion,
             wait_for_finalization=wait_for_finalization,
             period=period,
+            move_all_stake=move_all_stake,
         )
 
     async def register(

bittensor/core/extrinsics/asyncex/move_stake.py

@@ -305,28 +305,34 @@ async def move_stake_extrinsic(
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
     period: Optional[int] = None,
+    move_all_stake: bool = False,
 ) -> bool:
     """
     Moves stake from one hotkey to another within subnets in the Bittensor network.
 
     Args:
-        subtensor (Subtensor): The subtensor instance to interact with the blockchain.
-        wallet (Wallet): The wallet containing the coldkey to authorize the move.
-        origin_hotkey (str): SS58 address of the origin hotkey associated with the stake.
-        origin_netuid (int): Network UID of the origin subnet.
-        destination_hotkey (str): SS58 address of the destination hotkey.
-        destination_netuid (int): Network UID of the destination subnet.
-        amount (Balance): The amount of stake to move as a `Balance` object.
-        wait_for_inclusion (bool): If True, waits for transaction inclusion in a block. Defaults to True.
-        wait_for_finalization (bool): If True, waits for transaction finalization. Defaults to False.
-        period (Optional[int]): 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.
+        subtensor: The subtensor instance to interact with the blockchain.
+        wallet: The wallet containing the coldkey to authorize the move.
+        origin_hotkey: SS58 address of the origin hotkey associated with the stake.
+        origin_netuid: Network UID of the origin subnet.
+        destination_hotkey: SS58 address of the destination hotkey.
+        destination_netuid: Network UID of the destination subnet.
+        amount: The amount of stake to move as a `Balance` object.
+        wait_for_inclusion: If True, waits for transaction inclusion in a block. Defaults to True.
+        wait_for_finalization: If True, waits for transaction finalization. Defaults to False.
+        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.
+        move_all_stake: If true, moves all stake from the source hotkey to the destination hotkey.
 
     Returns:
         bool: True if the move was successful, False otherwise.
     """
-    amount.set_unit(netuid=origin_netuid)
+    if not amount and not move_all_stake:
+        logging.error(
+            ":cross_mark: [red]Failed[/red]: Please specify an `amount` or `move_all_stake` argument to move stake."
+        )
+        return False
 
     # Check sufficient stake
     stake_in_origin, stake_in_destination = await _get_stake_in_origin_and_dest(
@@ -338,13 +344,18 @@ async def move_stake_extrinsic(
         origin_netuid=origin_netuid,
         destination_netuid=destination_netuid,
     )
-    if stake_in_origin < amount:
+    if move_all_stake:
+        amount = stake_in_origin
+
+    elif stake_in_origin < amount:
         logging.error(
             f":cross_mark: [red]Failed[/red]: Insufficient stake in origin hotkey: {origin_hotkey}. "
             f"Stake: {stake_in_origin}, amount: {amount}"
         )
         return False
 
+    amount.set_unit(netuid=origin_netuid)
+
     try:
         logging.info(
             f"Moving stake from hotkey [blue]{origin_hotkey}[/blue] to hotkey [blue]{destination_hotkey}[/blue]\n"

bittensor/core/extrinsics/move_stake.py

@@ -301,29 +301,34 @@ def move_stake_extrinsic(
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
     period: Optional[int] = None,
+    move_all_stake: bool = False,
 ) -> bool:
     """
     Moves stake to a different hotkey and/or subnet while keeping the same coldkey owner.
 
     Args:
-        subtensor (Subtensor): Subtensor instance.
-        wallet (bittensor.wallet): The wallet to move stake from.
-        origin_hotkey (str): The SS58 address of the source hotkey.
-        origin_netuid (int): The netuid of the source subnet.
-        destination_hotkey (str): The SS58 address of the destination hotkey.
-        destination_netuid (int): The netuid of the destination subnet.
-        amount (Union[Balance, float]): Amount to move.
-        wait_for_inclusion (bool): If true, waits for inclusion before returning.
-        wait_for_finalization (bool): If true, waits for finalization before returning.
-        period (Optional[int]): 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.
+        subtensor: Subtensor instance.
+        wallet: The wallet to move stake from.
+        origin_hotkey: The SS58 address of the source hotkey.
+        origin_netuid: The netuid of the source subnet.
+        destination_hotkey: The SS58 address of the destination hotkey.
+        destination_netuid: The netuid of the destination subnet.
+        amount: Amount to move.
+        wait_for_inclusion: If true, waits for inclusion before returning.
+        wait_for_finalization: If true, waits for finalization before returning.
+        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.
+        move_all_stake: If true, moves all stake from the source hotkey to the destination hotkey.
 
     Returns:
-        success (bool): True if the move was successful.
+        success: True if the move was successful. Otherwise, False.
     """
-
-    amount.set_unit(netuid=origin_netuid)
+    if not amount and not move_all_stake:
+        logging.error(
+            ":cross_mark: [red]Failed[/red]: Please specify an `amount` or `move_all_stake` argument to move stake."
+        )
+        return False
 
     # Check sufficient stake
     stake_in_origin, stake_in_destination = _get_stake_in_origin_and_dest(
@@ -335,12 +340,18 @@ def move_stake_extrinsic(
         origin_coldkey_ss58=wallet.coldkeypub.ss58_address,
         destination_coldkey_ss58=wallet.coldkeypub.ss58_address,
     )
-    if stake_in_origin < amount:
+    if move_all_stake:
+        amount = stake_in_origin
+
+    elif stake_in_origin < amount:
         logging.error(
-            f":cross_mark: [red]Failed[/red]: Insufficient stake in origin hotkey: {origin_hotkey}. Stake: {stake_in_origin}, amount: {amount}"
+            f":cross_mark: [red]Failed[/red]: Insufficient stake in origin hotkey: {origin_hotkey}. "
+            f"Stake: {stake_in_origin}, amount: {amount}"
         )
         return False
 
+    amount.set_unit(netuid=origin_netuid)
+
     try:
         logging.info(
             f"Moving stake from hotkey [blue]{origin_hotkey}[/blue] to hotkey [blue]{destination_hotkey}[/blue]\n"

bittensor/core/subtensor.py

@@ -3493,10 +3493,11 @@ def move_stake(
         origin_netuid: int,
         destination_hotkey: str,
         destination_netuid: int,
-        amount: Balance,
+        amount: Optional[Balance] = None,
         wait_for_inclusion: bool = True,
         wait_for_finalization: bool = False,
         period: Optional[int] = None,
+        move_all_stake: bool = False,
     ) -> bool:
         """
         Moves stake to a different hotkey and/or subnet.
@@ -3513,6 +3514,7 @@ def move_stake(
             period (Optional[int]): 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.
+            move_all_stake: If true, moves all stake from the source hotkey to the destination hotkey.
 
         Returns:
             success (bool): True if the stake movement was successful.
@@ -3529,6 +3531,7 @@ def move_stake(
             wait_for_inclusion=wait_for_inclusion,
             wait_for_finalization=wait_for_finalization,
             period=period,
+            move_all_stake=move_all_stake,
         )
 
     def register(

bittensor/core/subtensor_api/staking.py

@@ -22,6 +22,7 @@ def __init__(self, subtensor: Union["_Subtensor", "_AsyncSubtensor"]):
         self.get_stake_operations_fee = subtensor.get_stake_operations_fee
         self.get_stake_weight = subtensor.get_stake_weight
         self.get_unstake_fee = subtensor.get_unstake_fee
+        self.move_stake = subtensor.move_stake
         self.unstake = subtensor.unstake
         self.unstake_all = subtensor.unstake_all
         self.unstake_multiple = subtensor.unstake_multiple

tests/e2e_tests/test_delegate.py

@@ -13,7 +13,7 @@
     vote,
 )
 from tests.e2e_tests.utils.e2e_test_utils import wait_to_start_call
-from tests.helpers.helpers import CLOSE_IN_VALUE
+from tests.helpers.helpers import CloseInValue
 
 DEFAULT_DELEGATE_TAKE = 0.179995422293431
 
@@ -449,7 +449,7 @@ def test_get_vote_data(subtensor, alice_wallet):
 
     assert proposal == ProposalVoteData(
         ayes=[],
-        end=CLOSE_IN_VALUE(1_000_000, subtensor.block),
+        end=CloseInValue(1_000_000, subtensor.block),
         index=0,
         nays=[],
         threshold=3,
@@ -475,7 +475,7 @@ def test_get_vote_data(subtensor, alice_wallet):
         ayes=[
             alice_wallet.hotkey.ss58_address,
         ],
-        end=CLOSE_IN_VALUE(1_000_000, subtensor.block),
+        end=CloseInValue(1_000_000, subtensor.block),
         index=0,
         nays=[],
         threshold=3,

tests/e2e_tests/test_staking.py

@@ -9,7 +9,7 @@
     sudo_set_admin_utils,
 )
 from tests.e2e_tests.utils.e2e_test_utils import wait_to_start_call
-from tests.helpers.helpers import CLOSE_IN_VALUE
+from tests.helpers.helpers import CloseInValue
 
 
 def test_single_operation(subtensor, alice_wallet, bob_wallet):
@@ -323,7 +323,7 @@ def test_batch_operations(subtensor, alice_wallet, bob_wallet):
         bob_wallet.coldkey.ss58_address,
     )
 
-    assert CLOSE_IN_VALUE(  # Make sure we are within 0.0001 TAO due to tx fees
+    assert CloseInValue(  # Make sure we are within 0.0001 TAO due to tx fees
         balances[bob_wallet.coldkey.ss58_address], Balance.from_rao(100_000)
     ) == Balance.from_tao(999_999.7994)
 
@@ -643,11 +643,12 @@ def test_safe_swap_stake_scenarios(subtensor, alice_wallet, bob_wallet):
     )
 
 
-def test_move_stake(subtensor, alice_wallet, bob_wallet):
+def test_move_stake(subtensor, alice_wallet, bob_wallet, dave_wallet):
     """
     Tests:
     - Adding stake
     - Moving stake from one hotkey-subnet pair to another
+    - Testing `move_stake` method with `move_all_stake=True` flag.
     """
 
     alice_subnet_netuid = subtensor.get_total_subnets()  # 2
@@ -658,16 +659,9 @@ def test_move_stake(subtensor, alice_wallet, bob_wallet):
 
     assert wait_to_start_call(subtensor, alice_wallet, alice_subnet_netuid)
 
-    subtensor.burned_register(
-        alice_wallet,
-        netuid=alice_subnet_netuid,
-        wait_for_inclusion=True,
-        wait_for_finalization=True,
-    )
-
     assert subtensor.add_stake(
-        alice_wallet,
-        alice_wallet.hotkey.ss58_address,
+        wallet=alice_wallet,
+        hotkey_ss58=alice_wallet.hotkey.ss58_address,
         netuid=alice_subnet_netuid,
         amount=Balance.from_tao(1_000),
         wait_for_inclusion=True,
@@ -697,8 +691,22 @@ def test_move_stake(subtensor, alice_wallet, bob_wallet):
 
     assert wait_to_start_call(subtensor, bob_wallet, bob_subnet_netuid)
 
+    subtensor.burned_register(
+        wallet=bob_wallet,
+        netuid=alice_subnet_netuid,
+        wait_for_inclusion=True,
+        wait_for_finalization=True,
+    )
+
+    subtensor.burned_register(
+        wallet=dave_wallet,
+        netuid=alice_subnet_netuid,
+        wait_for_inclusion=True,
+        wait_for_finalization=True,
+    )
+
     assert subtensor.move_stake(
-        alice_wallet,
+        wallet=alice_wallet,
         origin_hotkey=alice_wallet.hotkey.ss58_address,
         origin_netuid=alice_subnet_netuid,
         destination_hotkey=bob_wallet.hotkey.ss58_address,
@@ -743,9 +751,59 @@ def test_move_stake(subtensor, alice_wallet, bob_wallet):
     )
 
     expected_stakes += fast_block_stake
-
     assert stakes == expected_stakes
-    logging.console.success("✅ Test [green]test_move_stake[/green] passed")
+
+    # test move_stake with move_all_stake=True
+    dave_stake = subtensor.staking.get_stake(
+        coldkey_ss58=dave_wallet.coldkey.ss58_address,
+        hotkey_ss58=bob_wallet.hotkey.ss58_address,
+        netuid=bob_subnet_netuid,
+    )
+    logging.console.info(f"[orange]Dave stake before adding: {dave_stake}[orange]")
+
+    assert subtensor.staking.add_stake(
+        wallet=dave_wallet,
+        hotkey_ss58=dave_wallet.hotkey.ss58_address,
+        netuid=bob_subnet_netuid,
+        amount=Balance.from_tao(1000),
+        wait_for_inclusion=True,
+        wait_for_finalization=True,
+        allow_partial_stake=True,
+    )
+
+    dave_stake = subtensor.staking.get_stake(
+        coldkey_ss58=dave_wallet.coldkey.ss58_address,
+        hotkey_ss58=dave_wallet.hotkey.ss58_address,
+        netuid=bob_subnet_netuid,
+    )
+    logging.console.info(f"[orange]Dave stake after adding: {dave_stake}[orange]")
+
+    # let chain to process the transaction
+    subtensor.wait_for_block(
+        subtensor.block + subtensor.subnets.tempo(netuid=bob_subnet_netuid)
+    )
+
+    assert subtensor.staking.move_stake(
+        wallet=dave_wallet,
+        origin_hotkey=dave_wallet.hotkey.ss58_address,
+        origin_netuid=bob_subnet_netuid,
+        destination_hotkey=bob_wallet.hotkey.ss58_address,
+        destination_netuid=bob_subnet_netuid,
+        wait_for_inclusion=True,
+        wait_for_finalization=True,
+        move_all_stake=True,
+    )
+
+    dave_stake = subtensor.staking.get_stake(
+        coldkey_ss58=dave_wallet.coldkey.ss58_address,
+        hotkey_ss58=dave_wallet.hotkey.ss58_address,
+        netuid=bob_subnet_netuid,
+    )
+    logging.console.info(f"[orange]Dave stake after moving all: {dave_stake}[orange]")
+
+    assert dave_stake.rao == CloseInValue(0, 0.00001)
+
+    logging.console.success("✅ Test [green]test_move_stake[/green] passed.")
 
 
 def test_transfer_stake(subtensor, alice_wallet, bob_wallet, dave_wallet):

tests/helpers/__init__.py

@@ -1,6 +1,6 @@
 import os
 from .helpers import (  # noqa: F401
-    CLOSE_IN_VALUE,
+    CloseInValue,
     __mock_wallet_factory__,
 )
 from bittensor_wallet.mock.wallet_mock import (  # noqa: F401