diff --git a/bittensor/utils/mock/subtensor_mock.py b/bittensor/utils/mock/subtensor_mock.py index 9773d525a4..51925b4a88 100644 --- a/bittensor/utils/mock/subtensor_mock.py +++ b/bittensor/utils/mock/subtensor_mock.py @@ -1,3 +1,7 @@ +import base58 +import time +import bittensor.core.subtensor as subtensor_module + from collections.abc import Mapping from dataclasses import dataclass from hashlib import sha256 @@ -9,18 +13,20 @@ from async_substrate_interface import SubstrateInterface from bittensor_wallet import Wallet -import bittensor.core.subtensor as subtensor_module from bittensor.core.chain_data import ( NeuronInfo, NeuronInfoLite, PrometheusInfo, AxonInfo, + MetagraphInfo, + SelectiveMetagraphIndex, ) from bittensor.core.errors import ChainQueryError from bittensor.core.subtensor import Subtensor from bittensor.core.types import AxonServeCallParams, PrometheusServeCallParams from bittensor.utils import RAOPERTAO, u16_normalized_float -from bittensor.utils.balance import Balance +from bittensor.utils.balance import Balance, UNITS + # Mock Testing Constant __GLOBAL_MOCK_STATE__ = {} @@ -242,6 +248,7 @@ def setup(self) -> None: "Commits": {}, "AdjustmentAlpha": {}, "BondsMovingAverage": {}, + "NetworkRegisteredAt": {}, }, } @@ -349,6 +356,8 @@ def create_subnet(self, netuid: int) -> None: subtensor_state["BondsMovingAverage"][netuid] = {} subtensor_state["BondsMovingAverage"][netuid][0] = 1000 + + subtensor_state["NetworkRegisteredAt"][netuid] = time.time() else: raise Exception("Subnet already exists") @@ -389,6 +398,9 @@ def _register_neuron(self, netuid: int, hotkey: str, coldkey: str) -> int: subtensor_state["SubnetworkN"][netuid][self.block_number] = ( subnetwork_n + 1 ) + if uid == 0: + subtensor_state["SubnetOwner"][netuid] = {} + subtensor_state["SubnetOwner"][netuid][hotkey] = coldkey subtensor_state["Stake"][hotkey] = {} subtensor_state["Stake"][hotkey][coldkey] = {} @@ -784,6 +796,11 @@ def _get_most_recent_storage( block_number -= 1 return None + + @staticmethod + def _get_min_block_number(storage: dict[BlockNumber, Any]) -> int: + return min(storage.keys()) + def _get_axon_info( self, netuid: int, hotkey: str, block: Optional[int] = None @@ -1073,3 +1090,163 @@ def do_serve_axon( wait_for_finalization: bool = True, ) -> tuple[bool, Optional[str]]: return True, None + + def get_metagraph_info( + self, + netuid: int, + field_indices: Optional[Union[list[SelectiveMetagraphIndex], list[int]]] = None, + block: Optional[int] = None, + ) -> Optional[MetagraphInfo]: + """ + Mocks the get_metagraph_info method to return the state of the subnet. + """ + + subtensor_state = self.chain_state["SubtensorModule"] + if netuid not in subtensor_state["NetworksAdded"]: + raise Exception("Subnet does not exist") + subnetwork_n = self._get_most_recent_storage( + subtensor_state["SubnetworkN"][netuid] + ) + if subnetwork_n == 0: + raise Exception("Subnet is empty") + + subnet_owner_hotkey = next(iter(subtensor_state["SubnetOwner"][netuid])) + subnet_owner_coldkey = subtensor_state["SubnetOwner"][netuid][subnet_owner_hotkey] + + neurons: list[NeuronInfo] = [ + self._neuron_subnet_exists(uid, netuid, self.block_number) for uid in range(0, subnetwork_n) + ] + + megagraph_info_dict = { + "netuid": netuid, + "name": f"subnet_{netuid}".encode(), + "symbol": UNITS[netuid].encode(), + "network_registered_at": subtensor_state["NetworkRegisteredAt"][netuid], + "owner_hotkey": self._convert_key(subnet_owner_hotkey), + "owner_coldkey": self._convert_key(subnet_owner_coldkey), + "block": self.block_number, + "tempo": subtensor_state["Tempo"][netuid][self.block_number], + "last_step": 1, + "blocks_since_last_step": subtensor_state["BlocksSinceLastStep"][netuid][self.block_number], + "subnet_emission": sum(neurons[uid].emission for uid in range(0, subnetwork_n)), + "alpha_in": 1, + "alpha_out": 1, + "tao_in": 1, + "alpha_out_emission": 1, + "alpha_in_emission": 1, + "tao_in_emission": 1, + "pending_alpha_emission": 1, + "pending_root_emission": 1, + "subnet_volume": sum(neurons[uid].emission for uid in range(0, subnetwork_n)), + "moving_price": {"bits": 1}, + "rho": subtensor_state["Rho"][netuid][self.block_number], + "kappa": subtensor_state["Kappa"][netuid][self.block_number], + "min_allowed_weights": subtensor_state["MinAllowedWeights"][netuid][self.block_number], + "max_allowed_weights": subtensor_state["MaxWeightLimit"][netuid][self.block_number], + "weights_version": 0, + "weights_rate_limit": 1, + "activity_cutoff": 1, + "max_validators": subtensor_state["MaxAllowedValidators"][netuid][self.block_number], + "num_uids": subnetwork_n, + "max_uids": subtensor_state["MaxAllowedUids"][netuid][self.block_number], + "burn": subtensor_state["Burn"][netuid][self.block_number], + "difficulty": subtensor_state["Difficulty"][netuid][self.block_number], + "registration_allowed": True, + "pow_registration_allowed": True, + "immunity_period": subtensor_state["ImmunityPeriod"][netuid][self.block_number], + "min_difficulty": 0, + "max_difficulty": 0, + "max_weights_limit": 0, + "min_burn": 0, + "max_burn": 0, + "adjustment_alpha": 1, + "adjustment_interval": 1, + "target_regs_per_interval": 1, + "max_regs_per_block": 1, + "serving_rate_limit": 1, + "commit_reveal_weights_enabled": True, + "commit_reveal_period": 4, + "liquid_alpha_enabled": True, + "alpha_high": 1, + "alpha_low": 1, + "bonds_moving_avg": 1, + "hotkeys": [self._convert_key(neurons[uid].hotkey) for uid in range(0, subnetwork_n)], + "active": [neurons[uid].active for uid in range(0, subnetwork_n)], + "validator_permit": [neurons[uid].validator_permit for uid in range(0, subnetwork_n)], + "last_update": [neurons[uid].last_update for uid in range(0, subnetwork_n)], + "block_at_registration": [ + self._get_min_block_number(subtensor_state["LastUpdate"][netuid][uid]) + for uid in range(0, subnetwork_n) + ], + "alpha_stake": [neurons[uid].total_stake.tao for uid in range(0, subnetwork_n)], + "tao_stake": [neurons[uid].stake.tao for uid in range(0, subnetwork_n)], + "total_stake": [neurons[uid].total_stake.tao for uid in range(0, subnetwork_n)], + "tao_dividends_per_hotkey": [], + "alpha_dividends_per_hotkey": [], + "identity": None, + "identities": [], + } + return MetagraphInfo.from_dict(megagraph_info_dict) + + def query_runtime_api( + self, + runtime_api: str, + method: str, + params: Optional[Union[list[Any], dict[str, Any]]] = None, + block: Optional[int] = None, + ) -> Any: + """ + Mocks the query_runtime_api method to return the state of the subnet. + """ + block = self.block_number if block is None else block + if runtime_api == "SubnetInfoRuntimeApi" and method == "get_subnet_state" and isinstance(params, list) and params: + return self._get_subtensor_state_dict(netuid=params[0]) + super().query_runtime_api(runtime_api, method, params, block) + + def _get_subtensor_state_dict(self, netuid: int) -> dict[str, Any]: + """ + Returns the state of the subnet as a dictionary. + This is used to sync the metagraph with the subtensor. + """ + subtensor_state = self.chain_state["SubtensorModule"] + if netuid not in subtensor_state["NetworksAdded"]: + raise Exception("Subnet does not exist") + + subnetwork_n = self._get_most_recent_storage( + subtensor_state["SubnetworkN"][netuid] + ) + + neurons: list[NeuronInfo] = [ + self._neuron_subnet_exists(uid, netuid, self.block_number) for uid in range(0, subnetwork_n) + ] + + return { + "netuid": netuid, + "hotkeys": [self._convert_key(neurons[uid].hotkey) for uid in range(0, subnetwork_n)], + "coldkeys": [self._convert_key(neurons[uid].coldkey) for uid in range(0, subnetwork_n)], + "active": [neurons[uid].active for uid in range(0, subnetwork_n)], + "validator_permit": [neurons[uid].validator_permit for uid in range(0, subnetwork_n)], + "pruning_score": [neurons[uid].pruning_score for uid in range(0, subnetwork_n)], + "last_update": [neurons[uid].last_update for uid in range(0, subnetwork_n)], + "emission": [neurons[uid].emission for uid in range(0, subnetwork_n)], + "dividends": [neurons[uid].dividends for uid in range(0, subnetwork_n)], + "incentives": [neurons[uid].incentive for uid in range(0, subnetwork_n)], + "consensus": [neurons[uid].consensus for uid in range(0, subnetwork_n)], + "trust": [neurons[uid].trust for uid in range(0, subnetwork_n)], + "rank": [neurons[uid].rank for uid in range(0, subnetwork_n)], + "block_at_registration": [ + self._get_min_block_number(subtensor_state["LastUpdate"][netuid][uid]) + for uid in range(0, subnetwork_n) + ], + "alpha_stake": [neurons[uid].total_stake.tao for uid in range(0, subnetwork_n)], + "tao_stake": [neurons[uid].stake.tao for uid in range(0, subnetwork_n)], + "total_stake": [neurons[uid].total_stake.tao for uid in range(0, subnetwork_n)], + "emission_history": [ + [0] * self.block_number for _ in range(0, subnetwork_n) + ], + } + + @staticmethod + def _convert_key(key: str) -> tuple[int, int]: + """Convert key to be compatible with the metagraph pipeline.""" + return tuple(base58.b58decode(key[1:])[:-2]) diff --git a/tests/integration_tests/test_metagraph_integration.py b/tests/integration_tests/test_metagraph_integration.py index 005b1ae0c3..b9d6da4de0 100644 --- a/tests/integration_tests/test_metagraph_integration.py +++ b/tests/integration_tests/test_metagraph_integration.py @@ -6,6 +6,8 @@ import bittensor from bittensor.core.metagraph import METAGRAPH_STATE_DICT_NDARRAY_KEYS, get_save_dir from bittensor.utils.mock import MockSubtensor +from bittensor_wallet.mock import get_mock_wallet +from bittensor.utils.balance import Balance _subtensor_mock: MockSubtensor = MockSubtensor() @@ -36,6 +38,29 @@ def test_full_sync(self): ): self.metagraph.sync(lite=False, subtensor=self.sub) + def test_sync_with_stakes(self): + wallets = [get_mock_wallet() for _ in range(200)] + netuid = 1 + subtensor_mock = MockSubtensor() + subtensor_mock.setup() + subtensor_mock.create_subnet(netuid) + balance = 10000.0 + + for wallet in wallets: + subtensor_mock.force_register_neuron( + netuid=netuid, + hotkey=wallet.hotkey.ss58_address, + coldkey=wallet.coldkey.ss58_address, + stake=Balance(balance), + balance=Balance(balance), + ) + + metagraph = bittensor.Metagraph(netuid=netuid, network=subtensor_mock.network, sync=False) + metagraph.sync(subtensor=subtensor_mock) + assert len(metagraph.S) == len(wallets) + for stake in metagraph.S: + assert stake == balance + def test_sync_block_0(self): with mock.patch.object( self.sub, "get_metagraph_info", return_value=mock.MagicMock()