Skip to content

Commit

Permalink
add add stake multiple
Browse files Browse the repository at this point in the history
  • Loading branch information
@camfairchild
camfairchild committed Aug 15, 2022
1 parent d324d79 commit 6b772bf
Show file tree
Hide file tree
Showing 2 changed files with 198 additions and 4 deletions.
4 changes: 1 addition & 3 deletions bittensor/_cli/cli_impl.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -312,7 +312,6 @@ def unstake( self ):
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 )
Expand Down Expand Up @@ -379,8 +378,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 ):
Expand Down
198 changes: 197 additions & 1 deletion bittensor/_subtensor/subtensor_impl.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -729,6 +729,198 @@ 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 not None and 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

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 ]

neurons = []
with bittensor.__console__.status(":satellite: Syncing with chain: [white]{}[/white] ...".format(self.network)):
old_balance = self.get_balance( wallet.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]

successfull_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

elif not isinstance(amount, bittensor.Balance ):
staking_balance = bittensor.Balance.from_tao( amount )
else:
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
successfull_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
successfull_stakes += 1
if staking_all:
# If staked all, no need to continue
break

if successfull_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',
Expand Down Expand Up @@ -966,7 +1158,7 @@ def unstake_multiple (
wait_for_finalization: bool = False,
prompt: bool = False,
) -> bool:
r""" Removes stake from each wallet coldkey in the list, using each amount.
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.
Expand Down Expand Up @@ -1022,6 +1214,10 @@ def unstake_multiple (
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
Expand Down

0 comments on commit 6b772bf

Please sign in to comment.