Skip to content

Commit d64edf2

Browse files
thewhalekingthewhaleking
committed
stake add
1 parent 7dd4fe5 commit d64edf2

File tree

2 files changed

+129
-100
lines changed

2 files changed

+129
-100
lines changed

‎bittensor_cli/cli.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2940,6 +2940,7 @@ def stake_add(
29402940
prompt: bool = Options.prompt,
29412941
quiet: bool = Options.quiet,
29422942
verbose: bool = Options.verbose,
2943+
json_output: bool = Options.json_output,
29432944
):
29442945
"""
29452946
Stake TAO to one or more hotkeys on specific netuids with your coldkey.
@@ -2972,7 +2973,7 @@ def stake_add(
29722973
• [blue]--partial[/blue]: Complete partial stake if rates exceed tolerance
29732974
29742975
"""
2975-
self.verbosity_handler(quiet, verbose)
2976+
self.verbosity_handler(quiet, verbose, json_output)
29762977
safe_staking = self.ask_safe_staking(safe_staking)
29772978
if safe_staking:
29782979
rate_tolerance = self.ask_rate_tolerance(rate_tolerance)
@@ -3130,6 +3131,7 @@ def stake_add(
31303131
safe_staking,
31313132
rate_tolerance,
31323133
allow_partial_stake,
3134+
json_output,
31333135
)
31343136
)
31353137

‎bittensor_cli/src/commands/stake/add.py‎

Lines changed: 126 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
import asyncio
2+
import json
3+
from collections import defaultdict
24
from functools import partial
35

46
from typing import TYPE_CHECKING, Optional
@@ -17,6 +19,7 @@
1719
print_error,
1820
print_verbose,
1921
unlock_key,
22+
json_console,
2023
)
2124
from bittensor_wallet import Wallet
2225

@@ -38,6 +41,7 @@ async def stake_add(
3841
safe_staking: bool,
3942
rate_tolerance: float,
4043
allow_partial_stake: bool,
44+
json_output: bool,
4145
):
4246
"""
4347
Args:
@@ -55,11 +59,13 @@ async def stake_add(
5559
safe_staking: whether to use safe staking
5660
rate_tolerance: rate tolerance percentage for stake operations
5761
allow_partial_stake: whether to allow partial stake
62+
json_output: whether to output stake info in JSON format
5863
5964
Returns:
6065
bool: True if stake operation is successful, False otherwise
6166
"""
6267

68+
# TODO name shadowing
6369
async def safe_stake_extrinsic(
6470
netuid: int,
6571
amount: Balance,
@@ -69,25 +75,25 @@ async def safe_stake_extrinsic(
6975
wallet: Wallet,
7076
subtensor: "SubtensorInterface",
7177
status=None,
72-
) -> None:
78+
) -> bool:
7379
err_out = partial(print_error, status=status)
7480
failure_prelude = (
7581
f":cross_mark: [red]Failed[/red] to stake {amount} on Netuid {netuid}"
7682
)
77-
current_balance = await subtensor.get_balance(wallet.coldkeypub.ss58_address)
78-
next_nonce = await subtensor.substrate.get_account_next_index(
79-
wallet.coldkeypub.ss58_address
80-
)
81-
call = await subtensor.substrate.compose_call(
82-
call_module="SubtensorModule",
83-
call_function="add_stake_limit",
84-
call_params={
85-
"hotkey": hotkey_ss58,
86-
"netuid": netuid,
87-
"amount_staked": amount.rao,
88-
"limit_price": price_limit,
89-
"allow_partial": allow_partial_stake,
90-
},
83+
current_balance, next_nonce, call = await asyncio.gather(
84+
subtensor.get_balance(wallet.coldkeypub.ss58_address),
85+
subtensor.substrate.get_account_next_index(wallet.coldkeypub.ss58_address),
86+
subtensor.substrate.compose_call(
87+
call_module="SubtensorModule",
88+
call_function="add_stake_limit",
89+
call_params={
90+
"hotkey": hotkey_ss58,
91+
"netuid": netuid,
92+
"amount_staked": amount.rao,
93+
"limit_price": price_limit,
94+
"allow_partial": allow_partial_stake,
95+
},
96+
),
9197
)
9298
extrinsic = await subtensor.substrate.create_signed_extrinsic(
9399
call=call, keypair=wallet.coldkey, nonce=next_nonce
@@ -104,71 +110,78 @@ async def safe_stake_extrinsic(
104110
f"Either increase price tolerance or enable partial staking.",
105111
status=status,
106112
)
107-
return
113+
return False
108114
else:
109115
err_out(f"\n{failure_prelude} with error: {format_error_message(e)}")
110-
return
116+
return False
117+
if not await response.is_success:
118+
err_out(
119+
f"\n{failure_prelude} with error: {format_error_message(await response.error_message)}"
120+
)
121+
return False
111122
else:
112-
await response.process_events()
113-
if not await response.is_success:
114-
err_out(
115-
f"\n{failure_prelude} with error: {format_error_message(await response.error_message)}"
116-
)
117-
else:
118-
block_hash = await subtensor.substrate.get_chain_head()
119-
new_balance, new_stake = await asyncio.gather(
120-
subtensor.get_balance(wallet.coldkeypub.ss58_address, block_hash),
121-
subtensor.get_stake(
122-
hotkey_ss58=hotkey_ss58,
123-
coldkey_ss58=wallet.coldkeypub.ss58_address,
124-
netuid=netuid,
125-
block_hash=block_hash,
126-
),
127-
)
128-
console.print(
129-
f":white_heavy_check_mark: [dark_sea_green3]Finalized. Stake added to netuid: {netuid}[/dark_sea_green3]"
130-
)
131-
console.print(
132-
f"Balance:\n [blue]{current_balance}[/blue] :arrow_right: [{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_balance}"
133-
)
134-
135-
amount_staked = current_balance - new_balance
136-
if allow_partial_stake and (amount_staked != amount):
137-
console.print(
138-
"Partial stake transaction. Staked:\n"
139-
f" [{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{amount_staked}[/{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}] "
140-
f"instead of "
141-
f"[blue]{amount}[/blue]"
142-
)
123+
if json_output:
124+
# the rest of this checking is not necessary if using json_output
125+
return True
126+
block_hash = await subtensor.substrate.get_chain_head()
127+
new_balance, new_stake = await asyncio.gather(
128+
subtensor.get_balance(wallet.coldkeypub.ss58_address, block_hash),
129+
subtensor.get_stake(
130+
hotkey_ss58=hotkey_ss58,
131+
coldkey_ss58=wallet.coldkeypub.ss58_address,
132+
netuid=netuid,
133+
block_hash=block_hash,
134+
),
135+
)
136+
console.print(
137+
f":white_heavy_check_mark: [dark_sea_green3]Finalized. "
138+
f"Stake added to netuid: {netuid}[/dark_sea_green3]"
139+
)
140+
console.print(
141+
f"Balance:\n [blue]{current_balance}[/blue] :arrow_right: "
142+
f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_balance}"
143+
)
143144

145+
amount_staked = current_balance - new_balance
146+
if allow_partial_stake and (amount_staked != amount):
144147
console.print(
145-
f"Subnet: [{COLOR_PALETTE['GENERAL']['SUBHEADING']}]{netuid}[/{COLOR_PALETTE['GENERAL']['SUBHEADING']}] "
146-
f"Stake:\n"
147-
f" [blue]{current_stake}[/blue] "
148-
f":arrow_right: "
149-
f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_stake}\n"
148+
"Partial stake transaction. Staked:\n"
149+
f" [{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{amount_staked}"
150+
f"[/{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}] "
151+
f"instead of "
152+
f"[blue]{amount}[/blue]"
150153
)
151154

155+
console.print(
156+
f"Subnet: [{COLOR_PALETTE['GENERAL']['SUBHEADING']}]"
157+
f"{netuid}[/{COLOR_PALETTE['GENERAL']['SUBHEADING']}] "
158+
f"Stake:\n"
159+
f" [blue]{current_stake}[/blue] "
160+
f":arrow_right: "
161+
f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_stake}\n"
162+
)
163+
return True
164+
152165
async def stake_extrinsic(
153166
netuid_i, amount_, current, staking_address_ss58, status=None
154-
):
167+
) -> bool:
155168
err_out = partial(print_error, status=status)
156-
current_balance = await subtensor.get_balance(wallet.coldkeypub.ss58_address)
169+
current_balance, next_nonce, call = await asyncio.gather(
170+
subtensor.get_balance(wallet.coldkeypub.ss58_address),
171+
subtensor.substrate.get_account_next_index(wallet.coldkeypub.ss58_address),
172+
subtensor.substrate.compose_call(
173+
call_module="SubtensorModule",
174+
call_function="add_stake",
175+
call_params={
176+
"hotkey": staking_address_ss58,
177+
"netuid": netuid_i,
178+
"amount_staked": amount_.rao,
179+
},
180+
),
181+
)
157182
failure_prelude = (
158183
f":cross_mark: [red]Failed[/red] to stake {amount} on Netuid {netuid_i}"
159184
)
160-
next_nonce = await subtensor.substrate.get_account_next_index(
161-
wallet.coldkeypub.ss58_address
162-
)
163-
call = await subtensor.substrate.compose_call(
164-
call_module="SubtensorModule",
165-
call_function="add_stake",
166-
call_params={
167-
"hotkey": staking_address_ss58,
168-
"netuid": netuid_i,
169-
"amount_staked": amount_.rao,
170-
},
171-
)
172185
extrinsic = await subtensor.substrate.create_signed_extrinsic(
173186
call=call, keypair=wallet.coldkey, nonce=next_nonce
174187
)
@@ -178,35 +191,46 @@ async def stake_extrinsic(
178191
)
179192
except SubstrateRequestException as e:
180193
err_out(f"\n{failure_prelude} with error: {format_error_message(e)}")
181-
return
194+
return False
182195
else:
183-
await response.process_events()
184196
if not await response.is_success:
185197
err_out(
186198
f"\n{failure_prelude} with error: {format_error_message(await response.error_message)}"
187199
)
200+
return False
188201
else:
202+
if json_output:
203+
# the rest of this is not necessary if using json_output
204+
return True
205+
new_block_hash = await subtensor.substrate.get_chain_head()
189206
new_balance, new_stake = await asyncio.gather(
190-
subtensor.get_balance(wallet.coldkeypub.ss58_address),
207+
subtensor.get_balance(
208+
wallet.coldkeypub.ss58_address, block_hash=new_block_hash
209+
),
191210
subtensor.get_stake(
192211
hotkey_ss58=staking_address_ss58,
193212
coldkey_ss58=wallet.coldkeypub.ss58_address,
194213
netuid=netuid_i,
214+
block_hash=new_block_hash,
195215
),
196216
)
197217
console.print(
198-
f":white_heavy_check_mark: [dark_sea_green3]Finalized. Stake added to netuid: {netuid_i}[/dark_sea_green3]"
218+
f":white_heavy_check_mark: "
219+
f"[dark_sea_green3]Finalized. Stake added to netuid: {netuid_i}[/dark_sea_green3]"
199220
)
200221
console.print(
201-
f"Balance:\n [blue]{current_balance}[/blue] :arrow_right: [{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_balance}"
222+
f"Balance:\n [blue]{current_balance}[/blue] :arrow_right: "
223+
f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_balance}"
202224
)
203225
console.print(
204-
f"Subnet: [{COLOR_PALETTE['GENERAL']['SUBHEADING']}]{netuid_i}[/{COLOR_PALETTE['GENERAL']['SUBHEADING']}] "
226+
f"Subnet: [{COLOR_PALETTE['GENERAL']['SUBHEADING']}]"
227+
f"{netuid_i}[/{COLOR_PALETTE['GENERAL']['SUBHEADING']}] "
205228
f"Stake:\n"
206229
f" [blue]{current}[/blue] "
207230
f":arrow_right: "
208231
f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_stake}\n"
209232
)
233+
return True
210234

211235
netuids = (
212236
[int(netuid)]
@@ -322,7 +346,9 @@ async def stake_extrinsic(
322346
base_row.extend(
323347
[
324348
f"{rate_with_tolerance} {Balance.get_unit(netuid)}/{Balance.get_unit(0)} ",
325-
f"[{'dark_sea_green3' if allow_partial_stake else 'red'}]{allow_partial_stake}[/{'dark_sea_green3' if allow_partial_stake else 'red'}]", # safe staking
349+
f"[{'dark_sea_green3' if allow_partial_stake else 'red'}]"
350+
# safe staking
351+
f"{allow_partial_stake}[/{'dark_sea_green3' if allow_partial_stake else 'red'}]",
326352
]
327353
)
328354

@@ -341,7 +367,7 @@ async def stake_extrinsic(
341367
return False
342368

343369
if safe_staking:
344-
stake_coroutines = []
370+
stake_coroutines = {}
345371
for i, (ni, am, curr, price_with_tolerance) in enumerate(
346372
zip(
347373
netuids, amounts_to_stake, current_stake_balances, prices_with_tolerance
@@ -350,29 +376,25 @@ async def stake_extrinsic(
350376
for _, staking_address in hotkeys_to_stake_to:
351377
# Regular extrinsic for root subnet
352378
if ni == 0:
353-
stake_coroutines.append(
354-
stake_extrinsic(
355-
netuid_i=ni,
356-
amount_=am,
357-
current=curr,
358-
staking_address_ss58=staking_address,
359-
)
379+
stake_coroutines[(ni, staking_address)] = stake_extrinsic(
380+
netuid_i=ni,
381+
amount_=am,
382+
current=curr,
383+
staking_address_ss58=staking_address,
360384
)
361385
else:
362-
stake_coroutines.append(
363-
safe_stake_extrinsic(
364-
netuid=ni,
365-
amount=am,
366-
current_stake=curr,
367-
hotkey_ss58=staking_address,
368-
price_limit=price_with_tolerance,
369-
wallet=wallet,
370-
subtensor=subtensor,
371-
)
386+
stake_coroutines[(ni, staking_address)] = safe_stake_extrinsic(
387+
netuid=ni,
388+
amount=am,
389+
current_stake=curr,
390+
hotkey_ss58=staking_address,
391+
price_limit=price_with_tolerance,
392+
wallet=wallet,
393+
subtensor=subtensor,
372394
)
373395
else:
374-
stake_coroutines = [
375-
stake_extrinsic(
396+
stake_coroutines = {
397+
(ni, staking_address): stake_extrinsic(
376398
netuid_i=ni,
377399
amount_=am,
378400
current=curr,
@@ -382,12 +404,15 @@ async def stake_extrinsic(
382404
zip(netuids, amounts_to_stake, current_stake_balances)
383405
)
384406
for _, staking_address in hotkeys_to_stake_to
385-
]
386-
407+
}
408+
successes = defaultdict(dict)
387409
with console.status(f"\n:satellite: Staking on netuid(s): {netuids} ..."):
388410
# We can gather them all at once but balance reporting will be in race-condition.
389-
for coroutine in stake_coroutines:
390-
await coroutine
411+
for (ni, staking_address), coroutine in stake_coroutines.items():
412+
success = await coroutine
413+
successes[ni][staking_address] = success
414+
if json_output:
415+
json_console.print(json.dumps({"staking_success": successes}))
391416

392417

393418
# Helper functions
@@ -590,7 +615,9 @@ def _print_table_and_slippage(table: Table, max_slippage: float, safe_staking: b
590615
console.print(base_description + (safe_staking_description if safe_staking else ""))
591616

592617

593-
def _calculate_slippage(subnet_info, amount: Balance) -> tuple[Balance, str, float]:
618+
def _calculate_slippage(
619+
subnet_info, amount: Balance
620+
) -> tuple[Balance, str, float, str]:
594621
"""Calculate slippage when adding stake.
595622
596623
Args:

0 commit comments

Comments
 (0)