Merge pull request #620 from opentensor/arbitrage_coldkeys

Schedule Coldkey Swap

node/src/rpc.rs

@@ -60,6 +60,7 @@ where
     C::Api: subtensor_custom_rpc_runtime_api::NeuronInfoRuntimeApi<Block>,
     C::Api: subtensor_custom_rpc_runtime_api::SubnetInfoRuntimeApi<Block>,
     C::Api: subtensor_custom_rpc_runtime_api::SubnetRegistrationRuntimeApi<Block>,
+    C::Api: subtensor_custom_rpc_runtime_api::ColdkeySwapRuntimeApi<Block>,
     B: sc_client_api::Backend<Block> + Send + Sync + 'static,
     P: TransactionPool + 'static,
 {

pallets/admin-utils/tests/mock.rs

@@ -114,6 +114,7 @@ parameter_types! {
     pub const InitialAlphaHigh: u16 = 58982; // Represents 0.9 as per the production default
     pub const InitialAlphaLow: u16 = 45875; // Represents 0.7 as per the production default
     pub const InitialLiquidAlphaOn: bool = false; // Default value for LiquidAlphaOn
+    pub const InitialBaseDifficulty: u64 = 10_000; // Base difficulty
 }
 
 impl pallet_subtensor::Config for Test {
@@ -169,6 +170,7 @@ impl pallet_subtensor::Config for Test {
     type AlphaHigh = InitialAlphaHigh;
     type AlphaLow = InitialAlphaLow;
     type LiquidAlphaOn = InitialLiquidAlphaOn;
+    type InitialBaseDifficulty = InitialBaseDifficulty;
 }
 
 #[derive_impl(frame_system::config_preludes::TestDefaultConfig)]

pallets/subtensor/rpc/src/lib.rs

@@ -12,7 +12,7 @@ use std::sync::Arc;
 use sp_api::ProvideRuntimeApi;
 
 pub use subtensor_custom_rpc_runtime_api::{
-    DelegateInfoRuntimeApi, NeuronInfoRuntimeApi, SubnetInfoRuntimeApi,
+    ColdkeySwapRuntimeApi, DelegateInfoRuntimeApi, NeuronInfoRuntimeApi, SubnetInfoRuntimeApi,
     SubnetRegistrationRuntimeApi,
 };
 
@@ -51,6 +51,24 @@ pub trait SubtensorCustomApi<BlockHash> {
 
     #[method(name = "subnetInfo_getLockCost")]
     fn get_network_lock_cost(&self, at: Option<BlockHash>) -> RpcResult<u64>;
+    #[method(name = "coldkeySwap_getScheduledColdkeySwap")]
+    fn get_scheduled_coldkey_swap(
+        &self,
+        coldkey_account_vec: Vec<u8>,
+        at: Option<BlockHash>,
+    ) -> RpcResult<Vec<u8>>;
+    #[method(name = "coldkeySwap_getRemainingArbitrationPeriod")]
+    fn get_remaining_arbitration_period(
+        &self,
+        coldkey_account_vec: Vec<u8>,
+        at: Option<BlockHash>,
+    ) -> RpcResult<Vec<u8>>;
+    #[method(name = "coldkeySwap_getColdkeySwapDestinations")]
+    fn get_coldkey_swap_destinations(
+        &self,
+        coldkey_account_vec: Vec<u8>,
+        at: Option<BlockHash>,
+    ) -> RpcResult<Vec<u8>>;
 }
 
 pub struct SubtensorCustom<C, P> {
@@ -99,6 +117,7 @@ where
     C::Api: NeuronInfoRuntimeApi<Block>,
     C::Api: SubnetInfoRuntimeApi<Block>,
     C::Api: SubnetRegistrationRuntimeApi<Block>,
+    C::Api: ColdkeySwapRuntimeApi<Block>,
 {
     fn get_delegates(&self, at: Option<<Block as BlockT>::Hash>) -> RpcResult<Vec<u8>> {
         let api = self.client.runtime_api();
@@ -223,4 +242,51 @@ where
             Error::RuntimeError(format!("Unable to get subnet lock cost: {:?}", e)).into()
         })
     }
+
+    fn get_scheduled_coldkey_swap(
+        &self,
+        coldkey_account_vec: Vec<u8>,
+        at: Option<<Block as BlockT>::Hash>,
+    ) -> RpcResult<Vec<u8>> {
+        let api = self.client.runtime_api();
+        let at = at.unwrap_or_else(|| self.client.info().best_hash);
+
+        api.get_scheduled_coldkey_swap(at, coldkey_account_vec)
+            .map_err(|e| {
+                Error::RuntimeError(format!("Unable to get scheduled coldkey swap: {:?}", e)).into()
+            })
+    }
+
+    fn get_remaining_arbitration_period(
+        &self,
+        coldkey_account_vec: Vec<u8>,
+        at: Option<<Block as BlockT>::Hash>,
+    ) -> RpcResult<Vec<u8>> {
+        let api = self.client.runtime_api();
+        let at = at.unwrap_or_else(|| self.client.info().best_hash);
+
+        api.get_remaining_arbitration_period(at, coldkey_account_vec)
+            .map_err(|e| {
+                Error::RuntimeError(format!(
+                    "Unable to get remaining arbitration period: {:?}",
+                    e
+                ))
+                .into()
+            })
+    }
+
+    fn get_coldkey_swap_destinations(
+        &self,
+        coldkey_account_vec: Vec<u8>,
+        at: Option<<Block as BlockT>::Hash>,
+    ) -> RpcResult<Vec<u8>> {
+        let api = self.client.runtime_api();
+        let at = at.unwrap_or_else(|| self.client.info().best_hash);
+
+        api.get_coldkey_swap_destinations(at, coldkey_account_vec)
+            .map_err(|e| {
+                Error::RuntimeError(format!("Unable to get coldkey swap destinations: {:?}", e))
+                    .into()
+            })
+    }
 }

pallets/subtensor/runtime-api/src/lib.rs

@@ -32,4 +32,10 @@ sp_api::decl_runtime_apis! {
     pub trait SubnetRegistrationRuntimeApi {
         fn get_network_registration_cost() -> u64;
     }
+
+    pub trait ColdkeySwapRuntimeApi {
+        fn get_scheduled_coldkey_swap( coldkey_account_vec: Vec<u8> ) -> Vec<u8>;
+        fn get_remaining_arbitration_period( coldkey_account_vec: Vec<u8> ) -> Vec<u8>;
+        fn get_coldkey_swap_destinations( coldkey_account_vec: Vec<u8> ) -> Vec<u8>;
+    }
 }

pallets/subtensor/src/benchmarks.rs

@@ -428,4 +428,31 @@ reveal_weights {
     let _ = Subtensor::<T>::commit_weights(<T as frame_system::Config>::RuntimeOrigin::from(RawOrigin::Signed(hotkey.clone())), netuid, commit_hash);
 
   }: reveal_weights(RawOrigin::Signed(hotkey.clone()), netuid, uids, weight_values, salt, version_key)
+
+  schedule_coldkey_swap {
+    let seed: u32 = 1;
+    let old_coldkey: T::AccountId = account("OldColdkey", 0, seed);
+    let new_coldkey: T::AccountId = account("NewColdkey", 0, seed + 1);
+    let hotkey: T::AccountId = account("Hotkey", 0, seed);
+
+    let netuid = 1u16;
+    let tempo = 1u16;
+    let block_number: u64 = Subtensor::<T>::get_current_block_as_u64();
+    let nonce = 0;
+
+    // Initialize the network
+    Subtensor::<T>::init_new_network(netuid, tempo);
+    Subtensor::<T>::set_network_registration_allowed(netuid, true);
+
+    // Add balance to the old coldkey account
+    let amount_to_be_staked: u64 = 1000000u32.into();
+    Subtensor::<T>::add_balance_to_coldkey_account(&old_coldkey.clone(), amount_to_be_staked+1000000000);
+       // Burned register the hotkey with the old coldkey
+       assert_ok!(Subtensor::<T>::burned_register(
+        RawOrigin::Signed(old_coldkey.clone()).into(),
+        netuid,
+        hotkey.clone()
+    ));
+
+  }: schedule_coldkey_swap(RawOrigin::Signed(old_coldkey.clone()), new_coldkey.clone(), vec![], block_number, nonce)
 }

pallets/subtensor/src/delegate_info.rs

@@ -41,21 +41,19 @@ impl<T: Config> Pallet<T> {
         let mut emissions_per_day: U64F64 = U64F64::from_num(0);
 
         for netuid in registrations.iter() {
-            let _uid = Self::get_uid_for_net_and_hotkey(*netuid, &delegate.clone());
-            if _uid.is_err() {
-                continue; // this should never happen
-            } else {
-                let uid = _uid.expect("Delegate's UID should be ok");
+            if let Ok(uid) = Self::get_uid_for_net_and_hotkey(*netuid, &delegate.clone()) {
                 let validator_permit = Self::get_validator_permit_for_uid(*netuid, uid);
                 if validator_permit {
                     validator_permits.push((*netuid).into());
                 }
 
                 let emission: U64F64 = Self::get_emission_for_uid(*netuid, uid).into();
                 let tempo: U64F64 = Self::get_tempo(*netuid).into();
-                let epochs_per_day: U64F64 = U64F64::from_num(7200).saturating_div(tempo);
-                emissions_per_day =
-                    emissions_per_day.saturating_add(emission.saturating_mul(epochs_per_day));
+                if tempo > U64F64::from_num(0) {
+                    let epochs_per_day: U64F64 = U64F64::from_num(7200).saturating_div(tempo);
+                    emissions_per_day =
+                        emissions_per_day.saturating_add(emission.saturating_mul(epochs_per_day));
+                }
             }
         }
 
@@ -64,15 +62,15 @@ impl<T: Config> Pallet<T> {
 
         let total_stake: U64F64 = Self::get_total_stake_for_hotkey(&delegate.clone()).into();
 
-        let mut return_per_1000: U64F64 = U64F64::from_num(0);
-
-        if total_stake > U64F64::from_num(0) {
-            return_per_1000 = emissions_per_day
+        let return_per_1000: U64F64 = if total_stake > U64F64::from_num(0) {
+            emissions_per_day
                 .saturating_mul(U64F64::from_num(0.82))
-                .saturating_div(total_stake.saturating_div(U64F64::from_num(1000)));
-        }
+                .saturating_div(total_stake.saturating_div(U64F64::from_num(1000)))
+        } else {
+            U64F64::from_num(0)
+        };
 
-        return DelegateInfo {
+        DelegateInfo {
             delegate_ss58: delegate.clone(),
             take,
             nominators,
@@ -81,7 +79,7 @@ impl<T: Config> Pallet<T> {
             validator_permits,
             return_per_1000: U64F64::to_num::<u64>(return_per_1000).into(),
             total_daily_return: U64F64::to_num::<u64>(emissions_per_day).into(),
-        };
+        }
     }
 
     pub fn get_delegate(delegate_account_vec: Vec<u8>) -> Option<DelegateInfo<T>> {
@@ -133,4 +131,47 @@ impl<T: Config> Pallet<T> {
 
         delegates
     }
+
+    pub fn get_total_delegated_stake(coldkey: &T::AccountId) -> u64 {
+        let mut total_delegated = 0u64;
+
+        // Get all hotkeys associated with this coldkey
+        let hotkeys = StakingHotkeys::<T>::get(coldkey);
+
+        for hotkey in hotkeys {
+            let owner = Owner::<T>::get(&hotkey);
+
+            for (delegator, stake) in Stake::<T>::iter_prefix(&hotkey) {
+                if delegator != owner {
+                    total_delegated = total_delegated.saturating_add(stake);
+                }
+            }
+        }
+
+        log::info!(
+            "Total delegated stake for coldkey {:?}: {}",
+            coldkey,
+            total_delegated
+        );
+        total_delegated
+    }
+
+    // Helper function to get total delegated stake for a hotkey
+    pub fn get_total_hotkey_delegated_stake(hotkey: &T::AccountId) -> u64 {
+        let mut total_delegated = 0u64;
+
+        // Iterate through all delegators for this hotkey
+        for (delegator, stake) in Stake::<T>::iter_prefix(hotkey) {
+            if delegator != Self::get_coldkey_for_hotkey(hotkey) {
+                total_delegated = total_delegated.saturating_add(stake);
+            }
+        }
+
+        total_delegated
+    }
+
+    // Helper function to get the coldkey associated with a hotkey
+    pub fn get_coldkey_for_hotkey(hotkey: &T::AccountId) -> T::AccountId {
+        Owner::<T>::get(hotkey)
+    }
 }

pallets/subtensor/src/errors.rs

@@ -146,5 +146,15 @@ mod errors {
         NoBalanceToTransfer,
         /// Same coldkey
         SameColdkey,
+        /// The coldkey is in arbitration
+        ColdkeyIsInArbitration,
+        /// The new coldkey is already registered for the drain
+        DuplicateColdkey,
+        /// Error thrown on a coldkey swap.
+        ColdkeySwapError,
+        /// Insufficient Balance to Schedule coldkey swap
+        InsufficientBalanceToPerformColdkeySwap,
+        /// The maximum number of coldkey destinations has been reached
+        MaxColdkeyDestinationsReached,
     }
 }

pallets/subtensor/src/events.rs

@@ -150,5 +150,19 @@ mod events {
                 <T as frame_system::Config>::AccountId,
             >>::Balance,
         },
+        /// A coldkey swap has been scheduled
+        ColdkeySwapScheduled {
+            /// The account ID of the old coldkey
+            old_coldkey: T::AccountId,
+            /// The account ID of the new coldkey
+            new_coldkey: T::AccountId,
+            /// The arbitration block for the coldkey swap
+            arbitration_block: u64,
+        },
+        /// The arbitration period has been extended
+        ArbitrationPeriodExtended {
+            /// The account ID of the coldkey
+            coldkey: T::AccountId,
+        },
     }
 }