solana-program-library/stake-pool/py/bot/rebalance.py

133 lines
6.5 KiB
Python

import argparse
import asyncio
import json
from solana.keypair import Keypair
from solana.publickey import PublicKey
from solana.rpc.async_api import AsyncClient
from solana.rpc.commitment import Confirmed
from stake.constants import STAKE_LEN, LAMPORTS_PER_SOL
from stake_pool.actions import decrease_validator_stake, increase_validator_stake, update_stake_pool
from stake_pool.constants import MINIMUM_ACTIVE_STAKE
from stake_pool.state import StakePool, ValidatorList
async def get_client(endpoint: str) -> AsyncClient:
print(f'Connecting to network at {endpoint}')
async_client = AsyncClient(endpoint=endpoint, commitment=Confirmed)
total_attempts = 10
current_attempt = 0
while not await async_client.is_connected():
if current_attempt == total_attempts:
raise Exception("Could not connect to test validator")
else:
current_attempt += 1
await asyncio.sleep(1)
return async_client
async def rebalance(endpoint: str, stake_pool_address: PublicKey, staker: Keypair, retained_reserve_amount: float):
async_client = await get_client(endpoint)
resp = await async_client.get_epoch_info(commitment=Confirmed)
epoch = resp['result']['epoch']
resp = await async_client.get_account_info(stake_pool_address, commitment=Confirmed)
data = resp['result']['value']['data']
stake_pool = StakePool.decode(data[0], data[1])
print(f'Stake pool last update epoch {stake_pool.last_update_epoch}, current epoch {epoch}')
if stake_pool.last_update_epoch != epoch:
print('Updating stake pool')
await update_stake_pool(async_client, staker, stake_pool_address)
resp = await async_client.get_account_info(stake_pool_address, commitment=Confirmed)
data = resp['result']['value']['data']
stake_pool = StakePool.decode(data[0], data[1])
resp = await async_client.get_minimum_balance_for_rent_exemption(STAKE_LEN)
stake_rent_exemption = resp['result']
retained_reserve_lamports = int(retained_reserve_amount * LAMPORTS_PER_SOL)
resp = await async_client.get_account_info(stake_pool.validator_list, commitment=Confirmed)
data = resp['result']['value']['data']
validator_list = ValidatorList.decode(data[0], data[1])
print('Stake pool stats:')
print(f'* {stake_pool.total_lamports} total lamports')
num_validators = len(validator_list.validators)
print(f'* {num_validators} validators')
print(f'* Retaining {retained_reserve_lamports} lamports in the reserve')
lamports_per_validator = (stake_pool.total_lamports - retained_reserve_lamports) // num_validators
num_increases = sum([
1 for validator in validator_list.validators
if validator.transient_stake_lamports == 0 and validator.active_stake_lamports < lamports_per_validator
])
total_usable_lamports = stake_pool.total_lamports - retained_reserve_lamports - num_increases * stake_rent_exemption
lamports_per_validator = total_usable_lamports // num_validators
print(f'* {lamports_per_validator} lamports desired per validator')
futures = []
for validator in validator_list.validators:
if validator.transient_stake_lamports != 0:
print(f'Skipping {validator.vote_account_address}: {validator.transient_stake_lamports} transient lamports')
else:
if validator.active_stake_lamports > lamports_per_validator:
lamports_to_decrease = validator.active_stake_lamports - lamports_per_validator
if lamports_to_decrease <= stake_rent_exemption:
print(f'Skipping decrease on {validator.vote_account_address}, \
currently at {validator.active_stake_lamports} lamports, \
decrease of {lamports_to_decrease} below the rent exmption')
else:
futures.append(decrease_validator_stake(
async_client, staker, staker, stake_pool_address,
validator.vote_account_address, lamports_to_decrease
))
elif validator.active_stake_lamports < lamports_per_validator:
lamports_to_increase = lamports_per_validator - validator.active_stake_lamports
if lamports_to_increase < MINIMUM_ACTIVE_STAKE:
print(f'Skipping increase on {validator.vote_account_address}, \
currently at {validator.active_stake_lamports} lamports, \
increase of {lamports_to_increase} less than the minimum of {MINIMUM_ACTIVE_STAKE}')
else:
futures.append(increase_validator_stake(
async_client, staker, staker, stake_pool_address,
validator.vote_account_address, lamports_to_increase
))
else:
print(f'{validator.vote_account_address}: already at {lamports_per_validator}')
print('Executing strategy')
await asyncio.gather(*futures)
print('Done')
await async_client.close()
def keypair_from_file(keyfile_name: str) -> Keypair:
with open(keyfile_name, 'r') as keyfile:
data = keyfile.read()
int_list = json.loads(data)
bytes_list = [value.to_bytes(1, 'little') for value in int_list]
return Keypair.from_secret_key(b''.join(bytes_list))
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Rebalance stake evenly between all the validators in a stake pool.')
parser.add_argument('stake_pool', metavar='STAKE_POOL_ADDRESS', type=str,
help='Stake pool to rebalance, given by a public key in base-58,\
e.g. Zg5YBPAk8RqBR9kaLLSoN5C8Uv7nErBz1WC63HTsCPR')
parser.add_argument('staker', metavar='STAKER_KEYPAIR', type=str,
help='Staker for the stake pool, given by a keypair file, e.g. staker.json')
parser.add_argument('reserve_amount', metavar='RESERVE_AMOUNT', type=float,
help='Amount of SOL to keep in the reserve, e.g. 10.5')
parser.add_argument('--endpoint', metavar='ENDPOINT_URL', type=str,
default='https://api.mainnet-beta.solana.com',
help='RPC endpoint to use, e.g. https://api.mainnet-beta.solana.com')
args = parser.parse_args()
stake_pool = PublicKey(args.stake_pool)
staker = keypair_from_file(args.staker)
print(f'Rebalancing stake pool {stake_pool}')
print(f'Staker public key: {staker.public_key}')
print(f'Amount to leave in the reserve: {args.reserve_amount} SOL')
asyncio.run(rebalance(args.endpoint, stake_pool, staker, args.reserve_amount))