Merge pull request #2736 from opentensor/feat/safe-staking

Feat/safe staking

Author: ibraheem-abe

bittensor/core/async_subtensor.py

@@ -2810,6 +2810,9 @@ async def add_stake(
         amount: Optional[Balance] = None,
         wait_for_inclusion: bool = True,
         wait_for_finalization: bool = False,
+        safe_staking: bool = False,
+        allow_partial_stake: bool = False,
+        rate_threshold: float = 0.005,
     ) -> bool:
         """
         Adds the specified amount of stake to a neuron identified by the hotkey ``SS58`` address.
@@ -2823,12 +2826,20 @@ async def add_stake(
             amount (Balance): The amount of TAO to stake.
             wait_for_inclusion (bool): Waits for the transaction to be included in a block.
             wait_for_finalization (bool): Waits for the transaction to be finalized on the blockchain.
+            safe_staking (bool): If true, enables price safety checks to protect against fluctuating prices. The stake
+                will only execute if the price change doesn't exceed the rate threshold. Default is False.
+            allow_partial_stake (bool): If true and safe_staking is enabled, allows partial staking when
+                the full amount would exceed the price threshold. If false, the entire stake fails if it would
+                exceed the threshold. Default is False.
+            rate_threshold (float): The maximum allowed price change ratio when staking. For example,
+                0.005 = 0.5% maximum price increase. Only used when safe_staking is True. Default is 0.005.
 
         Returns:
             bool: ``True`` if the staking is successful, False otherwise.
 
-        This function enables neurons to increase their stake in the network, enhancing their influence and potential
-            rewards in line with Bittensor's consensus and reward mechanisms.
+        This function enables neurons to increase their stake in the network, enhancing their influence and potential.
+            When safe_staking is enabled, it provides protection against price fluctuations during the time stake is
+            executed and the time it is actually processed by the chain.
         """
         amount = check_and_convert_to_balance(amount)
         return await add_stake_extrinsic(
@@ -2839,6 +2850,9 @@ async def add_stake(
             amount=amount,
             wait_for_inclusion=wait_for_inclusion,
             wait_for_finalization=wait_for_finalization,
+            safe_staking=safe_staking,
+            allow_partial_stake=allow_partial_stake,
+            rate_threshold=rate_threshold,
         )
 
     async def add_stake_multiple(
@@ -3460,6 +3474,9 @@ async def swap_stake(
         amount: Balance,
         wait_for_inclusion: bool = True,
         wait_for_finalization: bool = False,
+        safe_staking: bool = False,
+        allow_partial_stake: bool = False,
+        rate_threshold: float = 0.005,
     ) -> bool:
         """
         Moves stake between subnets while keeping the same coldkey-hotkey pair ownership.
@@ -3473,9 +3490,25 @@ async def swap_stake(
             amount (Union[Balance, float]): The amount to swap.
             wait_for_inclusion (bool): Waits for the transaction to be included in a block.
             wait_for_finalization (bool): Waits for the transaction to be finalized on the blockchain.
+            safe_staking (bool): If true, enables price safety checks to protect against fluctuating prices. The swap
+                will only execute if the price ratio between subnets doesn't exceed the rate threshold.
+                Default is False.
+            allow_partial_stake (bool): If true and safe_staking is enabled, allows partial stake swaps when
+                the full amount would exceed the price threshold. If false, the entire swap fails if it would
+                exceed the threshold. Default is False.
+            rate_threshold (float): The maximum allowed increase in the price ratio between subnets
+                (origin_price/destination_price). For example, 0.005 = 0.5% maximum increase. Only used
+                when safe_staking is True. Default is 0.005.
 
         Returns:
             success (bool): True if the extrinsic was successful.
+
+        The price ratio for swap_stake in safe mode is calculated as: origin_subnet_price / destination_subnet_price
+        When safe_staking is enabled, the swap will only execute if:
+            - With allow_partial_stake=False: The entire swap amount can be executed without the price ratio
+            increasing more than rate_threshold
+            - With allow_partial_stake=True: A partial amount will be swapped up to the point where the
+            price ratio would increase by rate_threshold
         """
         amount = check_and_convert_to_balance(amount)
         return await swap_stake_extrinsic(
@@ -3487,6 +3520,9 @@ async def swap_stake(
             amount=amount,
             wait_for_inclusion=wait_for_inclusion,
             wait_for_finalization=wait_for_finalization,
+            safe_staking=safe_staking,
+            allow_partial_stake=allow_partial_stake,
+            rate_threshold=rate_threshold,
         )
 
     async def transfer_stake(
@@ -3575,6 +3611,9 @@ async def unstake(
         amount: Optional[Balance] = None,
         wait_for_inclusion: bool = True,
         wait_for_finalization: bool = False,
+        safe_staking: bool = False,
+        allow_partial_stake: bool = False,
+        rate_threshold: float = 0.005,
     ) -> bool:
         """
         Removes a specified amount of stake from a single hotkey account. This function is critical for adjusting
@@ -3584,10 +3623,17 @@ async def unstake(
             wallet (bittensor_wallet.wallet): The wallet associated with the neuron from which the stake is being
                 removed.
             hotkey_ss58 (Optional[str]): The ``SS58`` address of the hotkey account to unstake from.
-            netuid (Optional[int]): Subnet unique ID.
+            netuid (Optional[int]): The unique identifier of the subnet.
             amount (Balance): The amount of TAO to unstake. If not specified, unstakes all.
             wait_for_inclusion (bool): Waits for the transaction to be included in a block.
             wait_for_finalization (bool): Waits for the transaction to be finalized on the blockchain.
+            safe_staking (bool): If true, enables price safety checks to protect against fluctuating prices. The unstake
+                will only execute if the price change doesn't exceed the rate threshold. Default is False.
+            allow_partial_stake (bool): If true and safe_staking is enabled, allows partial unstaking when
+                the full amount would exceed the price threshold. If false, the entire unstake fails if it would
+                exceed the threshold. Default is False.
+            rate_threshold (float): The maximum allowed price change ratio when unstaking. For example,
+                0.005 = 0.5% maximum price decrease. Only used when safe_staking is True. Default is 0.005.
 
         Returns:
             bool: ``True`` if the unstaking process is successful, False otherwise.
@@ -3604,6 +3650,9 @@ async def unstake(
             amount=amount,
             wait_for_inclusion=wait_for_inclusion,
             wait_for_finalization=wait_for_finalization,
+            safe_staking=safe_staking,
+            allow_partial_stake=allow_partial_stake,
+            rate_threshold=rate_threshold,
         )
 
     async def unstake_multiple(

bittensor/core/chain_data/dynamic_info.py

@@ -77,9 +77,9 @@ def _from_dict(cls, decoded: dict) -> "DynamicInfo":
         price = (
             Balance.from_tao(1.0)
             if netuid == 0
-            else Balance.from_tao(tao_in.tao / alpha_in.tao)
+            else Balance.from_tao(tao_in.tao / alpha_in.tao).set_unit(netuid)
             if alpha_in.tao > 0
-            else Balance.from_tao(1)
+            else Balance.from_tao(1).set_unit(netuid)
         )  # Root always has 1-1 price
 
         if decoded.get("subnet_identity"):

bittensor/core/extrinsics/asyncex/move_stake.py

@@ -160,6 +160,9 @@ async def swap_stake_extrinsic(
     amount: Balance,
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
+    safe_staking: bool = False,
+    allow_partial_stake: bool = False,
+    rate_threshold: float = 0.005,
 ) -> bool:
     """
     Swaps stake from one subnet to another for a given hotkey in the Bittensor network.
@@ -173,6 +176,9 @@ async def swap_stake_extrinsic(
         amount (Balance): The amount of stake to swap 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.
+        safe_staking (bool): If true, enables price safety checks to protect against price impact.
+        allow_partial_stake (bool): If true, allows partial stake swaps when the full amount would exceed the price threshold.
+        rate_threshold (float): Maximum allowed increase in price ratio (0.005 = 0.5%).
 
     Returns:
         bool: True if the swap was successful, False otherwise.
@@ -205,20 +211,47 @@ async def swap_stake_extrinsic(
         return False
 
     try:
-        logging.info(
-            f"Swapping stake for hotkey [blue]{hotkey_ss58}[/blue]\n"
-            f"Amount: [green]{amount}[/green] from netuid [yellow]{origin_netuid}[/yellow] to netuid "
-            f"[yellow]{destination_netuid}[/yellow]"
-        )
+        call_params = {
+            "hotkey": hotkey_ss58,
+            "origin_netuid": origin_netuid,
+            "destination_netuid": destination_netuid,
+            "alpha_amount": amount.rao,
+        }
+
+        if safe_staking:
+            origin_pool, destination_pool = await asyncio.gather(
+                subtensor.subnet(netuid=origin_netuid),
+                subtensor.subnet(netuid=destination_netuid),
+            )
+            swap_rate_ratio = origin_pool.price.rao / destination_pool.price.rao
+            swap_rate_ratio_with_tolerance = swap_rate_ratio * (1 + rate_threshold)
+
+            logging.info(
+                f"Swapping stake with safety for hotkey [blue]{hotkey_ss58}[/blue]\n"
+                f"Amount: [green]{amount}[/green] from netuid [green]{origin_netuid}[/green] to netuid "
+                f"[green]{destination_netuid}[/green]\n"
+                f"Current price ratio: [green]{swap_rate_ratio:.4f}[/green], "
+                f"Ratio with tolerance: [green]{swap_rate_ratio_with_tolerance:.4f}[/green]"
+            )
+            call_params.update(
+                {
+                    "limit_price": swap_rate_ratio_with_tolerance,
+                    "allow_partial": allow_partial_stake,
+                }
+            )
+            call_function = "swap_stake_limit"
+        else:
+            logging.info(
+                f"Swapping stake for hotkey [blue]{hotkey_ss58}[/blue]\n"
+                f"Amount: [green]{amount}[/green] from netuid [green]{origin_netuid}[/green] to netuid "
+                f"[green]{destination_netuid}[/green]"
+            )
+            call_function = "swap_stake"
+
         call = await subtensor.substrate.compose_call(
             call_module="SubtensorModule",
-            call_function="swap_stake",
-            call_params={
-                "hotkey": hotkey_ss58,
-                "origin_netuid": origin_netuid,
-                "destination_netuid": destination_netuid,
-                "alpha_amount": amount.rao,
-            },
+            call_function=call_function,
+            call_params=call_params,
         )
 
         success, err_msg = await subtensor.sign_and_send_extrinsic(
@@ -253,7 +286,12 @@ async def swap_stake_extrinsic(
 
             return True
         else:
-            logging.error(f":cross_mark: [red]Failed[/red]: {err_msg}")
+            if safe_staking and "Custom error: 8" in err_msg:
+                logging.error(
+                    ":cross_mark: [red]Failed[/red]: Price ratio exceeded tolerance limit. Either increase price tolerance or enable partial staking."
+                )
+            else:
+                logging.error(f":cross_mark: [red]Failed[/red]: {err_msg}")
             return False
 
     except Exception as e:

bittensor/core/extrinsics/asyncex/staking.py

@@ -21,6 +21,9 @@ async def add_stake_extrinsic(
     amount: Optional[Balance] = None,
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
+    safe_staking: bool = False,
+    allow_partial_stake: bool = False,
+    rate_threshold: float = 0.005,
 ) -> bool:
     """
     Adds the specified amount of stake to passed hotkey `uid`.
@@ -36,6 +39,9 @@ async def add_stake_extrinsic(
             `False` if the extrinsic fails to enter the block within the timeout.
         wait_for_finalization: If set, waits for the extrinsic to be finalized on the chain before returning `True`,
             or returns `False` if the extrinsic fails to be finalized within the timeout.
+        safe_staking: If set, uses safe staking logic
+        allow_partial_stake: If set, allows partial stake
+        rate_threshold: The rate threshold for safe staking
 
     Returns:
         success: Flag is `True` if extrinsic was finalized or included in the block. If we did not wait for
@@ -97,19 +103,48 @@ async def add_stake_extrinsic(
         return False
 
     try:
-        logging.info(
-            f":satellite: [magenta]Staking to:[/magenta] "
-            f"[blue]netuid: {netuid}, amount: {staking_balance} "
-            f"on {subtensor.network}[/blue] [magenta]...[/magenta]"
-        )
+        call_params = {
+            "hotkey": hotkey_ss58,
+            "netuid": netuid,
+            "amount_staked": staking_balance.rao,
+        }
+
+        if safe_staking:
+            pool = await subtensor.subnet(netuid=netuid)
+            base_price = pool.price.rao
+            price_with_tolerance = base_price * (1 + rate_threshold)
+            call_params.update(
+                {
+                    "limit_price": price_with_tolerance,
+                    "allow_partial": allow_partial_stake,
+                }
+            )
+            call_function = "add_stake_limit"
+
+            # For logging
+            base_rate = pool.price.tao
+            rate_with_tolerance = base_rate * (1 + rate_threshold)
+            logging.info(
+                f":satellite: [magenta]Safe Staking to:[/magenta] "
+                f"[blue]netuid: [green]{netuid}[/green], amount: [green]{staking_balance}[/green], "
+                f"tolerance percentage: [green]{rate_threshold*100}%[/green], "
+                f"price limit: [green]{rate_with_tolerance}[/green], "
+                f"original price: [green]{base_rate}[/green], "
+                f"with partial stake: [green]{allow_partial_stake}[/green] "
+                f"on [blue]{subtensor.network}[/blue][/magenta]...[/magenta]"
+            )
+        else:
+            logging.info(
+                f":satellite: [magenta]Staking to:[/magenta] "
+                f"[blue]netuid: [green]{netuid}[/green], amount: [green]{staking_balance}[/green] "
+                f"on [blue]{subtensor.network}[/blue][magenta]...[/magenta]"
+            )
+            call_function = "add_stake"
+
         call = await subtensor.substrate.compose_call(
             call_module="SubtensorModule",
-            call_function="add_stake",
-            call_params={
-                "hotkey": hotkey_ss58,
-                "amount_staked": staking_balance.rao,
-                "netuid": netuid,
-            },
+            call_function=call_function,
+            call_params=call_params,
         )
         staking_response, err_msg = await subtensor.sign_and_send_extrinsic(
             call,
@@ -152,7 +187,12 @@ async def add_stake_extrinsic(
             )
             return True
         else:
-            logging.error(f":cross_mark: [red]Failed: {err_msg}.[/red]")
+            if safe_staking and "Custom error: 8" in err_msg:
+                logging.error(
+                    ":cross_mark: [red]Failed[/red]: Price exceeded tolerance limit. Either increase price tolerance or enable partial staking."
+                )
+            else:
+                logging.error(f":cross_mark: [red]Failed: {err_msg}.[/red]")
             return False
 
     except NotRegisteredError: