diff --git a/bittensor/_cli/__init__.py b/bittensor/_cli/__init__.py index 83fe2b8f26..c9e5ecedfe 100644 --- a/bittensor/_cli/__init__.py +++ b/bittensor/_cli/__init__.py @@ -18,15 +18,18 @@ # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER # DEALINGS IN THE SOFTWARE. +import argparse import os import sys -import argparse +from typing import List import bittensor -from rich.prompt import Prompt -from rich.prompt import Confirm +import torch +from rich.prompt import Confirm, Prompt from substrateinterface.utils.ss58 import ss58_decode, ss58_encode + from . import cli_impl + console = bittensor.__console__ class cli: @@ -88,6 +91,27 @@ def config() -> 'bittensor.config': help='''Set the output width of the overview. Defaults to automatic width from terminal.''', default=None, ) + overview_parser.add_argument( + '--sort_by', + '--wallet.sort_by', + dest='sort_by', + required=False, + action='store', + default="", + type=str, + help='''Sort the hotkeys by the specified column title (e.g. name, uid, axon).''' + ) + overview_parser.add_argument( + '--sort_order', + '--wallet.sort_order', + dest="sort_order", + required=False, + action='store', + default="ascending", + type=str, + help='''Sort the hotkeys in the specified ordering. (ascending/asc or descending/desc/reverse)''' + ) + bittensor.wallet.add_args( overview_parser ) bittensor.subtensor.add_args( overview_parser ) @@ -250,6 +274,7 @@ def config() -> 'bittensor.config': 'register', help='''Register a wallet to a network.''' ) + unstake_parser = cmd_parsers.add_parser( 'unstake', help='''Unstake from hotkey accounts.''' @@ -497,6 +522,7 @@ def config() -> 'bittensor.config': help='''Set true to avoid prompting the user.''', default=False, ) + bittensor.wallet.add_args( unstake_parser ) bittensor.subtensor.add_args( unstake_parser ) @@ -535,6 +561,7 @@ def config() -> 'bittensor.config': help='''Set true to avoid prompting the user.''', default=False, ) + bittensor.wallet.add_args( stake_parser ) bittensor.subtensor.add_args( stake_parser ) @@ -571,23 +598,6 @@ def config() -> 'bittensor.config': help='''Set true to avoid prompting the user.''', default=False, ) - register_parser.add_argument( - '--num_processes', - '--num', - '-n', - dest='num_processes', - help="Number of processors to use for registration", - type=int, - default=None, - ) - register_parser.add_argument( - '--update_interval', - '-u', - dest='update_interval', - help="The number of nonces to process before checking for next block during registration", - type=int, - default=None, - ) bittensor.wallet.add_args( register_parser ) bittensor.subtensor.add_args( register_parser ) @@ -709,7 +719,7 @@ def check_unstake_config( config: 'bittensor.Config' ): if config.wallet.get('all_hotkeys'): hotkeys = "all hotkeys" elif config.wallet.get('hotkeys'): - hotkeys = str(config.hotkeys).replace('[', '').replace(']', '') + hotkeys = str(config.wallet.hotkeys).replace('[', '').replace(']', '') else: hotkeys = str(config.wallet.hotkey) if not Confirm.ask("Unstake all Tao from: [bold]'{}'[/bold]?".format(hotkeys)): @@ -819,6 +829,28 @@ def check_register_config( config: 'bittensor.Config' ): hotkey = Prompt.ask("Enter hotkey name", default = bittensor.defaults.wallet.hotkey) config.wallet.hotkey = str(hotkey) + if not config.no_prompt and config.subtensor.register.cuda.use_cuda == bittensor.defaults.subtensor.register.cuda.use_cuda: + # Ask about cuda registration only if a CUDA device is available. + if torch.cuda.is_available(): + cuda = Confirm.ask("Detected CUDA device, use CUDA for registration?\n") + config.subtensor.register.cuda.use_cuda = cuda + # Only ask about which CUDA device if the user has more than one CUDA device. + if cuda and config.subtensor.register.cuda.get('dev_id') is None and torch.cuda.device_count() > 0: + devices: List[str] = [str(x) for x in range(torch.cuda.device_count())] + device_names: List[str] = [torch.cuda.get_device_name(x) for x in range(torch.cuda.device_count())] + console.print("Available CUDA devices:") + choices_str: str = "" + for i, device in enumerate(devices): + choices_str += (" {}: {}\n".format(device, device_names[i])) + console.print(choices_str) + dev_id = Prompt.ask("Which GPU would you like to use?", choices=devices, default=str(bittensor.defaults.subtensor.register.cuda.dev_id)) + try: + dev_id = int(dev_id) + except ValueError: + console.error(":cross_mark:[red]Invalid GPU device[/red] [bold white]{}[/bold white]\nAvailable CUDA devices:{}".format(dev_id, choices_str)) + sys.exit(1) + config.subtensor.register.cuda.dev_id = dev_id + def check_new_coldkey_config( config: 'bittensor.Config' ): if config.wallet.get('name') == bittensor.defaults.wallet.name and not config.no_prompt: wallet_name = Prompt.ask("Enter wallet name", default = bittensor.defaults.wallet.name) @@ -902,4 +934,4 @@ def check_help_config( config: 'bittensor.Config'): def check_update_config( config: 'bittensor.Config'): if not config.no_prompt: answer = Prompt.ask('This will update the local bittensor package', choices = ['Y','N'], default = 'Y') - config.answer = answer \ No newline at end of file + config.answer = answer diff --git a/bittensor/_cli/cli_impl.py b/bittensor/_cli/cli_impl.py index 245b943885..bdce4358dc 100644 --- a/bittensor/_cli/cli_impl.py +++ b/bittensor/_cli/cli_impl.py @@ -17,7 +17,7 @@ import os import sys -from typing import List, Union +from typing import List, Union, Optional from cachetools import Cache @@ -239,7 +239,15 @@ def register( self ): """ wallet = bittensor.wallet( config = self.config ) subtensor = bittensor.subtensor( config = self.config ) - subtensor.register( wallet = wallet, prompt = not self.config.no_prompt, num_processes = self.config.num_processes, update_interval = self.config.update_interval ) + subtensor.register( + wallet = wallet, + prompt = not self.config.no_prompt, + TPB = self.config.subtensor.register.cuda.get('TPB', None), + update_interval = self.config.subtensor.register.get('update_interval', None), + num_processes = self.config.subtensor.register.get('num_processes', None), + cuda = self.config.subtensor.register.cuda.get('use_cuda', None), + dev_id = self.config.subtensor.register.cuda.get('dev_id', None) + ) def transfer( self ): r""" Transfer token of amount to destination. @@ -251,7 +259,6 @@ def transfer( self ): def unstake( self ): r""" Unstake token of amount from hotkey(s). """ - # TODO: Implement this without re-unlocking the coldkey. config = self.config.copy() config.hotkey = None wallet = bittensor.wallet( config = self.config ) @@ -263,7 +270,7 @@ def unstake( self ): all_hotkeys: List[bittensor.wallet] = self._get_hotkey_wallets_for_wallet( wallet = wallet ) # Exclude hotkeys that are specified. wallets_to_unstake_from = [ - wallet for wallet in all_hotkeys if wallet.hotkey_str not in self.config.wallet.get('hotkeys') + wallet for wallet in all_hotkeys if wallet.hotkey_str not in self.config.wallet.get('hotkeys', []) ] elif self.config.wallet.get('hotkeys'): @@ -276,9 +283,7 @@ def unstake( self ): subtensor.unstake( wallet, amount = None if self.config.get('unstake_all') else self.config.get('amount'), wait_for_inclusion = True, prompt = not self.config.no_prompt ) return None - wallet_0: 'bittensor.wallet' = wallets_to_unstake_from[0] - # Decrypt coldkey for all wallet(s) to use - wallet_0.coldkey + final_wallets: List['bittensor.wallet'] = [] final_amounts: List[Union[float, Balance]] = [] @@ -287,9 +292,6 @@ def unstake( self ): if not wallet.is_registered(): # Skip unregistered hotkeys. continue - # Assign decrypted coldkey from wallet_0 - # so we don't have to decrypt again - wallet._coldkey = wallet_0._coldkey unstake_amount_tao: float = self.config.get('amount') if self.config.get('max_stake'): @@ -307,19 +309,17 @@ def unstake( self ): if not self.config.no_prompt: if not Confirm.ask("Do you want to unstake from the following keys:\n" + \ "".join([ - f" [bold white]- {wallet.hotkey_str}: {amount}𝜏[/bold white]\n" for wallet, amount in zip(final_wallets, final_amounts) + f" [bold white]- {wallet.hotkey_str}: {amount.tao}𝜏[/bold white]\n" for wallet, amount in zip(final_wallets, final_amounts) ]) ): return None - - for wallet, amount in zip(final_wallets, final_amounts): - subtensor.unstake( wallet, amount = None if self.config.get('unstake_all') else amount, wait_for_inclusion = True, prompt = False ) + + subtensor.unstake_multiple( wallets = final_wallets, amounts = None if self.config.get('unstake_all') else final_amounts, wait_for_inclusion = True, prompt = False ) def stake( self ): r""" Stake token of amount to hotkey(s). """ - # TODO: Implement this without re-unlocking the coldkey. config = self.config.copy() config.hotkey = None wallet = bittensor.wallet( config = config ) @@ -331,7 +331,7 @@ def stake( self ): all_hotkeys: List[bittensor.wallet] = self._get_hotkey_wallets_for_wallet( wallet = wallet ) # Exclude hotkeys that are specified. wallets_to_stake_to = [ - wallet for wallet in all_hotkeys if wallet.hotkey_str not in self.config.wallet.get('hotkeys') + wallet for wallet in all_hotkeys if wallet.hotkey_str not in self.config.wallet.get('hotkeys', []) ] elif self.config.wallet.get('hotkeys'): @@ -386,8 +386,7 @@ def stake( self ): ): return None - for wallet, amount in zip(final_wallets, final_amounts): - subtensor.add_stake( wallet, amount = None if self.config.get('stake_all') else amount, wait_for_inclusion = True, prompt = False ) + subtensor.add_stake_multiple( wallets = final_wallets, amounts = None if self.config.get('stake_all') else final_amounts, wait_for_inclusion = True, prompt = False ) def set_weights( self ): @@ -596,24 +595,14 @@ def overview(self): all_hotkeys = [] total_balance = bittensor.Balance(0) - - # We are printing for every wallet. + + # We are printing for every coldkey. if self.config.all: cold_wallets = CLI._get_coldkey_wallets_for_path(self.config.wallet.path) for cold_wallet in tqdm(cold_wallets, desc="Pulling balances"): if cold_wallet.coldkeypub_file.exists_on_device() and not cold_wallet.coldkeypub_file.is_encrypted(): total_balance = total_balance + subtensor.get_balance( cold_wallet.coldkeypub.ss58_address ) all_hotkeys = CLI._get_all_wallets_for_path( self.config.wallet.path ) - - # We are printing for a select number of hotkeys. - elif self.config.wallet.hotkeys: - # Only show hotkeys for wallets in the list - all_hotkeys = [hotkey for hotkey in all_hotkeys if hotkey.hotkey_str in self.config.wallet.hotkeys] - coldkey_wallet = bittensor.wallet( config = self.config ) - if coldkey_wallet.coldkeypub_file.exists_on_device() and not coldkey_wallet.coldkeypub_file.is_encrypted(): - total_balance = subtensor.get_balance( coldkey_wallet.coldkeypub.ss58_address ) - - # We are printing for all keys under the wallet. else: # We are only printing keys for a single coldkey coldkey_wallet = bittensor.wallet( config = self.config ) @@ -624,6 +613,16 @@ def overview(self): return all_hotkeys = CLI._get_hotkey_wallets_for_wallet( coldkey_wallet ) + # We are printing for a select number of hotkeys from all_hotkeys. + + if self.config.wallet.get('hotkeys', []): + if not self.config.get('all_hotkeys', False): + # We are only showing hotkeys that are specified. + all_hotkeys = [hotkey for hotkey in all_hotkeys if hotkey.hotkey_str in self.config.wallet.hotkeys] + else: + # We are excluding the specified hotkeys from all_hotkeys. + all_hotkeys = [hotkey for hotkey in all_hotkeys if hotkey.hotkey_str not in self.config.wallet.hotkeys] + # Check we have keys to display. if len(all_hotkeys) == 0: console.print("[red]No wallets found.[/red]") @@ -732,10 +731,10 @@ def overview(self): console.clear() - sort_by: str = self.config.wallet.sort_by - sort_order: str = self.config.wallet.sort_order + sort_by: Optional[str] = self.config.get('sort_by', None) + sort_order: Optional[str] = self.config.get('sort_order', None) - if sort_by != "": + if sort_by is not None and sort_by != "": column_to_sort_by: int = 0 highest_matching_ratio: int = 0 sort_descending: bool = False # Default sort_order to ascending diff --git a/bittensor/_dataset/dataset_impl.py b/bittensor/_dataset/dataset_impl.py index f86f0c7aff..f104b632bf 100644 --- a/bittensor/_dataset/dataset_impl.py +++ b/bittensor/_dataset/dataset_impl.py @@ -160,8 +160,7 @@ def __init__( self.build_hash_table() - if not os.path.isdir(os.path.expanduser(data_dir)): - os.makedirs(os.path.expanduser(data_dir)) + os.makedirs(os.path.expanduser(data_dir), exist_ok=True) self.data_queue = ThreadQueue( producer_target = self.reserve_multiple_data, diff --git a/bittensor/_subtensor/__init__.py b/bittensor/_subtensor/__init__.py index 9819c897ed..a994c1e79a 100644 --- a/bittensor/_subtensor/__init__.py +++ b/bittensor/_subtensor/__init__.py @@ -184,6 +184,14 @@ def add_args(cls, parser: argparse.ArgumentParser, prefix: str = None ): help='''The subtensor endpoint flag. If set, overrides the --network flag. ''') parser.add_argument('--' + prefix_str + 'subtensor._mock', action='store_true', help='To turn on subtensor mocking for testing purposes.', default=bittensor.defaults.subtensor._mock) + + parser.add_argument('--' + prefix_str + 'subtensor.register.num_processes', '-n', dest='subtensor.register.num_processes', help="Number of processors to use for registration", type=int, default=bittensor.defaults.subtensor.register.num_processes) + parser.add_argument('--' + prefix_str + 'subtensor.register.update_interval', '--' + prefix_str + 'subtensor.register.cuda.update_interval', '--' + prefix_str + 'cuda.update_interval', '-u', help="The number of nonces to process before checking for next block during registration", type=int, default=bittensor.defaults.subtensor.register.update_interval) + # registration args. Used for register and re-register and anything that calls register. + parser.add_argument( '--' + prefix_str + 'subtensor.register.cuda.use_cuda', '--' + prefix_str + 'cuda', '--' + prefix_str + 'cuda.use_cuda', default=bittensor.defaults.subtensor.register.cuda.use_cuda, help='''Set true to use CUDA.''', action='store_true', required=False ) + parser.add_argument( '--' + prefix_str + 'subtensor.register.cuda.dev_id', '--' + prefix_str + 'cuda.dev_id', type=int, default=argparse.SUPPRESS, help='''Set the CUDA device id. Goes by the order of speed. (i.e. 0 is the fastest).''', required=False ) + parser.add_argument( '--' + prefix_str + 'subtensor.register.cuda.TPB', '--' + prefix_str + 'cuda.TPB', type=int, default=bittensor.defaults.subtensor.register.cuda.TPB, help='''Set the number of Threads Per Block for CUDA.''', required=False ) + except argparse.ArgumentError: # re-parsing arguments. pass @@ -197,6 +205,15 @@ def add_defaults(cls, defaults ): defaults.subtensor.chain_endpoint = os.getenv('BT_SUBTENSOR_CHAIN_ENDPOINT') if os.getenv('BT_SUBTENSOR_CHAIN_ENDPOINT') != None else None defaults.subtensor._mock = os.getenv('BT_SUBTENSOR_MOCK') if os.getenv('BT_SUBTENSOR_MOCK') != None else False + defaults.subtensor.register = bittensor.Config() + defaults.subtensor.register.num_processes = os.getenv('BT_SUBTENSOR_REGISTER_NUM_PROCESSES') if os.getenv('BT_SUBTENSOR_REGISTER_NUM_PROCESSES') != None else None # uses processor count by default within the function + defaults.subtensor.register.update_interval = os.getenv('BT_SUBTENSOR_REGISTER_UPDATE_INTERVAL') if os.getenv('BT_SUBTENSOR_REGISTER_UPDATE_INTERVAL') != None else 50_000 + + defaults.subtensor.register.cuda = bittensor.Config() + defaults.subtensor.register.cuda.dev_id = 0 + defaults.subtensor.register.cuda.use_cuda = False + defaults.subtensor.register.cuda.TPB = 256 + @staticmethod def check_config( config: 'bittensor.Config' ): assert config.subtensor diff --git a/bittensor/_subtensor/subtensor_impl.py b/bittensor/_subtensor/subtensor_impl.py index e71fcca7fa..48d67d07e7 100644 --- a/bittensor/_subtensor/subtensor_impl.py +++ b/bittensor/_subtensor/subtensor_impl.py @@ -15,8 +15,9 @@ # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER # DEALINGS IN THE SOFTWARE. import torch -from rich.prompt import Confirm +from rich.prompt import Confirm, Prompt from typing import List, Dict, Union, Optional +from multiprocessing import Process import bittensor from tqdm import tqdm @@ -440,6 +441,9 @@ def register ( wait_for_finalization: bool = True, prompt: bool = False, max_allowed_attempts: int = 3, + cuda: bool = False, + dev_id: int = 0, + TPB: int = 256, num_processes: Optional[int] = None, update_interval: Optional[int] = None, ) -> bool: @@ -455,6 +459,18 @@ def register ( or returns false if the extrinsic fails to be finalized within the timeout. prompt (bool): If true, the call waits for confirmation from the user before proceeding. + max_allowed_attempts (int): + Maximum number of attempts to register the wallet. + cuda (bool): + If true, the wallet should be registered on the cuda device. + dev_id (int): + The cuda device id. + TPB (int): + The number of threads per block (cuda). + num_processes (int): + The number of processes to use to register. + update_interval (int): + The number of nonces to solve between updates. Returns: success (bool): flag is true if extrinsic was finalized or uncluded in the block. @@ -474,9 +490,16 @@ def register ( # Attempt rolling registration. attempts = 1 while True: + bittensor.__console__.print(":satellite: Registering...({}/{})".format(attempts, max_allowed_attempts)) # Solve latest POW. - pow_result = bittensor.utils.create_pow( self, wallet, num_processes=num_processes, update_interval=update_interval ) - with bittensor.__console__.status(":satellite: Registering...({}/{})".format(attempts,max_allowed_attempts)) as status: + if cuda: + if not torch.cuda.is_available(): + if prompt: + bittensor.__console__.error('CUDA is not available.') + return False + pow_result = bittensor.utils.create_pow( self, wallet, cuda, dev_id, TPB, num_processes=num_processes, update_interval=update_interval ) + else: + pow_result = bittensor.utils.create_pow( self, wallet, num_processes=num_processes, update_interval=update_interval) # pow failed if not pow_result: @@ -487,52 +510,64 @@ def register ( # pow successful, proceed to submit pow to chain for registration else: - #check if pow result is still valid - while pow_result['block_number'] >= self.get_current_block() - 3: - with self.substrate as substrate: - # create extrinsic call - call = substrate.compose_call( - call_module='SubtensorModule', - call_function='register', - call_params={ - 'block_number': pow_result['block_number'], - 'nonce': pow_result['nonce'], - 'work': bittensor.utils.hex_bytes_to_u8_list( pow_result['work'] ), - 'hotkey': wallet.hotkey.ss58_address, - 'coldkey': wallet.coldkeypub.ss58_address - } - ) - extrinsic = substrate.create_signed_extrinsic( call = call, keypair = wallet.hotkey ) - response = substrate.submit_extrinsic( extrinsic, wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization ) - - # We only wait here if we expect finalization. - if not wait_for_finalization and not wait_for_inclusion: - bittensor.__console__.print(":white_heavy_check_mark: [green]Sent[/green]") - return True - - # process if registration successful, try again if pow is still valid - response.process_events() - if not response.is_success: - bittensor.__console__.print(":cross_mark: [red]Failed[/red]: error:{}".format(response.error_message)) - time.sleep(1) - continue - - # Successful registration, final check for neuron and pubkey - else: - bittensor.__console__.print(":satellite: Checking Balance...") - neuron = self.neuron_for_pubkey( wallet.hotkey.ss58_address ) - bittensor.__console__.print(":white_heavy_check_mark: [green]Registered[/green]") - return True - - #Failed registration, retry pow - attempts += 1 - if attempts > max_allowed_attempts: - bittensor.__console__.print( "[red]No more attempts.[/red]" ) - return False + with bittensor.__console__.status(":satellite: Submitting POW..."): + # check if pow result is still valid + while bittensor.utils.POWNotStale(self, pow_result): + with self.substrate as substrate: + # create extrinsic call + call = substrate.compose_call( + call_module='SubtensorModule', + call_function='register', + call_params={ + 'block_number': pow_result['block_number'], + 'nonce': pow_result['nonce'], + 'work': bittensor.utils.hex_bytes_to_u8_list( pow_result['work'] ), + 'hotkey': wallet.hotkey.ss58_address, + 'coldkey': wallet.coldkeypub.ss58_address + } + ) + extrinsic = substrate.create_signed_extrinsic( call = call, keypair = wallet.hotkey ) + response = substrate.submit_extrinsic( extrinsic, wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization ) + + # We only wait here if we expect finalization. + if not wait_for_finalization and not wait_for_inclusion: + bittensor.__console__.print(":white_heavy_check_mark: [green]Sent[/green]") + return True + + # process if registration successful, try again if pow is still valid + response.process_events() + if not response.is_success: + if 'key is already registered' in response.error_message: + # Error meant that the key is already registered. + bittensor.__console__.print(":white_heavy_check_mark: [green]Already Registered[/green]") + return True + + bittensor.__console__.print(":cross_mark: [red]Failed[/red]: error:{}".format(response.error_message)) + time.sleep(0.5) + + # Successful registration, final check for neuron and pubkey + else: + bittensor.__console__.print(":satellite: Checking Balance...") + neuron = self.neuron_for_pubkey( wallet.hotkey.ss58_address ) + if not neuron.is_null: + bittensor.__console__.print(":white_heavy_check_mark: [green]Registered[/green]") + return True + else: + # neuron not found, try again + bittensor.__console__.print(":cross_mark: [red]Unknown error. Neuron not found.[/red]") + continue + else: + # Exited loop because pow is no longer valid. + bittensor.__console__.print( "[red]POW is stale.[/red]" ) + return False + if attempts < max_allowed_attempts: + #Failed registration, retry pow + attempts += 1 + bittensor.__console__.print( ":satellite: Failed registration, retrying pow ...({}/{})".format(attempts, max_allowed_attempts)) else: - status.update( ":satellite: Failed registration, retrying pow ...({}/{})".format(attempts, max_allowed_attempts)) - continue - + # Failed to register after max attempts. + bittensor.__console__.print( "[red]No more attempts.[/red]" ) + return False def serve ( self, @@ -729,6 +764,196 @@ def add_stake( return False + def add_stake_multiple ( + self, + wallets: List['bittensor.wallet'], + amounts: List[Union[Balance, float]] = None, + wait_for_inclusion: bool = True, + wait_for_finalization: bool = False, + prompt: bool = False, + ) -> bool: + r""" Adds stake to each wallet hotkey in the list, using each amount, from the common coldkey. + Args: + wallets (List[bittensor.wallet]): + List of wallets to stake. + amounts (List[Union[Balance, float]]): + List of amounts to stake. If None, stake all to the first hotkey. + wait_for_inclusion (bool): + if set, waits for the extrinsic to enter a block before returning true, + or returns false if the extrinsic fails to enter the block within the timeout. + wait_for_finalization (bool): + 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. + prompt (bool): + If true, the call waits for confirmation from the user before proceeding. + Returns: + success (bool): + flag is true if extrinsic was finalized or included in the block. + flag is true if any wallet was staked. + If we did not wait for finalization / inclusion, the response is true. + """ + if not isinstance(wallets, list): + raise TypeError("wallets must be a list of bittensor.wallet") + + if len(wallets) == 0: + return True + + if amounts is not None and len(amounts) != len(wallets): + raise ValueError("amounts must be a list of the same length as wallets") + + if amounts is not None and not all(isinstance(amount, (Balance, float)) for amount in amounts): + raise TypeError("amounts must be a [list of bittensor.Balance or float] or None") + + if amounts is None: + amounts = [None] * len(wallets) + else: + # Convert to Balance + amounts = [bittensor.Balance.from_tao(amount) if isinstance(amount, float) else amount for amount in amounts ] + + if sum(amount.tao for amount in amounts) == 0: + # Staking 0 tao + return True + + wallet_0: 'bittensor.wallet' = wallets[0] + # Decrypt coldkey for all wallet(s) to use + wallet_0.coldkey + + neurons = [] + with bittensor.__console__.status(":satellite: Syncing with chain: [white]{}[/white] ...".format(self.network)): + old_balance = self.get_balance( wallet_0.coldkey.ss58_address ) + + for wallet in wallets: + neuron = self.neuron_for_pubkey( ss58_hotkey = wallet.hotkey.ss58_address ) + + if neuron.is_null: + neurons.append( None ) + continue + + neurons.append( neuron ) + + # Remove existential balance to keep key alive. + ## Keys must maintain a balance of at least 1000 rao to stay alive. + total_staking_rao = sum([amount.rao if amount is not None else 0 for amount in amounts]) + if total_staking_rao == 0: + # Staking all to the first wallet. + if old_balance.rao > 1000: + old_balance -= bittensor.Balance.from_rao(1000) + + elif total_staking_rao < 1000: + # Staking less than 1000 rao to the wallets. + pass + else: + # Staking more than 1000 rao to the wallets. + ## Reduce the amount to stake to each wallet to keep the balance above 1000 rao. + percent_reduction = 1 - (1000 / total_staking_rao) + amounts = [amount * percent_reduction for amount in amounts] + + successful_stakes = 0 + for wallet, amount, neuron in zip(wallets, amounts, neurons): + if neuron is None: + bittensor.__console__.print(":cross_mark: [red]Hotkey: {} is not registered. Skipping ...[/red]".format( wallet.hotkey_str )) + continue + + if wallet.coldkeypub.ss58_address != wallet_0.coldkeypub.ss58_address: + bittensor.__console__.print(":cross_mark: [red]Hotkey: {} is not under the same coldkey. Skipping ...[/red]".format( wallet.hotkey_str )) + continue + + # Assign decrypted coldkey from wallet_0 + # so we don't have to decrypt again + wallet._coldkey = wallet_0._coldkey + staking_all = False + # Convert to bittensor.Balance + if amount == None: + # Stake it all. + staking_balance = bittensor.Balance.from_tao( old_balance.tao ) + staking_all = True + else: + # Amounts are cast to balance earlier in the function + assert isinstance(amount, bittensor.Balance) + staking_balance = amount + + # Estimate staking fee. + stake_fee = None # To be filled. + with bittensor.__console__.status(":satellite: Estimating Staking Fees..."): + with self.substrate as substrate: + call = substrate.compose_call( + call_module='SubtensorModule', + call_function='add_stake', + call_params={ + 'hotkey': wallet.hotkey.ss58_address, + 'ammount_staked': staking_balance.rao + } + ) + payment_info = substrate.get_payment_info(call = call, keypair = wallet.coldkey) + if payment_info: + stake_fee = bittensor.Balance.from_rao(payment_info['partialFee']) + bittensor.__console__.print("[green]Estimated Fee: {}[/green]".format( stake_fee )) + else: + stake_fee = bittensor.Balance.from_tao( 0.2 ) + bittensor.__console__.print(":cross_mark: [red]Failed[/red]: could not estimate staking fee, assuming base fee of 0.2") + + # Check enough to stake + if staking_all: + staking_balance -= stake_fee + max(staking_balance, bittensor.Balance.from_tao(0)) + + if staking_balance > old_balance - stake_fee: + bittensor.__console__.print(":cross_mark: [red]Not enough balance[/red]: [green]{}[/green] to stake: [blue]{}[/blue] from coldkey: [white]{}[/white]".format(old_balance, staking_balance, wallet.name)) + continue + + # Ask before moving on. + if prompt: + if not Confirm.ask("Do you want to stake:\n[bold white] amount: {}\n hotkey: {}\n fee: {}[/bold white ]?".format( staking_balance, wallet.hotkey_str, stake_fee) ): + continue + + with bittensor.__console__.status(":satellite: Staking to chain: [white]{}[/white] ...".format(self.network)): + with self.substrate as substrate: + call = substrate.compose_call( + call_module='SubtensorModule', + call_function='add_stake', + call_params={ + 'hotkey': wallet.hotkey.ss58_address, + 'ammount_staked': staking_balance.rao + } + ) + extrinsic = substrate.create_signed_extrinsic( call = call, keypair = wallet.coldkey ) + response = substrate.submit_extrinsic( extrinsic, wait_for_inclusion = wait_for_inclusion, wait_for_finalization = wait_for_finalization ) + # We only wait here if we expect finalization. + if not wait_for_finalization and not wait_for_inclusion: + bittensor.__console__.print(":white_heavy_check_mark: [green]Sent[/green]") + old_balance -= staking_balance + stake_fee + successful_stakes += 1 + if staking_all: + # If staked all, no need to continue + break + + continue + + response.process_events() + if response.is_success: + bittensor.__console__.print(":white_heavy_check_mark: [green]Finalized[/green]") + else: + bittensor.__console__.print(":cross_mark: [red]Failed[/red]: error:{}".format(response.error_message)) + + if response.is_success: + block = self.get_current_block() + new_stake = bittensor.Balance.from_tao( self.neuron_for_uid( uid = neuron.uid, block = block ).stake) + new_balance = self.get_balance( wallet.coldkey.ss58_address ) + bittensor.__console__.print("Stake ({}): [blue]{}[/blue] :arrow_right: [green]{}[/green]".format( neuron.uid, neuron.stake, new_stake )) + old_balance = new_balance + successful_stakes += 1 + if staking_all: + # If staked all, no need to continue + break + + if successful_stakes != 0: + with bittensor.__console__.status(":satellite: Checking Balance on: ([white]{}[/white] ...".format(self.network)): + new_balance = self.get_balance( wallet.coldkey.ss58_address ) + bittensor.__console__.print("Balance: [blue]{}[/blue] :arrow_right: [green]{}[/green]".format( old_balance, new_balance )) + return True + + return False + def transfer( self, wallet: 'bittensor.wallet', @@ -957,6 +1182,166 @@ def unstake ( return True return False + + def unstake_multiple ( + self, + wallets: List['bittensor.wallet'], + amounts: List[Union[Balance, float]] = None, + wait_for_inclusion: bool = True, + wait_for_finalization: bool = False, + prompt: bool = False, + ) -> bool: + r""" Removes stake from each wallet hotkey in the list, using each amount, to their common coldkey. + Args: + wallets (List[bittensor.wallet]): + List of wallets to unstake. + amounts (List[Union[Balance, float]]): + List of amounts to unstake. If None, unstake all. + wait_for_inclusion (bool): + if set, waits for the extrinsic to enter a block before returning true, + or returns false if the extrinsic fails to enter the block within the timeout. + wait_for_finalization (bool): + 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. + prompt (bool): + If true, the call waits for confirmation from the user before proceeding. + Returns: + success (bool): + flag is true if extrinsic was finalized or included in the block. + flag is true if any wallet was unstaked. + If we did not wait for finalization / inclusion, the response is true. + """ + if not isinstance(wallets, list): + raise TypeError("wallets must be a list of bittensor.wallet") + + if len(wallets) == 0: + return True + + if amounts is not None and len(amounts) != len(wallets): + raise ValueError("amounts must be a list of the same length as wallets") + + if amounts is not None and not all(isinstance(amount, (Balance, float)) for amount in amounts): + raise TypeError("amounts must be a [list of bittensor.Balance or float] or None") + + if amounts is None: + amounts = [None] * len(wallets) + else: + # Convert to Balance + amounts = [bittensor.Balance.from_tao(amount) if isinstance(amount, float) else amount for amount in amounts ] + + if sum(amount.tao for amount in amounts) == 0: + # Staking 0 tao + return True + + + wallet_0: 'bittensor.wallet' = wallets[0] + # Decrypt coldkey for all wallet(s) to use + wallet_0.coldkey + + neurons = [] + with bittensor.__console__.status(":satellite: Syncing with chain: [white]{}[/white] ...".format(self.network)): + old_balance = self.get_balance( wallet_0.coldkey.ss58_address ) + + for wallet in wallets: + neuron = self.neuron_for_pubkey( ss58_hotkey = wallet.hotkey.ss58_address ) + + if neuron.is_null: + neurons.append( None ) + continue + + neurons.append( neuron ) + + successful_unstakes = 0 + for wallet, amount, neuron in zip(wallets, amounts, neurons): + if neuron is None: + bittensor.__console__.print(":cross_mark: [red]Hotkey: {} is not registered. Skipping ...[/red]".format( wallet.hotkey_str )) + continue + + if wallet.coldkeypub.ss58_address != wallet_0.coldkeypub.ss58_address: + bittensor.__console__.print(":cross_mark: [red]Hotkey: {} is not under the same coldkey. Skipping ...[/red]".format( wallet.hotkey_str )) + continue + + # Assign decrypted coldkey from wallet_0 + # so we don't have to decrypt again + wallet._coldkey = wallet_0._coldkey + + # Covert to bittensor.Balance + if amount == None: + # Unstake it all. + unstaking_balance = bittensor.Balance.from_tao( neuron.stake ) + elif not isinstance(amount, bittensor.Balance ): + unstaking_balance = bittensor.Balance.from_tao( amount ) + else: + unstaking_balance = amount + + # Check enough to unstake. + stake_on_uid = bittensor.Balance.from_tao( neuron.stake ) + if unstaking_balance > stake_on_uid: + bittensor.__console__.print(":cross_mark: [red]Not enough stake[/red]: [green]{}[/green] to unstake: [blue]{}[/blue] from hotkey: [white]{}[/white]".format(stake_on_uid, unstaking_balance, wallet.hotkey_str)) + continue + + # Estimate unstaking fee. + unstake_fee = None # To be filled. + with bittensor.__console__.status(":satellite: Estimating Staking Fees..."): + with self.substrate as substrate: + call = substrate.compose_call( + call_module='SubtensorModule', + call_function='remove_stake', + call_params={ + 'hotkey': wallet.hotkey.ss58_address, + 'ammount_unstaked': unstaking_balance.rao + } + ) + payment_info = substrate.get_payment_info(call = call, keypair = wallet.coldkey) + if payment_info: + unstake_fee = bittensor.Balance.from_rao(payment_info['partialFee']) + bittensor.__console__.print("[green]Estimated Fee: {}[/green]".format( unstake_fee )) + else: + unstake_fee = bittensor.Balance.from_tao( 0.2 ) + bittensor.__console__.print(":cross_mark: [red]Failed[/red]: could not estimate staking fee, assuming base fee of 0.2") + + # Ask before moving on. + if prompt: + if not Confirm.ask("Do you want to unstake:\n[bold white] amount: {}\n hotkey: {}\n fee: {}[/bold white ]?".format( unstaking_balance, wallet.hotkey_str, unstake_fee) ): + continue + + with bittensor.__console__.status(":satellite: Unstaking from chain: [white]{}[/white] ...".format(self.network)): + with self.substrate as substrate: + call = substrate.compose_call( + call_module='SubtensorModule', + call_function='remove_stake', + call_params={ + 'hotkey': wallet.hotkey.ss58_address, + 'ammount_unstaked': unstaking_balance.rao + } + ) + extrinsic = substrate.create_signed_extrinsic( call = call, keypair = wallet.coldkey ) + response = substrate.submit_extrinsic( extrinsic, wait_for_inclusion = wait_for_inclusion, wait_for_finalization = wait_for_finalization ) + # We only wait here if we expect finalization. + if not wait_for_finalization and not wait_for_inclusion: + bittensor.__console__.print(":white_heavy_check_mark: [green]Sent[/green]") + successful_unstakes += 1 + continue + + response.process_events() + if response.is_success: + bittensor.__console__.print(":white_heavy_check_mark: [green]Finalized[/green]") + else: + bittensor.__console__.print(":cross_mark: [red]Failed[/red]: error:{}".format(response.error_message)) + + if response.is_success: + block = self.get_current_block() + new_stake = bittensor.Balance.from_tao( self.neuron_for_uid( uid = neuron.uid, block = block ).stake) + bittensor.__console__.print("Stake ({}): [blue]{}[/blue] :arrow_right: [green]{}[/green]".format( neuron.uid, stake_on_uid, new_stake )) + successful_unstakes += 1 + + if successful_unstakes != 0: + with bittensor.__console__.status(":satellite: Checking Balance on: ([white]{}[/white] ...".format(self.network)): + new_balance = self.get_balance( wallet.coldkey.ss58_address ) + bittensor.__console__.print("Balance: [blue]{}[/blue] :arrow_right: [green]{}[/green]".format( old_balance, new_balance )) + return True + + return False def set_weights( self, diff --git a/bittensor/_wallet/__init__.py b/bittensor/_wallet/__init__.py index df8fb4fdf7..9849d50664 100644 --- a/bittensor/_wallet/__init__.py +++ b/bittensor/_wallet/__init__.py @@ -111,11 +111,11 @@ def add_args(cls, parser: argparse.ArgumentParser, prefix: str = None ): parser.add_argument('--' + prefix_str + 'wallet.hotkey', required=False, default=bittensor.defaults.wallet.hotkey, help='''The name of wallet's hotkey.''') parser.add_argument('--' + prefix_str + 'wallet.path', required=False, default=bittensor.defaults.wallet.path, help='''The path to your bittensor wallets''') parser.add_argument('--' + prefix_str + 'wallet._mock', action='store_true', default=bittensor.defaults.wallet._mock, help='To turn on wallet mocking for testing purposes.') + parser.add_argument('--' + prefix_str + 'wallet.hotkeys', '--' + prefix_str + 'wallet.exclude_hotkeys', required=False, action='store', default=bittensor.defaults.wallet.hotkeys, type=str, nargs='*', help='''Specify the hotkeys by name. (e.g. hk1 hk2 hk3)''') parser.add_argument('--' + prefix_str + 'wallet.all_hotkeys', required=False, action='store_true', default=bittensor.defaults.wallet.all_hotkeys, help='''To specify all hotkeys. Specifying hotkeys will exclude them from this all.''') - parser.add_argument('--' + prefix_str + 'wallet.sort_by', required=False, action='store', default=bittensor.defaults.wallet.sort_by, type=str, help='''Sort the hotkeys by the specified column title (e.g. name, uid, axon).''') - parser.add_argument('--' + prefix_str + 'wallet.sort_order', required=False, action='store', default=bittensor.defaults.wallet.sort_order, type=str, help='''Sort the hotkeys in the specified ordering. (ascending/asc or descending/desc/reverse)''') parser.add_argument('--' + prefix_str + 'wallet.reregister', required=False, action='store', default=bittensor.defaults.wallet.reregister, type=bool, help='''Whether to reregister the wallet if it is not already registered.''') + except argparse.ArgumentError as e: import pdb #pdb.set_trace() @@ -134,8 +134,6 @@ def add_defaults(cls, defaults): # CLI defaults for Overview defaults.wallet.hotkeys = [] defaults.wallet.all_hotkeys = False - defaults.wallet.sort_by = "" - defaults.wallet.sort_order = "ascending" # Defaults for registration defaults.wallet.reregister = True @@ -148,6 +146,5 @@ def check_config(cls, config: 'bittensor.Config' ): assert isinstance(config.wallet.get('hotkey', bittensor.defaults.wallet.hotkey), str ) or config.wallet.get('hotkey', bittensor.defaults.wallet.hotkey) == None assert isinstance(config.wallet.path, str) assert isinstance(config.wallet.hotkeys, list) - assert isinstance(config.wallet.sort_by, str) - assert isinstance(config.wallet.sort_order, str) assert isinstance(config.wallet.reregister, bool) + assert isinstance(config.wallet.all_hotkeys, bool) diff --git a/bittensor/_wallet/wallet_impl.py b/bittensor/_wallet/wallet_impl.py index 158f8d211f..993b09930a 100644 --- a/bittensor/_wallet/wallet_impl.py +++ b/bittensor/_wallet/wallet_impl.py @@ -245,33 +245,76 @@ def reregister( # Check if the wallet should reregister if not self.config.wallet.get('reregister'): sys.exit(0) - return self.register(subtensor=subtensor, wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, prompt=prompt) + + subtensor.register( + wallet = self, + prompt = prompt, + TPB = self.config.subtensor.register.cuda.get('TPB', None), + update_interval = self.config.subtensor.register.cuda.get('update_interval', None), + num_processes = self.config.subtensor.register.get('num_processes', None), + cuda = self.config.subtensor.register.cuda.get('use_cuda', None), + dev_id = self.config.subtensor.register.cuda.get('dev_id', None), + wait_for_inclusion = wait_for_inclusion, + wait_for_finalization = wait_for_finalization, + ) + + return self def register ( self, subtensor: 'bittensor.Subtensor' = None, wait_for_inclusion: bool = False, wait_for_finalization: bool = True, - prompt: bool = False + prompt: bool = False, + max_allowed_attempts: int = 3, + cuda: bool = False, + dev_id: int = 0, + TPB: int = 256, + num_processes: Optional[int] = None, + update_interval: Optional[int] = None, ) -> 'bittensor.Wallet': - """ Registers this wallet on the chain. - Args: - wait_for_inclusion (bool): - if set, waits for the extrinsic to enter a block before returning true, - or returns false if the extrinsic fails to enter the block within the timeout. - wait_for_finalization (bool): - 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. - subtensor( 'bittensor.Subtensor' ): - Bittensor subtensor connection. Overrides with defaults if None. - prompt (bool): - If true, the call waits for confirmation from the user before proceeding. - Return: - This wallet. + """ Registers the wallet to chain. + Args: + subtensor( 'bittensor.Subtensor' ): + Bittensor subtensor connection. Overrides with defaults if None. + wait_for_inclusion (bool): + If set, waits for the extrinsic to enter a block before returning true, + or returns false if the extrinsic fails to enter the block within the timeout. + wait_for_finalization (bool): + 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. + prompt (bool): + If true, the call waits for confirmation from the user before proceeding. + max_allowed_attempts (int): + Maximum number of attempts to register the wallet. + cuda (bool): + If true, the wallet should be registered on the cuda device. + dev_id (int): + The cuda device id. + TPB (int): + The number of threads per block (cuda). + num_processes (int): + The number of processes to use to register. + update_interval (int): + The number of nonces to solve between updates. + Returns: + success (bool): + flag is true if extrinsic was finalized or uncluded in the block. + If we did not wait for finalization / inclusion, the response is true. """ # Get chain connection. if subtensor == None: subtensor = bittensor.subtensor() - subtensor.register( wallet = self, wait_for_inclusion = wait_for_inclusion, wait_for_finalization = wait_for_finalization, prompt=prompt ) + subtensor.register( + wallet = self, + wait_for_inclusion = wait_for_inclusion, + wait_for_finalization = wait_for_finalization, + prompt=prompt, max_allowed_attempts=max_allowed_attempts, + cuda=cuda, + dev_id=dev_id, + TPB=TPB, + num_processes=num_processes, + update_interval=update_interval + ) return self diff --git a/bittensor/utils/__init__.py b/bittensor/utils/__init__.py index a5379e5428..3a5b353b8d 100644 --- a/bittensor/utils/__init__.py +++ b/bittensor/utils/__init__.py @@ -1,4 +1,5 @@ import binascii +import datetime import hashlib import math import multiprocessing @@ -10,6 +11,7 @@ from queue import Empty from typing import Any, Dict, Optional, Tuple, Union +import backoff import bittensor import pandas import requests @@ -18,6 +20,8 @@ from substrateinterface import Keypair from substrateinterface.utils import ss58 +from .register_cuda import reset_cuda, solve_cuda + def indexed_values_to_dataframe ( prefix: Union[str, int], @@ -122,6 +126,12 @@ def millify(n: int): return '{:.0f}{}'.format(n / 10**(3 * millidx), millnames[millidx]) +def POWNotStale(subtensor: 'bittensor.Subtensor', pow_result: Dict) -> bool: + """Returns True if the POW is not stale. + This means the block the POW is solved for is within 3 blocks of the current block. + """ + return pow_result['block_number'] >= subtensor.get_current_block() - 3 + @dataclass class POWSolution: """A solution to the registration PoW problem.""" @@ -220,7 +230,7 @@ def run(self): with self.check_block: block_number = self.curr_block_num.value block_bytes = bytes(self.curr_block) - block_difficulty = int(self.curr_diff[0] >> 32 | self.curr_diff[1]) + block_difficulty = registration_diff_unpack(self.curr_diff) self.newBlockEvent.clear() # reset nonces to start from random point @@ -257,7 +267,6 @@ def solve_for_nonce_block(solver: Solver, nonce_start: int, nonce_end: int, bloc # Check if seal meets difficulty product = seal_number * difficulty if product < limit: - print(f"{solver.proc_num} found a solution: {nonce}, {block_number}, {str(block_bytes)}, {str(seal)}, {difficulty}") # Found a solution, save it. return POWSolution(nonce, block_number, difficulty, seal), time.time() - start @@ -270,6 +279,25 @@ def solve_for_nonce_block(solver: Solver, nonce_start: int, nonce_end: int, bloc return None, time.time() - start +def registration_diff_unpack(packed_diff: multiprocessing.Array) -> int: + """Unpacks the packed two 32-bit integers into one 64-bit integer. Little endian.""" + return int(packed_diff[0] << 32 | packed_diff[1]) + + +def registration_diff_pack(diff: int, packed_diff: multiprocessing.Array): + """Packs the difficulty into two 32-bit integers. Little endian.""" + packed_diff[0] = diff >> 32 + packed_diff[1] = diff & 0xFFFFFFFF # low 32 bits + + +def update_curr_block(curr_diff: multiprocessing.Array, curr_block: multiprocessing.Array, curr_block_num: multiprocessing.Value, block_number: int, block_bytes: bytes, diff: int, lock: multiprocessing.Lock): + with lock: + curr_block_num.value = block_number + for i in range(64): + curr_block[i] = block_bytes[i] + registration_diff_pack(diff, curr_diff) + + def solve_for_difficulty_fast( subtensor, wallet, num_processes: Optional[int] = None, update_interval: Optional[int] = None ) -> Optional[POWSolution]: """ Solves the POW for registration using multiprocessing. @@ -306,15 +334,7 @@ def solve_for_difficulty_fast( subtensor, wallet, num_processes: Optional[int] = curr_block = multiprocessing.Array('h', 64, lock=True) # byte array curr_block_num = multiprocessing.Value('i', 0, lock=True) # int curr_diff = multiprocessing.Array('Q', [0, 0], lock=True) # [high, low] - - def update_curr_block(block_number: int, block_bytes: bytes, diff: int, lock: multiprocessing.Lock): - with lock: - curr_block_num.value = block_number - for i in range(64): - curr_block[i] = block_bytes[i] - curr_diff[0] = diff >> 32 - curr_diff[1] = diff & 0xFFFFFFFF # low 32 bits - + status.start() # Establish communication queues @@ -339,7 +359,7 @@ def update_curr_block(block_number: int, block_bytes: bytes, diff: int, lock: mu block_bytes = block_hash.encode('utf-8')[2:] old_block_number = block_number # Set to current block - update_curr_block(block_number, block_bytes, difficulty, check_block) + update_curr_block(curr_diff, curr_block, curr_block_num, block_number, block_bytes, difficulty, check_block) # Set new block events for each solver to start for w in solvers: @@ -373,7 +393,7 @@ def update_curr_block(block_number: int, block_bytes: bytes, diff: int, lock: mu block_bytes = block_hash.encode('utf-8')[2:] difficulty = subtensor.difficulty - update_curr_block(block_number, block_bytes, difficulty, check_block) + update_curr_block(curr_diff, curr_block, curr_block_num, block_number, block_bytes, difficulty, check_block) # Set new block events for each solver for w in solvers: w.newBlockEvent.set() @@ -405,7 +425,7 @@ def update_curr_block(block_number: int, block_bytes: bytes, diff: int, lock: mu except Empty: break - + message = f"""Solving time spent: {time.time() - start_time} Difficulty: [bold white]{millify(difficulty)}[/bold white] @@ -413,23 +433,136 @@ def update_curr_block(block_number: int, block_bytes: bytes, diff: int, lock: mu Block: [bold white]{block_number}[/bold white] Block_hash: [bold white]{block_hash.encode('utf-8')}[/bold white] Best: [bold white]{binascii.hexlify(bytes(best_seal) if best_seal else bytes(0))}[/bold white]""" - status.update(message.replace(" ", "")) - + status.update(message.replace(" ", "")) + # exited while, solution contains the nonce or wallet is registered stopEvent.set() # stop all other processes status.stop() return solution + +def get_human_readable(num, suffix="H"): + for unit in ["", "K", "M", "G", "T", "P", "E", "Z"]: + if abs(num) < 1000.0: + return f"{num:3.1f}{unit}{suffix}" + num /= 1000.0 + return f"{num:.1f}Y{suffix}" + +def millify(n: int): + millnames = ['',' K',' M',' B',' T'] + n = float(n) + millidx = max(0,min(len(millnames)-1, + int(math.floor(0 if n == 0 else math.log10(abs(n))/3)))) + + return '{:.0f}{}'.format(n / 10**(3 * millidx), millnames[millidx]) +@backoff.on_exception(backoff.constant, + Exception, + interval=1, + max_tries=3) +def get_block_with_retry(subtensor: 'bittensor.Subtensor') -> Tuple[int, int, bytes]: + block_number = subtensor.get_current_block() + difficulty = subtensor.difficulty + block_hash = subtensor.substrate.get_block_hash( block_number ) + if block_hash is None: + raise Exception("Network error. Could not connect to substrate to get block hash") + return block_number, difficulty, block_hash + +def solve_for_difficulty_fast_cuda( subtensor: 'bittensor.Subtensor', wallet: 'bittensor.Wallet', update_interval: int = 50_000, TPB: int = 512, dev_id: int = 0 ) -> Optional[POWSolution]: + """ + Solves the registration fast using CUDA + Args: + subtensor: bittensor.Subtensor + The subtensor node to grab blocks + wallet: bittensor.Wallet + The wallet to register + update_interval: int + The number of nonces to try before checking for more blocks + TPB: int + The number of threads per block. CUDA param that should match the GPU capability + dev_id: int + The CUDA device ID to execute the registration on + """ + if not torch.cuda.is_available(): + raise Exception("CUDA not available") + + if update_interval is None: + update_interval = 50_000 + + block_number, difficulty, block_hash = get_block_with_retry(subtensor) + block_bytes = block_hash.encode('utf-8')[2:] + + nonce = 0 + limit = int(math.pow(2,256)) - 1 + start_time = time.time() + + console = bittensor.__console__ + status = console.status("Solving") + + solution = -1 + start_time = time.time() + interval_time = start_time + + status.start() + while solution == -1 and not wallet.is_registered(subtensor): + solution, seal = solve_cuda(nonce, + update_interval, + TPB, + block_bytes, + block_number, + difficulty, + limit, + dev_id) + + if (solution != -1): + # Attempt to reset CUDA device + reset_cuda() + status.stop() + new_bn = subtensor.get_current_block() + print(f"Found solution for bn: {block_number}; Newest: {new_bn}") + return POWSolution(solution, block_number, difficulty, seal) + + nonce += (TPB * update_interval) + if (nonce >= int(math.pow(2,63))): + nonce = 0 + itrs_per_sec = (TPB * update_interval) / (time.time() - interval_time) + interval_time = time.time() + + block_number, difficulty, block_hash = get_block_with_retry(subtensor) + block_bytes = block_hash.encode('utf-8')[2:] + + message = f"""Solving + time spent: {datetime.timedelta(seconds=time.time() - start_time)} + Nonce: [bold white]{nonce}[/bold white] + Difficulty: [bold white]{millify(difficulty)}[/bold white] + Iters: [bold white]{get_human_readable(int(itrs_per_sec), "H")}/s[/bold white] + Block: [bold white]{block_number}[/bold white] + Block_hash: [bold white]{block_hash.encode('utf-8')}[/bold white]""" + status.update(message.replace(" ", "")) + + # exited while, found_solution contains the nonce or wallet is registered + if solution == -1: # didn't find solution + reset_cuda() + status.stop() + return None + + else: + reset_cuda() + # Shouldn't get here + status.stop() + return None + +def create_pow( subtensor, wallet, cuda: bool = False, dev_id: int = 0, tpb: int = 256, num_processes: int = None, update_interval: int = None ) -> Optional[Dict[str, Any]]: + if cuda: + solution: POWSolution = solve_for_difficulty_fast_cuda( subtensor, wallet, dev_id=dev_id, TPB=tpb, update_interval=update_interval ) + else: + solution: POWSolution = solve_for_difficulty_fast( subtensor, wallet, num_processes=num_processes, update_interval=update_interval ) -def create_pow( subtensor, wallet, num_processes: Optional[int] = None, update_interval: Optional[int] = None ) -> Optional[Dict[str, Any]]: - solution: POWSolution = solve_for_difficulty_fast( subtensor, wallet, num_processes=num_processes, update_interval=update_interval ) - nonce, block_number, difficulty, seal = solution.nonce, solution.block_number, solution.difficulty, solution.seal - return None if nonce is None else { - 'nonce': nonce, - 'difficulty': difficulty, - 'block_number': block_number, - 'work': binascii.hexlify(seal) + return None if solution is None else { + 'nonce': solution.nonce, + 'difficulty': solution.difficulty, + 'block_number': solution.block_number, + 'work': binascii.hexlify(solution.seal) } def version_checking(): diff --git a/bittensor/utils/register_cuda.py b/bittensor/utils/register_cuda.py new file mode 100644 index 0000000000..f64f4777b4 --- /dev/null +++ b/bittensor/utils/register_cuda.py @@ -0,0 +1,87 @@ +import binascii +import hashlib +import math +from typing import Tuple + +import numpy as np +from Crypto.Hash import keccak + + +def solve_cuda(nonce_start: np.int64, update_interval: np.int64, TPB: int, block_bytes: bytes, bn: int, difficulty: int, limit: int, dev_id: int = 0) -> Tuple[np.int64, bytes]: + """ + Solves the PoW problem using CUDA. + Args: + nonce_start: int64 + Starting nonce. + update_interval: int64 + Number of nonces to solve before updating block information. + TPB: int + Threads per block. + block_bytes: bytes + Bytes of the block hash. 64 bytes. + difficulty: int256 + Difficulty of the PoW problem. + limit: int256 + Upper limit of the nonce. + dev_id: int (default=0) + The CUDA device ID + Returns: + Tuple[int64, bytes] + Tuple of the nonce and the seal corresponding to the solution. + Returns -1 for nonce if no solution is found. + """ + + try: + import cubit + except ImportError: + raise ImportError("Please install cubit") + + + upper = int(limit // difficulty) + + upper_bytes = upper.to_bytes(32, byteorder='little', signed=False) + + def seal_meets_difficulty( seal:bytes, difficulty:int ): + seal_number = int.from_bytes(seal, "big") + product = seal_number * difficulty + limit = int(math.pow(2,256))- 1 + + return product < limit + + def hex_bytes_to_u8_list( hex_bytes: bytes ): + hex_chunks = [int(hex_bytes[i:i+2], 16) for i in range(0, len(hex_bytes), 2)] + return hex_chunks + + def create_seal_hash( block_bytes:bytes, nonce:int ) -> bytes: + nonce_bytes = binascii.hexlify(nonce.to_bytes(8, 'little')) + pre_seal = nonce_bytes + block_bytes + seal_sh256 = hashlib.sha256( bytearray(hex_bytes_to_u8_list(pre_seal)) ).digest() + kec = keccak.new(digest_bits=256) + seal = kec.update( seal_sh256 ).digest() + return seal + + # Call cython function + # int blockSize, uint64 nonce_start, uint64 update_interval, const unsigned char[:] limit, + # const unsigned char[:] block_bytes, int dev_id + solution = cubit.solve_cuda(TPB, nonce_start, update_interval, upper_bytes, block_bytes, dev_id) # 0 is first GPU + seal = None + if solution != -1: + print(f"Checking solution: {solution} for bn: {bn}") + seal = create_seal_hash(block_bytes, solution) + if seal_meets_difficulty(seal, difficulty): + return solution, seal + else: + return -1, b'\x00' * 32 + + return solution, seal + +def reset_cuda(): + """ + Resets the CUDA environment. + """ + try: + import cubit + except ImportError: + raise ImportError("Please install cubit") + + cubit.reset_cuda() diff --git a/requirements.txt b/requirements.txt index 1e2367e28e..1c7f9b8fd6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,12 +1,13 @@ argparse base58>=2.0.1 +backoff>=2.1.0 certifi>=2020.11.8 cryptography>=3.1.1 idna>=2.10 jinja2>=3.0 fuzzywuzzy==0.18.0 +google-api-python-client>=2.6.0 python-levenshtein==0.12.1 -google-api-python-client grpcio==1.42.0 grpcio-tools==1.42.0 hypothesis>=6.47.4 diff --git a/setup.py b/setup.py index 4168052160..d4a9723bcd 100644 --- a/setup.py +++ b/setup.py @@ -69,4 +69,7 @@ 'Topic :: Software Development :: Libraries :: Python Modules', ], python_requires='>=3.7', + extras_requires={ + 'cubit': ['cubit>=1.0.5 @ git+https://github.com/opentensor/cubit.git'] + } ) diff --git a/tests/integration_tests/test_cli.py b/tests/integration_tests/test_cli.py index 77269d366b..d8276422e0 100644 --- a/tests/integration_tests/test_cli.py +++ b/tests/integration_tests/test_cli.py @@ -442,7 +442,7 @@ def test_unstake_with_specific_hotkeys( self ): with patch('bittensor.wallet') as mock_create_wallet: mock_create_wallet.side_effect = mock_wallets - with patch('bittensor.Subtensor.unstake', return_value=True) as mock_unstake: + with patch('bittensor.Subtensor.unstake_multiple', return_value=True) as mock_unstake: cli.run() mock_create_wallet.assert_has_calls( [ @@ -451,7 +451,7 @@ def test_unstake_with_specific_hotkeys( self ): any_order=True ) mock_unstake.assert_has_calls( - [call(mock_wallets[i+1], amount=5.0, wait_for_inclusion=True, prompt=False) for i in range(len(mock_wallets[1:]))], + [call(wallets=mock_wallets[1:], amounts=[5.0]*len(mock_wallets[1:]), wait_for_inclusion=True, prompt=False)], any_order = True ) @@ -500,11 +500,11 @@ def test_unstake_with_all_hotkeys( self ): with patch.object(cli, '_get_hotkey_wallets_for_wallet') as mock_get_all_wallets: mock_get_all_wallets.return_value = mock_wallets - with patch('bittensor.Subtensor.unstake', return_value=True) as mock_unstake: + with patch('bittensor.Subtensor.unstake_multiple', return_value=True) as mock_unstake: cli.run() mock_get_all_wallets.assert_called_once() mock_unstake.assert_has_calls( - [call(mock_wallets[i+1], amount=5.0, wait_for_inclusion=True, prompt=False) for i in range(len(mock_wallets[1:]))], + [call(wallets=mock_wallets, amounts=[5.0]*len(mock_wallets), wait_for_inclusion=True, prompt=False)], any_order = True ) @@ -553,13 +553,11 @@ def test_unstake_with_exclude_hotkeys_from_all( self ): with patch.object(cli, '_get_hotkey_wallets_for_wallet') as mock_get_all_wallets: mock_get_all_wallets.return_value = mock_wallets - with patch('bittensor.Subtensor.unstake', return_value=True) as mock_unstake: + with patch('bittensor.Subtensor.unstake_multiple', return_value=True) as mock_unstake: cli.run() mock_get_all_wallets.assert_called_once() mock_unstake.assert_has_calls( - [call(mock_wallets[i], amount=5.0, wait_for_inclusion=True, prompt=False) - for i in (0, 2) # Don't call for hk1 - ], + [call(wallets=[mock_wallets[0], mock_wallets[2]], amounts=[5.0, 5.0], wait_for_inclusion=True, prompt=False)], any_order = True ) @@ -628,7 +626,7 @@ def test_unstake_with_multiple_hotkeys_max_stake( self ): with patch('bittensor.wallet') as mock_create_wallet: mock_create_wallet.side_effect = mock_wallets - with patch('bittensor.Subtensor.unstake', return_value=True) as mock_unstake: + with patch('bittensor.Subtensor.unstake_multiple', return_value=True) as mock_unstake: cli.run() mock_create_wallet.assert_has_calls( [ @@ -637,7 +635,7 @@ def test_unstake_with_multiple_hotkeys_max_stake( self ): any_order=True ) mock_unstake.assert_has_calls( - [call(mock_wallet, amount=CLOSE_IN_VALUE((mock_stakes[mock_wallet.hotkey_str] - config.max_stake).tao, 0.001), wait_for_inclusion=True, prompt=False) for mock_wallet in mock_wallets[1:]], + [call(wallets=mock_wallets[1:], amounts=[CLOSE_IN_VALUE((mock_stakes[mock_wallet.hotkey_str] - config.max_stake).tao, 0.001) for mock_wallet in mock_wallets[1:]], wait_for_inclusion=True, prompt=False)], any_order = True ) @@ -706,7 +704,7 @@ def test_unstake_with_multiple_hotkeys_max_stake_not_enough_stake( self ): with patch('bittensor.wallet') as mock_create_wallet: mock_create_wallet.side_effect = mock_wallets - with patch('bittensor.Subtensor.unstake', return_value=True) as mock_unstake: + with patch('bittensor.Subtensor.unstake_multiple', return_value=True) as mock_unstake: cli.run() mock_create_wallet.assert_has_calls( [ @@ -715,15 +713,16 @@ def test_unstake_with_multiple_hotkeys_max_stake_not_enough_stake( self ): any_order=True ) mock_unstake.assert_called() - for mock_call in mock_unstake.mock_calls: - # Python 3.7 - ## https://docs.python.org/3.7/library/unittest.mock.html#call - ## Uses the 1st index as args list - ## call.args only works in Python 3.8+ - mock_wallet = mock_call[1][0] - # We shouldn't unstake from hk1 as it has less than max_stake staked - assert mock_wallet.hotkey_str != 'hk1' + # Python 3.7 + ## https://docs.python.org/3.7/library/unittest.mock.html#call + ## Uses the 1st index as args list + ## call.args only works in Python 3.8+ + mock_wallets_ = mock_unstake.mock_calls[0][2]['wallets'] + + + # We shouldn't unstake from hk1 as it has less than max_stake staked + assert all(mock_wallet.hotkey_str != 'hk1' for mock_wallet in mock_wallets_) def test_stake_with_specific_hotkeys( self ): bittensor.subtensor.register = MagicMock(return_value = True) @@ -782,7 +781,7 @@ def test_stake_with_specific_hotkeys( self ): with patch('bittensor.wallet') as mock_create_wallet: mock_create_wallet.side_effect = mock_wallets - with patch('bittensor.Subtensor.add_stake', return_value=True) as mock_add_stake: + with patch('bittensor.Subtensor.add_stake_multiple', return_value=True) as mock_add_stake: cli.run() mock_create_wallet.assert_has_calls( [ @@ -791,7 +790,7 @@ def test_stake_with_specific_hotkeys( self ): any_order=True ) mock_add_stake.assert_has_calls( - [call(mock_wallets[i+1], amount=5.0, wait_for_inclusion=True, prompt=False) for i in range(len(mock_wallets[1:]))], + [call(wallets=mock_wallets[1:], amounts=[5.0] * len(mock_wallets[1:]), wait_for_inclusion=True, prompt=False)], any_order = True ) @@ -826,7 +825,6 @@ def test_stake_with_all_hotkeys( self ): ) for hk in mock_hotkeys ] - # The 0th wallet is created twice during unstake mock_wallets[0]._coldkey = mock_coldkey mock_wallets[0].coldkey = MagicMock( return_value=mock_coldkey @@ -839,11 +837,11 @@ def test_stake_with_all_hotkeys( self ): with patch.object(cli, '_get_hotkey_wallets_for_wallet') as mock_get_all_wallets: mock_get_all_wallets.return_value = mock_wallets - with patch('bittensor.Subtensor.add_stake', return_value=True) as mock_add_stake: + with patch('bittensor.Subtensor.add_stake_multiple', return_value=True) as mock_add_stake: cli.run() mock_get_all_wallets.assert_called_once() mock_add_stake.assert_has_calls( - [call(mock_wallets[i+1], amount=5.0, wait_for_inclusion=True, prompt=False) for i in range(len(mock_wallets[1:]))], + [call(wallets=mock_wallets, amounts=[5.0] * len(mock_wallets), wait_for_inclusion=True, prompt=False)], any_order = True ) @@ -891,13 +889,11 @@ def test_stake_with_exclude_hotkeys_from_all( self ): with patch.object(cli, '_get_hotkey_wallets_for_wallet') as mock_get_all_wallets: mock_get_all_wallets.return_value = mock_wallets - with patch('bittensor.Subtensor.add_stake', return_value=True) as mock_add_stake: + with patch('bittensor.Subtensor.add_stake_multiple', return_value=True) as mock_add_stake: cli.run() mock_get_all_wallets.assert_called_once() mock_add_stake.assert_has_calls( - [call(mock_wallets[i], amount=5.0, wait_for_inclusion=True, prompt=False) - for i in (0, 2) # Don't call stake for hk1 - ], + [call(wallets=[mock_wallets[0], mock_wallets[2]], amounts=[5.0, 5.0], wait_for_inclusion=True, prompt=False)], any_order = True ) @@ -971,7 +967,7 @@ def test_stake_with_multiple_hotkeys_max_stake( self ): with patch('bittensor.wallet') as mock_create_wallet: mock_create_wallet.side_effect = mock_wallets - with patch('bittensor.Subtensor.add_stake', return_value=True) as mock_add_stake: + with patch('bittensor.Subtensor.add_stake_multiple', return_value=True) as mock_add_stake: cli.run() mock_create_wallet.assert_has_calls( [ @@ -980,7 +976,7 @@ def test_stake_with_multiple_hotkeys_max_stake( self ): any_order=True ) mock_add_stake.assert_has_calls( - [call(mock_wallet, amount=CLOSE_IN_VALUE((config.max_stake - mock_stakes[mock_wallet.hotkey_str].tao), 0.001), wait_for_inclusion=True, prompt=False) for mock_wallet in mock_wallets[1:]], + [call(wallets=mock_wallets[1:], amounts=[CLOSE_IN_VALUE((config.max_stake - mock_stakes[mock_wallet.hotkey_str].tao), 0.001) for mock_wallet in mock_wallets[1:]], wait_for_inclusion=True, prompt=False)], any_order = True ) @@ -1054,7 +1050,7 @@ def test_stake_with_multiple_hotkeys_max_stake_not_enough_balance( self ): with patch('bittensor.wallet') as mock_create_wallet: mock_create_wallet.side_effect = mock_wallets - with patch('bittensor.Subtensor.add_stake', return_value=True) as mock_add_stake: + with patch('bittensor.Subtensor.add_stake_multiple', return_value=True) as mock_add_stake: cli.run() mock_create_wallet.assert_has_calls( [ @@ -1066,14 +1062,15 @@ def test_stake_with_multiple_hotkeys_max_stake_not_enough_balance( self ): mock_add_stake.assert_called_once() total_staked = 0.0 - for mock_call in mock_add_stake.mock_calls: - # Python 3.7 - ## https://docs.python.org/3.7/library/unittest.mock.html#call - ## Uses the 2nd index as kwargs dict - ## call.kwargs only works in Python 3.8+ - staked_this_call = mock_call[2]['amount'] - - total_staked += staked_this_call + + # Python 3.7 + ## https://docs.python.org/3.7/library/unittest.mock.html#call + ## Uses the 2nd index as kwargs dict + ## call.kwargs only works in Python 3.8+ + amounts_passed = mock_add_stake.mock_calls[0][2]['amounts'] + + total_staked = sum(amounts_passed) + # We should not try to stake more than the mock_balance assert CLOSE_IN_VALUE(total_staked, 0.001) == mock_balance.tao @@ -1082,8 +1079,8 @@ def test_register( self ): config = self.config config.subtensor._mock = True config.command = "register" - config.num_processes = 1 - config.update_interval = 50_000 + config.subtensor.register.num_processes = 1 + config.subtensor.register.update_interval = 50_000 config.subtensor.network = "mock" config.no_prompt = True @@ -1265,7 +1262,7 @@ def test_inspect( self ): def test_list( self ): # Mock IO for wallet - with patch('bittensor.wallet.__new__', side_effect=[MagicMock( + with patch('bittensor.wallet', side_effect=[MagicMock( coldkeypub_file=MagicMock( exists_on_device=MagicMock( return_value=True # Wallet exists diff --git a/tests/integration_tests/test_subtensor.py b/tests/integration_tests/test_subtensor.py index 2d411f1fc0..8c1ae1967e 100644 --- a/tests/integration_tests/test_subtensor.py +++ b/tests/integration_tests/test_subtensor.py @@ -15,17 +15,17 @@ # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER # DEALINGS IN THE SOFTWARE. -from typing import DefaultDict -from unittest import mock + +import multiprocessing from unittest.mock import patch import bittensor import pytest -import psutil import unittest import time import random from unittest.mock import MagicMock from bittensor.utils.balance import Balance +from bittensor.utils import Solver, update_curr_block from substrateinterface import Keypair from bittensor._subtensor.subtensor_mock import mock_subtensor class TestSubtensor(unittest.TestCase): @@ -402,20 +402,29 @@ def process_events(self): mock_neuron.is_null = True with patch('bittensor.Subtensor.difficulty'): + # patch solution queue to return None + with patch('multiprocessing.queues.Queue.get', return_value=None) as mock_queue_get: + # patch time queue size check + with patch('multiprocessing.queues.Queue.qsize', return_value=0): - wallet = bittensor.wallet(_mock=True) - wallet.is_registered = MagicMock( side_effect=is_registered_return_values ) + wallet = bittensor.wallet(_mock=True) + wallet.is_registered = MagicMock( side_effect=is_registered_return_values ) - self.subtensor.difficulty= 1 - self.subtensor.neuron_for_pubkey = MagicMock( return_value=mock_neuron ) - self.subtensor.substrate.submit_extrinsic = MagicMock(return_value = success()) + self.subtensor.difficulty= 1 + self.subtensor.neuron_for_pubkey = MagicMock( return_value=mock_neuron ) + self.subtensor.substrate.submit_extrinsic = MagicMock(return_value = success()) - # should return True - assert self.subtensor.register(wallet=wallet,) - # calls until True and once again before exiting subtensor class - # This assertion is currently broken when difficulty is too low - #assert wallet.is_registered.call_count == workblocks_before_is_registered + 2 + with patch('bittensor.__console__.status') as mock_set_status: + # Need to patch the console status to avoid opening a parallel live display + mock_set_status.__enter__ = MagicMock(return_value=True) + mock_set_status.__exit__ = MagicMock(return_value=True) + + # should return True + assert self.subtensor.register(wallet=wallet, num_processes=3, update_interval=5 ) == True + # calls until True and once again before exiting subtensor class + # This assertion is currently broken when difficulty is too low + assert wallet.is_registered.call_count == workblocks_before_is_registered + 2 def test_registration_partly_failed( self ): class failed(): @@ -431,24 +440,29 @@ def __init__(self): def process_events(self): return True - is_registered_return_values = [False for _ in range(100)] - submit_extrinsic = [failed(), failed(), success()] + submit_extrinsic_mock = MagicMock( side_effect = [failed(), failed(), success()]) + + def is_registered_side_effect(*args, **kwargs): + nonlocal submit_extrinsic_mock + return submit_extrinsic_mock.call_count < 3 + current_block = [i for i in range(0,100)] - mock_neuron = MagicMock() - mock_neuron.is_null = True with patch('bittensor.Subtensor.difficulty'): wallet = bittensor.wallet(_mock=True) - wallet.is_registered = MagicMock( side_effect=is_registered_return_values ) + wallet.is_registered = MagicMock( side_effect=is_registered_side_effect ) - - self.subtensor.difficulty= 1 + self.subtensor.difficulty = 1 self.subtensor.get_current_block = MagicMock(side_effect=current_block) - self.subtensor.neuron_for_pubkey = MagicMock( return_value=mock_neuron ) - self.subtensor.substrate.submit_extrinsic = MagicMock(side_effect = submit_extrinsic) + self.subtensor.substrate.submit_extrinsic = submit_extrinsic_mock - # should return True - assert self.subtensor.register(wallet=wallet,) == True + with patch('bittensor.__console__.status') as mock_set_status: + # Need to patch the console status to avoid opening a parallel live display + mock_set_status.__enter__ = MagicMock(return_value=True) + mock_set_status.__exit__ = MagicMock(return_value=True) + + # should return True + assert self.subtensor.register(wallet=wallet, num_processes=3, update_interval=5) == True def test_registration_failed( self ): class failed(): diff --git a/tests/unit_tests/bittensor_tests/test_neuron.py b/tests/unit_tests/bittensor_tests/test_neuron.py index 02ee588d22..9fa79e7768 100644 --- a/tests/unit_tests/bittensor_tests/test_neuron.py +++ b/tests/unit_tests/bittensor_tests/test_neuron.py @@ -1,3 +1,4 @@ +from atexit import register from types import SimpleNamespace from unittest.mock import MagicMock, patch from more_itertools import side_effect @@ -67,6 +68,12 @@ def test_coreserver_reregister_flag_false_exit(): config.wallet = bittensor.Config() config.wallet.reregister = False # don't reregister the wallet + config.subtensor = bittensor.Config() + config.subtensor.register = bittensor.Config() + config.subtensor.register.cuda = bittensor.Config() + config.subtensor.register.cuda.use_cuda = False # don't use cuda on test + # No need to specify the other config options as they are default to None + mock_wallet = bittensor.wallet.mock() mock_wallet.config = config @@ -85,15 +92,15 @@ def exit_early(*args, **kwargs): metagraph=MagicMock(), spec=bittensor.neurons.core_server.neuron, subtensor=MagicMock( - network="mock" + network="mock", + register=mock_register ), config=config, ) - with patch.multiple( - 'bittensor.Wallet', - register=mock_register, - is_registered=MagicMock(return_value=False), # mock the wallet as not registered + with patch.object( + mock_wallet, + 'is_registered', MagicMock(return_value=False), # mock the wallet as not registered ): # Should exit without calling register @@ -113,6 +120,12 @@ def test_coreserver_reregister_flag_true(): config.wallet = bittensor.Config() config.wallet.reregister = True # try to reregister the wallet + config.subtensor = bittensor.Config() + config.subtensor.register = bittensor.Config() + config.subtensor.register.cuda = bittensor.Config() + config.subtensor.register.cuda.use_cuda = False # don't use cuda on test + # No need to specify the other config options as they are default to None + mock_wallet = bittensor.wallet.mock() mock_wallet.config = config @@ -131,17 +144,16 @@ def exit_early(*args, **kwargs): metagraph=MagicMock(), spec=bittensor.neurons.core_server.neuron, subtensor=MagicMock( - network="mock" + network="mock", + register=mock_register, ), config=config, ) - with patch.multiple( - 'bittensor.Wallet', - register=mock_register, - is_registered=MagicMock(return_value=False), # mock the wallet as not registered + with patch.object( + mock_wallet, + 'is_registered', MagicMock(return_value=False), # mock the wallet as not registered ): - # Should not exit with pytest.raises(MockException): # Should raise MockException @@ -157,6 +169,12 @@ def test_corevalidator_reregister_flag_false_exit(): config.wallet = bittensor.Config() config.wallet.reregister = False # don't reregister the wallet + config.subtensor = bittensor.Config() + config.subtensor.register = bittensor.Config() + config.subtensor.register.cuda = bittensor.Config() + config.subtensor.register.cuda.use_cuda = False # don't use cuda on test + # No need to specify the other config options as they are default to None + mock_wallet = bittensor.wallet.mock() mock_wallet.config = config @@ -172,17 +190,17 @@ def exit_early(*args, **kwargs): wallet=mock_wallet, spec=bittensor.neurons.core_validator.neuron, subtensor=MagicMock( - network="mock" + network="mock", + register=mock_register, ), config=config, ) - with patch.multiple( - 'bittensor.Wallet', - register=mock_register, - is_registered=MagicMock(return_value=False), # mock the wallet as not registered + with patch.object( + mock_wallet, + 'is_registered', MagicMock(return_value=False), # mock the wallet as not registered ): - + # Should exit without calling register with pytest.raises(SystemExit) as pytest_wrapped_e: # Should not raise MockException @@ -200,6 +218,12 @@ def test_corevalidator_reregister_flag_true(): config.wallet = bittensor.Config() config.wallet.reregister = True # try to reregister the wallet + config.subtensor = bittensor.Config() + config.subtensor.register = bittensor.Config() + config.subtensor.register.cuda = bittensor.Config() + config.subtensor.register.cuda.use_cuda = False # don't use cuda on test + # No need to specify the other config options as they are default to None + mock_wallet = bittensor.wallet.mock() mock_wallet.config = config @@ -215,15 +239,15 @@ def exit_early(*args, **kwargs): wallet=mock_wallet, spec=bittensor.neurons.core_validator.neuron, subtensor=MagicMock( - network="mock" + network="mock", + register=mock_register, ), config=config, ) - with patch.multiple( - 'bittensor.Wallet', - register=mock_register, - is_registered=MagicMock(return_value=False), # mock the wallet as not registered + with patch.object( + mock_wallet, + 'is_registered', MagicMock(return_value=False), # mock the wallet as not registered ): # Should not exit diff --git a/tests/unit_tests/bittensor_tests/utils/test_utils.py b/tests/unit_tests/bittensor_tests/utils/test_utils.py index 03065ccdca..5d0643bc08 100644 --- a/tests/unit_tests/bittensor_tests/utils/test_utils.py +++ b/tests/unit_tests/bittensor_tests/utils/test_utils.py @@ -1,5 +1,6 @@ import binascii import hashlib +import unittest import bittensor import sys import subprocess @@ -8,13 +9,14 @@ import os import random import torch +import multiprocessing from sys import platform from substrateinterface.base import Keypair from _pytest.fixtures import fixture from loguru import logger -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch @@ -233,5 +235,117 @@ def test_is_valid_ed25519_pubkey(): assert bittensor.utils.is_valid_ed25519_pubkey(good_pubkey) assert not bittensor.utils.is_valid_ed25519_pubkey(bad_pubkey) +def test_registration_diff_pack_unpack_under_32_bits(): + fake_diff = pow(2, 31)# this is under 32 bits + + mock_diff = multiprocessing.Array('Q', [0, 0], lock=True) # [high, low] + + bittensor.utils.registration_diff_pack(fake_diff, mock_diff) + assert bittensor.utils.registration_diff_unpack(mock_diff) == fake_diff + +def test_registration_diff_pack_unpack_over_32_bits(): + mock_diff = multiprocessing.Array('Q', [0, 0], lock=True) # [high, low] + fake_diff = pow(2, 32) * pow(2, 4) # this should be too large if the bit shift is wrong (32 + 4 bits) + + bittensor.utils.registration_diff_pack(fake_diff, mock_diff) + assert bittensor.utils.registration_diff_unpack(mock_diff) == fake_diff + +class TestGetBlockWithRetry(unittest.TestCase): + def test_get_block_with_retry_network_error_exit(self): + mock_subtensor = MagicMock( + get_current_block=MagicMock(return_value=1), + difficulty=1, + substrate=MagicMock( + get_block_hash=MagicMock(side_effect=Exception('network error')) + ) + ) + with pytest.raises(Exception): + # this should raise an exception because the network error is retried only 3 times + bittensor.utils.get_block_with_retry(mock_subtensor) + + def test_get_block_with_retry_network_error_no_error(self): + mock_subtensor = MagicMock( + get_current_block=MagicMock(return_value=1), + difficulty=1, + substrate=MagicMock( + get_block_hash=MagicMock(return_value=b'ba7ea4eb0b16dee271dbef5911838c3f359fcf598c74da65a54b919b68b67279') + ) + ) + + # this should not raise an exception because there is no error + bittensor.utils.get_block_with_retry(mock_subtensor) + + def test_get_block_with_retry_network_error_none_twice(self): + # Should retry twice then succeed on the third try + tries = 0 + def block_none_twice(block_hash: bytes): + nonlocal tries + if tries == 1: + return block_hash + else: + tries += 1 + return None + + + mock_subtensor = MagicMock( + get_current_block=MagicMock(return_value=1), + difficulty=1, + substrate=MagicMock( + get_block_hash=MagicMock(side_effect=block_none_twice(b'ba7ea4eb0b16dee271dbef5911838c3f359fcf598c74da65a54b919b68b67279')) + ) + ) + + # this should not raise an exception because there is no error on the third try + bittensor.utils.get_block_with_retry(mock_subtensor) +class TestPOWNotStale(unittest.TestCase): + def test_pow_not_stale_same_block_number(self): + mock_subtensor = MagicMock( + get_current_block=MagicMock(return_value=1), + ) + mock_solution = { + "block_number": 1, + } + + assert bittensor.utils.POWNotStale(mock_subtensor, mock_solution) + + def test_pow_not_stale_diff_block_number(self): + mock_subtensor = MagicMock( + get_current_block=MagicMock(return_value=2), + ) + mock_solution = { + "block_number": 1, # 1 less than current block number + } + + assert bittensor.utils.POWNotStale(mock_subtensor, mock_solution) + + mock_subtensor = MagicMock( + get_current_block=MagicMock(return_value=3), + ) + mock_solution = { + "block_number": 1, # 2 less than current block number + } + + assert bittensor.utils.POWNotStale(mock_subtensor, mock_solution) + + mock_subtensor = MagicMock( + get_current_block=MagicMock(return_value=4), + ) + mock_solution = { + "block_number": 1, # 3 less than current block number + } + + assert bittensor.utils.POWNotStale(mock_subtensor, mock_solution) + + def test_pow_not_stale_diff_block_number_too_old(self): + mock_subtensor = MagicMock( + get_current_block=MagicMock(return_value=5), + ) + mock_solution = { + "block_number": 1, # 4 less than current block number, stale + } + + assert not bittensor.utils.POWNotStale(mock_subtensor, mock_solution) + + if __name__ == "__main__": test_solve_for_difficulty_fast_registered_already() \ No newline at end of file